timothyafolami
84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
import sys, os
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|
from fastapi import FastAPI, File, UploadFile, BackgroundTasks, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pathlib import Path
|
|
from utils import load_documents_from_directory, create_vector_store, save_embedded_data, process_directory
|
|
from pydantic import BaseModel
|
|
from search import search_and_summarize
|
|
from typing import List
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
# Define allowed origins for CORS
|
|
origins = [
|
|
"http://localhost",
|
|
"http://localhost:8000",
|
|
"http://localhost:3000",
|
|
# Add other allowed origins here
|
|
]
|
|
|
|
# Add CORS middleware
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=origins, # Allows requests from listed origins
|
|
allow_credentials=True,
|
|
allow_methods=["*"], # Allows all HTTP methods
|
|
allow_headers=["*"], # Allows all headers
|
|
)
|
|
|
|
# Define the directory where you want to save uploaded files
|
|
UPLOAD_DIR = Path("./uploads")
|
|
|
|
# Ensure the directory exists
|
|
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
class SearchRequest(BaseModel):
|
|
query: str
|
|
|
|
def load_documents(directory: str):
|
|
|
|
# loading the documents from the directory
|
|
documents, docs_id, num_pages = load_documents_from_directory(directory)
|
|
# embedding the documents
|
|
embed_db = create_vector_store(documents, docs_id, num_pages)
|
|
# saving the embedded data
|
|
status = save_embedded_data(embed_db)
|
|
# creating the thumbnails
|
|
status = process_directory(directory)
|
|
|
|
return {"status": "Documents loaded successfully"}
|
|
|
|
class SearchRequest(BaseModel):
|
|
query: str
|
|
|
|
@app.post("/search/")
|
|
def search(request: SearchRequest):
|
|
|
|
# Perform search using the utility function
|
|
results = search_and_summarize(request.query)
|
|
|
|
return {"results": results}
|
|
|
|
@app.post("/upload/")
|
|
async def upload_file(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
|
|
|
|
file_location = UPLOAD_DIR/file.filename
|
|
|
|
# Save the uploaded file to the specified location
|
|
with open(file_location, "wb") as buffer:
|
|
buffer.write(await file.read())
|
|
|
|
# Add the load_documents function to the background tasks
|
|
background_tasks.add_task(load_documents, str(UPLOAD_DIR))
|
|
|
|
# Return the location of the saved file and inform about the successful upload
|
|
return {"message": "Upload successful. Document loading will begin shortly.", "file_location": str(UPLOAD_DIR)}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|