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 ) from sentence_transformers import SentenceTransformer # 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 # Create a copy of data that is safe to pickle pickle_safe_data = { "file_path": data.get("file_path"), "file_name": data.get("file_name"), "chunks": data.get("chunks"), "chat_history": data.get("chat_history", []) } # Persist to disk 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="meta-llama/llama-4-scout-17b-16e-instruct"): try: # Check if session is already in memory if session_id in sessions: # Ensure the LLM in the cached session matches the requested model_name # If not, update it. This handles cases where model_name might change for an existing session. if sessions[session_id].get("llm") is None or sessions[session_id]["llm"].model_name != model_name: try: sessions[session_id]["llm"] = model_selection(model_name) except Exception as e: print(f"Error updating LLM for in-memory session {session_id} to {model_name}: {str(e)}") # Decide if this is a critical error; for now, we'll proceed with the old LLM or handle as error # For simplicity, if LLM update fails, we might want to indicate session load failure or use existing. # Here, we'll let it proceed, but this could be a point of further refinement. return sessions[session_id], True # Try to load from disk 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) # This is pickle_safe_data # Recreate non-pickled objects # Ensure 'chunks' and 'file_path' (for the original PDF) are present in the loaded data # and the original PDF file still exists. 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') # data["chunks"] is already the list of dicts: {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"), # These are 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 # Store in memory cache return full_session_data, True else: # If essential data for reconstruction is missing from pickle or the original PDF is gone print(f"Warning: Session data for {session_id} is incomplete or its PDF file '{original_pdf_path}' is missing. Cannot reconstruct session.") # Optionally, remove the stale .pkl file # os.remove(file_path_pkl) return None, False return None, False # Session not in memory and not found on disk, or reconstruction failed except Exception as e: print(f"Error loading session {session_id}: {str(e)}") print(traceback.format_exc()) # Print full traceback for debugging 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("meta-llama/llama-4-scout-17b-16e-instruct") ): # Generate a unique session ID session_id = str(uuid.uuid4()) file_path = None try: # Save the uploaded file file_path = f"{UPLOAD_DIR}/{session_id}_{file.filename}" with open(file_path, "wb") as buffer: shutil.copyfileobj(file.file, buffer) # Check if API keys are set if not os.getenv("GROQ_API_KEY"): raise ValueError("GROQ_API_KEY is not set in the environment variables") # Process the PDF documents = process_pdf_file(file_path) # Returns list of Document objects chunks = chunk_text(documents, max_length=1500) # Updated to handle documents # Create embeddings model = SentenceTransformer('BAAI/bge-large-en-v1.5') # Updated embedding model embeddings, chunks_with_metadata = create_embeddings(chunks, model) # Unpack tuple # Build FAISS index index = build_faiss_index(embeddings) # Pass only embeddings array # Initialize LLM llm = model_selection(model_name) # Save session data session_data = { "file_path": file_path, "file_name": file.filename, "chunks": chunks_with_metadata, # Store chunks with metadata "model": model, "index": index, "llm": llm, "chat_history": [] } save_session(session_id, session_data) return {"status": "success", "session_id": session_id, "message": f"Processed {file.filename}"} except Exception as e: # Clean up on error 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}") print(f"Stack trace: {stack_trace}") return JSONResponse( status_code=400, content={ "status": "error", "detail": error_msg, "type": type(e).__name__ } ) # Route to chat with the document @app.post("/chat") async def chat(request: ChatRequest): # Try to load session if not in memory session, found = load_session(request.session_id, model_name=request.model_name) if not found: raise HTTPException(status_code=404, detail="Session not found. Please upload a document first.") try: from langchain.memory import ConversationBufferMemory agent_memory = ConversationBufferMemory(memory_key="chat_history", 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"]) # Retrieve similar chunks similar_chunks = retrieve_similar_chunks( request.query, session["index"], session["chunks"], session["model"], k=10 ) # Generate response using agentic_rag response = agentic_rag( session["llm"], tools, query=request.query, context_chunks=similar_chunks, # Pass the list of tuples Use_Tavily=request.use_search, memory=agent_memory ) # Update chat history session["chat_history"].append({"user": request.query, "assistant": response["output"]}) save_session(request.session_id, session) return { "status": "success", "answer": response["output"], "context_used": [{"text": chunk, "score": float(score)} for chunk, score, _ in similar_chunks] } except Exception as e: 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)