Refactor FastAPI application for improved modularity and maintainability

app.py

@@ -1,357 +0,0 @@
-import os
-import dotenv
-import pickle
-import uuid
-import shutil
-import traceback
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.staticfiles import StaticFiles
-from pydantic import BaseModel
-import uvicorn
-from preprocessing import (
-    model_selection,
-    process_pdf_file,
-    chunk_text,
-    create_embeddings,
-    build_faiss_index,
-    retrieve_similar_chunks,
-    agentic_rag,
-    tools as global_base_tools,  
-    create_vector_search_tool  
-)
-from sentence_transformers import SentenceTransformer
-from langchain.memory import ConversationBufferMemory
-
-# Load environment variables
-dotenv.load_dotenv()
-
-# Initialize FastAPI app
-app = FastAPI(title="PDF Insight Beta", description="Agentic RAG for PDF documents")
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Create upload directory if it doesn't exist
-UPLOAD_DIR = "uploads"
-if not os.path.exists(UPLOAD_DIR):
-    os.makedirs(UPLOAD_DIR)
-
-# Store active sessions
-sessions = {}
-
-# Define model for chat request
-class ChatRequest(BaseModel):
-    session_id: str
-    query: str
-    use_search: bool = False
-    model_name: str = "meta-llama/llama-4-scout-17b-16e-instruct"
-
-class SessionRequest(BaseModel):
-    session_id: str
-
-# Function to save session data
-def save_session(session_id, data):
-    sessions[session_id] = data # Keep non-picklable in memory for active session
-    
-    pickle_safe_data = {
-        "file_path": data.get("file_path"),
-        "file_name": data.get("file_name"),
-        "chunks": data.get("chunks"), # Chunks with metadata (list of dicts)
-        "chat_history": data.get("chat_history", [])
-        # FAISS index, embedding model, and LLM model are not pickled, will be reloaded/recreated
-    }
-    
-    with open(f"{UPLOAD_DIR}/{session_id}_session.pkl", "wb") as f:
-        pickle.dump(pickle_safe_data, f)
-
-
-# Function to load session data
-def load_session(session_id, model_name="llama3-8b-8192"): # Ensure model_name matches default
-    try:
-        if session_id in sessions:
-            cached_session = sessions[session_id]
-            # Ensure LLM and potentially other non-pickled parts are up-to-date or loaded
-            if cached_session.get("llm") is None or (hasattr(cached_session["llm"], "model_name") and cached_session["llm"].model_name != model_name):
-                 cached_session["llm"] = model_selection(model_name)
-            if cached_session.get("model") is None: # Embedding model
-                 cached_session["model"] = SentenceTransformer('BAAI/bge-large-en-v1.5')
-            if cached_session.get("index") is None and cached_session.get("chunks"): # FAISS index
-                embeddings, _ = create_embeddings(cached_session["chunks"], cached_session["model"])
-                cached_session["index"] = build_faiss_index(embeddings)
-            return cached_session, True
-        
-        file_path_pkl = f"{UPLOAD_DIR}/{session_id}_session.pkl"
-        if os.path.exists(file_path_pkl):
-            with open(file_path_pkl, "rb") as f:
-                data = pickle.load(f)
-            
-            original_pdf_path = data.get("file_path")
-            if data.get("chunks") and original_pdf_path and os.path.exists(original_pdf_path):
-                embedding_model_instance = SentenceTransformer('BAAI/bge-large-en-v1.5')
-                # Chunks are already {text: ..., metadata: ...}
-                recreated_embeddings, _ = create_embeddings(data["chunks"], embedding_model_instance)
-                recreated_index = build_faiss_index(recreated_embeddings)
-                recreated_llm = model_selection(model_name)
-
-                full_session_data = {
-                    "file_path": original_pdf_path,
-                    "file_name": data.get("file_name"),
-                    "chunks": data.get("chunks"), # chunks_with_metadata
-                    "chat_history": data.get("chat_history", []),
-                    "model": embedding_model_instance,    # SentenceTransformer model
-                    "index": recreated_index,            # FAISS index
-                    "llm": recreated_llm                 # LLM
-                }
-                sessions[session_id] = full_session_data
-                return full_session_data, True
-            else:
-                print(f"Warning: Session data for {session_id} is incomplete or PDF missing. Cannot reconstruct.")
-                if os.path.exists(file_path_pkl): os.remove(file_path_pkl) # Clean up stale pkl
-                return None, False
-        
-        return None, False
-    except Exception as e:
-        print(f"Error loading session {session_id}: {str(e)}")
-        print(traceback.format_exc())
-        return None, False
-
-# Function to remove PDF file
-def remove_pdf_file(session_id):
-    try:
-        # Check if the session exists
-        session_path = f"{UPLOAD_DIR}/{session_id}_session.pkl"
-        if os.path.exists(session_path):
-            # Load session data
-            with open(session_path, "rb") as f:
-                data = pickle.load(f)
-            
-            # Delete PDF file if it exists
-            if data.get("file_path") and os.path.exists(data["file_path"]):
-                os.remove(data["file_path"])
-            
-            # Remove session file
-            os.remove(session_path)
-        
-        # Remove from memory if exists
-        if session_id in sessions:
-            del sessions[session_id]
-            
-        return True
-    except Exception as e:
-        print(f"Error removing PDF file: {str(e)}")
-        return False
-
-# Mount static files (we'll create these later)
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-# Route for the home page
-@app.get("/")
-async def read_root():
-    from fastapi.responses import RedirectResponse
-    return RedirectResponse(url="/static/index.html")
-
-# Route to upload a PDF file
-@app.post("/upload-pdf")
-async def upload_pdf(
-    file: UploadFile = File(...), 
-    model_name: str = Form("llama3-8b-8192") # Default model
-):
-    session_id = str(uuid.uuid4())
-    file_path = None
-    
-    try:
-        file_path = f"{UPLOAD_DIR}/{session_id}_{file.filename}"
-        with open(file_path, "wb") as buffer:
-            shutil.copyfileobj(file.file, buffer)
-        
-        if not os.getenv("GROQ_API_KEY") and "llama" in model_name: # Llama specific check for Groq
-             raise ValueError("GROQ_API_KEY is not set for Groq Llama models.")
-        if not os.getenv("TAVILY_API_KEY"): # Needed for TavilySearchResults
-            print("Warning: TAVILY_API_KEY is not set. Web search will not function.")
-
-        documents = process_pdf_file(file_path)
-        chunks_with_metadata = chunk_text(documents, max_length=1000) # Increased from 256 to 1000 tokens for better context
-        
-        embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5')
-        embeddings, _ = create_embeddings(chunks_with_metadata, embedding_model) # Chunks are already with metadata
-        
-        index = build_faiss_index(embeddings)
-        llm = model_selection(model_name)
-        
-        session_data = {
-            "file_path": file_path,
-            "file_name": file.filename,
-            "chunks": chunks_with_metadata, # Store chunks with metadata
-            "model": embedding_model,       # SentenceTransformer instance
-            "index": index,                 # FAISS index instance
-            "llm": llm,                     # LLM instance
-            "chat_history": []
-        }
-        save_session(session_id, session_data)
-        
-        return {"status": "success", "session_id": session_id, "message": f"Processed {file.filename}"}
-    
-    except Exception as e:
-        if file_path and os.path.exists(file_path):
-            os.remove(file_path)
-        error_msg = str(e)
-        stack_trace = traceback.format_exc()
-        print(f"Error processing PDF: {error_msg}\nStack trace: {stack_trace}")
-        return JSONResponse(
-            status_code=500, # Internal server error for processing issues
-            content={"status": "error", "detail": error_msg, "type": type(e).__name__}
-        )
-
-# Route to chat with the document
-@app.post("/chat")
-async def chat(request: ChatRequest):
-    # Validate query
-    if not request.query or not request.query.strip():
-        raise HTTPException(status_code=400, detail="Query cannot be empty")
-    
-    if len(request.query.strip()) < 3:
-        raise HTTPException(status_code=400, detail="Query must be at least 3 characters long")
-    
-    session, found = load_session(request.session_id, model_name=request.model_name)
-    if not found:
-        raise HTTPException(status_code=404, detail="Session not found or expired. Please upload a document first.")
-    
-    try:
-        # Validate session data integrity
-        required_keys = ["index", "chunks", "model", "llm"]
-        missing_keys = [key for key in required_keys if key not in session]
-        if missing_keys:
-            print(f"Warning: Session {request.session_id} missing required data: {missing_keys}")
-            raise HTTPException(status_code=500, detail="Session data is incomplete. Please upload the document again.")
-        
-        # Per-request memory to ensure chat history is correctly loaded for the agent
-        agent_memory = ConversationBufferMemory(memory_key="chat_history", input_key="input", return_messages=True)
-        for entry in session.get("chat_history", []):
-            agent_memory.chat_memory.add_user_message(entry["user"])
-            agent_memory.chat_memory.add_ai_message(entry["assistant"])
-
-        # Prepare tools for the agent for THIS request
-        current_request_tools = []
-
-        # 1. Add the document-specific vector search tool
-        vector_search_tool_instance = create_vector_search_tool(
-            faiss_index=session["index"],
-            document_chunks_with_metadata=session["chunks"], # Pass the correct variable
-            embedding_model=session["model"], # This is the SentenceTransformer model
-            max_chunk_length=1000,
-            k=10
-        )
-        current_request_tools.append(vector_search_tool_instance)
-
-        # 2. Conditionally add Tavily (web search) tool
-        if request.use_search:
-            if os.getenv("TAVILY_API_KEY"):
-                tavily_tool = next((tool for tool in global_base_tools if tool.name == "tavily_search_results_json"), None)
-                if tavily_tool:
-                    current_request_tools.append(tavily_tool)
-                else: # Should not happen if global_base_tools is defined correctly
-                    print("Warning: Tavily search requested, but tool misconfigured.")
-            else:
-                print("Warning: Tavily search requested, but TAVILY_API_KEY is not set.")
-        
-        # Retrieve initial similar chunks for RAG context (can be empty if no good match)
-        # This context is given to the agent *before* it decides to use tools.
-        # k=5 means we retrieve up to 5 chunks for initial context.
-        # The agent can then use `vector_database_search` to search more if needed.
-        initial_similar_chunks = retrieve_similar_chunks(
-            request.query, 
-            session["index"], 
-            session["chunks"], # list of dicts {text:..., metadata:...}
-            session["model"], # SentenceTransformer model
-            k=5 # Number of chunks for initial context
-        )
-        
-        print(f"Query: '{request.query}' - Found {len(initial_similar_chunks)} initial chunks")
-        if initial_similar_chunks:
-            print(f"Best chunk score: {initial_similar_chunks[0][1]:.4f}")
-        
-        response = agentic_rag(
-            session["llm"], 
-            current_request_tools, # Pass the dynamically assembled list of tools
-            query=request.query, 
-            context_chunks=initial_similar_chunks,
-            Use_Tavily=request.use_search, # Still passed to agentic_rag for potential fine-grained logic, though prompt adapts to tools
-            memory=agent_memory
-        )
-        
-        response_output = response.get("output", "Sorry, I could not generate a response.")
-        print(f"Generated response length: {len(response_output)} characters")
-        
-        session["chat_history"].append({"user": request.query, "assistant": response_output})
-        save_session(request.session_id, session) # Save updated history and potentially other modified session state
-        
-        return {
-            "status": "success", 
-            "answer": response_output,
-            # Return context that was PRE-FETCHED for the agent, not necessarily all context it might have used via tools
-            "context_used": [{"text": chunk, "score": float(score), "metadata": meta} for chunk, score, meta in initial_similar_chunks]
-        }
-            
-    except Exception as e:
-        print(f"Error processing chat query: {str(e)}\nTraceback: {traceback.format_exc()}")
-        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")
-
-
-# Route to get chat history
-@app.post("/chat-history")
-async def get_chat_history(request: SessionRequest):
-    # Try to load session if not in memory
-    session, found = load_session(request.session_id)
-    if not found:
-        raise HTTPException(status_code=404, detail="Session not found")
-    
-    return {
-        "status": "success", 
-        "history": session.get("chat_history", [])
-    }
-
-# Route to clear chat history
-@app.post("/clear-history")
-async def clear_history(request: SessionRequest):
-    # Try to load session if not in memory
-    session, found = load_session(request.session_id)
-    if not found:
-        raise HTTPException(status_code=404, detail="Session not found")
-    
-    session["chat_history"] = []
-    save_session(request.session_id, session)
-    
-    return {"status": "success", "message": "Chat history cleared"}
-
-# Route to remove PDF from session
-@app.post("/remove-pdf")
-async def remove_pdf(request: SessionRequest):
-    success = remove_pdf_file(request.session_id)
-    
-    if success:
-        return {"status": "success", "message": "PDF file and session removed successfully"}
-    else:
-        raise HTTPException(status_code=404, detail="Session not found or could not be removed")
-
-# Route to list available models
-@app.get("/models")
-async def get_models():
-    # You can expand this list as needed
-    models = [
-        {"id": "meta-llama/llama-4-scout-17b-16e-instruct", "name": "Llama 4 Scout 17B"},
-        {"id": "llama-3.1-8b-instant", "name": "Llama 3.1 8B Instant"},
-        {"id": "llama-3.3-70b-versatile", "name": "Llama 3.3 70B Versatile"},
-    ]
-    return {"models": models}
-
-# Run the application if this file is executed directly
-if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)

