diff --git a/examples/authentication-methods/no-auth/main.py b/examples/authentication-methods/no-auth/main.py index f7dc46921..1fd5458b4 100644 --- a/examples/authentication-methods/no-auth/main.py +++ b/examples/authentication-methods/no-auth/main.py @@ -1,7 +1,8 @@ # This is a version of the main.py file found in ../../../server/main.py without authentication. # Copy and paste this into the main file at ../../../server/main.py if you choose to use no authentication for your retrieval plugin. +from typing import Optional import uvicorn -from fastapi import FastAPI, File, HTTPException, Body, UploadFile +from fastapi import FastAPI, File, Form, HTTPException, Body, UploadFile from fastapi.staticfiles import StaticFiles from models.api import ( @@ -15,6 +16,8 @@ from datastore.factory import get_datastore from services.file import get_document_from_file +from models.models import DocumentMetadata, Source + app = FastAPI() app.mount("/.well-known", StaticFiles(directory=".well-known"), name="static") @@ -35,8 +38,18 @@ ) async def upsert_file( file: UploadFile = File(...), + metadata: Optional[str] = Form(None), ): - document = await get_document_from_file(file) + try: + metadata_obj = ( + DocumentMetadata.parse_raw(metadata) + if metadata + else DocumentMetadata(source=Source.file) + ) + except: + metadata_obj = DocumentMetadata(source=Source.file) + + document = await get_document_from_file(file, metadata_obj) try: ids = await datastore.upsert([document]) diff --git a/examples/memory/main.py b/examples/memory/main.py index cdbab3ea5..5c96e4289 100644 --- a/examples/memory/main.py +++ b/examples/memory/main.py @@ -3,8 +3,9 @@ # Copy and paste this into the main file at ../../server/main.py if you choose to give the model access to the upsert endpoint # and want to access the openapi.json when you run the app locally at http://0.0.0.0:8000/sub/openapi.json. import os +from typing import Optional import uvicorn -from fastapi import FastAPI, File, HTTPException, Depends, Body, UploadFile +from fastapi import FastAPI, File, Form, HTTPException, Depends, Body, UploadFile from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from fastapi.staticfiles import StaticFiles @@ -19,6 +20,8 @@ from datastore.factory import get_datastore from services.file import get_document_from_file +from models.models import DocumentMetadata, Source + bearer_scheme = HTTPBearer() BEARER_TOKEN = os.environ.get("BEARER_TOKEN") @@ -51,9 +54,18 @@ def validate_token(credentials: HTTPAuthorizationCredentials = Depends(bearer_sc ) async def upsert_file( file: UploadFile = File(...), - token: HTTPAuthorizationCredentials = Depends(validate_token), + metadata: Optional[str] = Form(None), ): - document = await get_document_from_file(file) + try: + metadata_obj = ( + DocumentMetadata.parse_raw(metadata) + if metadata + else DocumentMetadata(source=Source.file) + ) + except: + metadata_obj = DocumentMetadata(source=Source.file) + + document = await get_document_from_file(file, metadata_obj) try: ids = await datastore.upsert([document]) diff --git a/server/main.py b/server/main.py index 0e510cb9d..3d44ced4f 100644 --- a/server/main.py +++ b/server/main.py @@ -1,6 +1,7 @@ import os +from typing import Optional import uvicorn -from fastapi import FastAPI, File, HTTPException, Depends, Body, UploadFile +from fastapi import FastAPI, File, Form, HTTPException, Depends, Body, UploadFile from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from fastapi.staticfiles import StaticFiles @@ -15,6 +16,8 @@ from datastore.factory import get_datastore from services.file import get_document_from_file +from models.models import DocumentMetadata, Source + bearer_scheme = HTTPBearer() BEARER_TOKEN = os.environ.get("BEARER_TOKEN") assert BEARER_TOKEN is not None @@ -46,8 +49,18 @@ def validate_token(credentials: HTTPAuthorizationCredentials = Depends(bearer_sc ) async def upsert_file( file: UploadFile = File(...), + metadata: Optional[str] = Form(None), ): - document = await get_document_from_file(file) + try: + metadata_obj = ( + DocumentMetadata.parse_raw(metadata) + if metadata + else DocumentMetadata(source=Source.file) + ) + except: + metadata_obj = DocumentMetadata(source=Source.file) + + document = await get_document_from_file(file, metadata_obj) try: ids = await datastore.upsert([document]) diff --git a/services/file.py b/services/file.py index 3cea13b80..90e0e5ea0 100644 --- a/services/file.py +++ b/services/file.py @@ -8,14 +8,14 @@ import csv import pptx -from models.models import Document, DocumentMetadata, Source +from models.models import Document, DocumentMetadata -async def get_document_from_file(file: UploadFile) -> Document: +async def get_document_from_file( + file: UploadFile, metadata: DocumentMetadata +) -> Document: extracted_text = await extract_text_from_form_file(file) - metadata = DocumentMetadata( - source=Source.file, - ) + doc = Document(text=extracted_text, metadata=metadata) return doc