Files
ds_fire_fighter/main.py
T
timothyafolami ddd0dda276 last commi
2024-08-16 21:39:28 +01:00

59 lines
1.7 KiB
Python

import sys, os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from utils import load_embedded_data, load_documents_from_directory, create_vector_store, save_embedded_data, process_directory
from search import search_and_summarize
from data_ingest import load_data
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
)
class SearchRequest(BaseModel):
query: str
@app.post("/load_documents")
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"}
@app.post("/search")
def search(request: SearchRequest):
# Perform search using the utility function
results = search_and_summarize(request.query)
return {"results": results}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)