last commi

ai_index.py

@@ -0,0 +1,83 @@
+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)