Update backend/main.py

backend/main.py

@@ -1 +1,144 @@
- 
+import os
+import tempfile
+import shutil
+from typing import List, Dict, Any
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader, PDFLoader
+from aimakerspace.openai_utils.prompts import (
+    UserRolePrompt,
+    SystemRolePrompt,
+    AssistantRolePrompt,
+)
+from aimakerspace.openai_utils.embedding import EmbeddingModel
+from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.openai_utils.chatmodel import ChatOpenAI
+
+app = FastAPI()
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, replace with your frontend URL
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize components
+text_splitter = CharacterTextSplitter()
+vector_db = None
+chat_openai = ChatOpenAI(model_name="gpt-3.5-turbo")
+
+system_template = """\
+Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
+system_role_prompt = SystemRolePrompt(system_template)
+
+user_prompt_template = """\
+Context:
+{context}
+
+Question:
+{question}
+"""
+user_role_prompt = UserRolePrompt(user_prompt_template)
+
+class QuestionRequest(BaseModel):
+    question: str
+
+class QuestionResponse(BaseModel):
+    response: str
+    context: List[tuple]
+
+def process_file(file_path: str, file_name: str):
+    print(f"Processing file: {file_name}")
+    
+    # Create appropriate loader
+    if file_name.lower().endswith('.pdf'):
+        loader = PDFLoader(file_path)
+    else:
+        loader = TextFileLoader(file_path)
+        
+    # Load and process the documents
+    documents = loader.load_documents()
+    texts = text_splitter.split_texts(documents)
+    return texts
+
+class RetrievalAugmentedQAPipeline:
+    def __init__(self, llm: ChatOpenAI, vector_db_retriever: VectorDatabase) -> None:
+        self.llm = llm
+        self.vector_db_retriever = vector_db_retriever
+
+    async def run_pipeline(self, user_query: str):
+        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
+
+        context_prompt = ""
+        for context in context_list:
+            context_prompt += context[0] + "\n"
+
+        formatted_system_prompt = system_role_prompt.create_message()
+        formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)
+
+        # Remove await since run is not an async method
+        response = self.llm.run([formatted_system_prompt, formatted_user_prompt])
+        
+        return {"response": response, "context": context_list}
+
+@app.post("/upload")
+async def upload_file(file: UploadFile = File(...)):
+    global vector_db
+    
+    # Validate file type
+    if not file.filename.lower().endswith(('.txt', '.pdf')):
+        raise HTTPException(status_code=400, detail="Only .txt and .pdf files are allowed")
+    
+    # Create a temporary file
+    suffix = f".{file.filename.split('.')[-1]}"
+    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
+        # Copy the uploaded file content to the temporary file
+        content = await file.read()
+        with open(temp_file.name, "wb") as f:
+            f.write(content)
+        
+        try:
+            # Process the file
+            texts = process_file(temp_file.name, file.filename)
+            print(f"Processing {len(texts)} text chunks")
+            
+            # Create a vector store
+            vector_db = VectorDatabase()
+            vector_db = await vector_db.abuild_from_list(texts)
+            
+            return {"message": f"File {file.filename} processed successfully", "chunks": len(texts)}
+        finally:
+            # Clean up the temporary file
+            try:
+                os.unlink(temp_file.name)
+            except Exception as e:
+                print(f"Error cleaning up temporary file: {e}")
+
+@app.post("/ask", response_model=QuestionResponse)
+async def ask_question(request: QuestionRequest):
+    global vector_db
+    
+    if vector_db is None:
+        raise HTTPException(status_code=400, detail="Please upload a file first")
+    
+    # Create a chain
+    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
+        vector_db_retriever=vector_db,
+        llm=chat_openai
+    )
+    
+    # Run the pipeline
+    result = await retrieval_augmented_qa_pipeline.run_pipeline(request.question)
+    
+    return QuestionResponse(
+        response=result["response"],
+        context=result["context"]
+    )
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}