app_refactored.py

@@ -0,0 +1,107 @@
+"""
+Refactored FastAPI application for PDF Insight Beta.
+
+This is the main application file that sets up the FastAPI app with modular components.
+The core logic has been preserved while improving code organization and maintainability.
+"""
+
+import uvicorn
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+
+from configs.config import Config
+from models.models import (
+    ChatRequest, SessionRequest, UploadResponse, ChatResponse, 
+    ChatHistoryResponse, StatusResponse, ModelsResponse
+)
+from api import (
+    upload_pdf_handler, chat_handler, get_chat_history_handler,
+    clear_history_handler, remove_pdf_handler, root_handler, get_models_handler
+)
+
+
+def create_app() -> FastAPI:
+    """
+    Create and configure the FastAPI application.
+    
+    Returns:
+        Configured FastAPI application instance
+    """
+    # Initialize FastAPI app
+    app = FastAPI(
+        title="PDF Insight Beta", 
+        description="Agentic RAG for PDF documents"
+    )
+    
+    # Add CORS middleware
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=Config.CORS_ORIGINS,
+        allow_credentials=Config.CORS_CREDENTIALS,
+        allow_methods=Config.CORS_METHODS,
+        allow_headers=Config.CORS_HEADERS,
+    )
+    
+    # Mount static files
+    app.mount("/static", StaticFiles(directory="static"), name="static")
+    
+    return app
+
+
+# Create app instance
+app = create_app()
+
+
+# Route definitions
+@app.get("/")
+async def read_root():
+    """Root endpoint that redirects to the main application."""
+    return await root_handler()
+
+
+@app.post("/upload-pdf", response_model=UploadResponse)
+async def upload_pdf(file: UploadFile = File(...), model_name: str = Form(Config.DEFAULT_MODEL)):
+    """Upload and process a PDF file."""
+    return await upload_pdf_handler(file, model_name)
+
+
+@app.post("/chat", response_model=ChatResponse)
+async def chat(request: ChatRequest):
+    """Chat with the uploaded document."""
+    return await chat_handler(request)
+
+
+@app.post("/chat-history", response_model=ChatHistoryResponse)
+async def get_chat_history(request: SessionRequest):
+    """Get chat history for a session."""
+    return await get_chat_history_handler(request)
+
+
+@app.post("/clear-history", response_model=StatusResponse)
+async def clear_history(request: SessionRequest):
+    """Clear chat history for a session."""
+    return await clear_history_handler(request)
+
+
+@app.post("/remove-pdf", response_model=StatusResponse)
+async def remove_pdf(request: SessionRequest):
+    """Remove PDF file and session data."""
+    return await remove_pdf_handler(request)
+
+
+@app.get("/models", response_model=ModelsResponse)
+async def get_models():
+    """Get list of available models."""
+    return await get_models_handler()
+
+
+def main():
+    """
+    Main entry point for running the application.
+    """
+    uvicorn.run("app_refactored:app", host="0.0.0.0", port=8000, reload=True)
+
+
+if __name__ == "__main__":
+    main()