# main_api.py import os import logging import traceback from typing import Optional, List, Dict, Any import tiktoken from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, Form from fastapi.responses import JSONResponse, FileResponse from fastapi.staticfiles import StaticFiles from pydantic import BaseModel import uvicorn # import your existing modules (assumed in same directory) from memory_store import init_db, save_message, get_last_messages, clear_user_memory, build_gradio_history # :contentReference[oaicite:4]{index=4} from chatbot_retriever import build_or_load_indexes, hybrid_retrieve, retrieve_node_from_rows, load_all_docs, ensure_data_dir # :contentReference[oaicite:5]{index=5} from chatbot_graph import SYSTEM_PROMPT, call_llm, _extract_answer_from_response # :contentReference[oaicite:6]{index=6} # ----------------- CORS SETUP ----------------- from fastapi.middleware.cors import CORSMiddleware app = FastAPI(title="RAG Chat Backend", version="1.0") from fastapi.middleware.cors import CORSMiddleware app.add_middleware( CORSMiddleware, allow_origins=[ "http://localhost:5173", "http://127.0.0.1:5173", ], allow_credentials=True, allow_methods=["*"], # ✅ lowercase 'allow_' allow_headers=["*"], # ✅ lowercase 'allow_' ) # ------------------------------------------------ from dotenv import load_dotenv load_dotenv() logger = logging.getLogger("rag_api") logging.basicConfig(level=logging.INFO) logger.setLevel(logging.INFO) # initialize DB now init_db() # Global in-memory flag/object to check indexes loaded (populated by build_or_load_indexes) INDEXES = {"built": False, "info": None} # ---------- Pydantic models ---------- class ChatRequest(BaseModel): user_id: Optional[str] = None message: str class ChatResponse(BaseModel): user_id: str message: str assistant: str history: List[Dict[str, str]] class RetrieveResponse(BaseModel): query: str context: Optional[str] meta: List[Dict[str, Any]] # ---------- helpers ---------- def ensure_indexes(force_reindex: bool = False): """ Build or load indexes synchronously. This wraps build_or_load_indexes from chatbot_retriever. Also initializes hybrid_retrieve module variables to avoid reindexing. """ # Check if hybrid_retrieve already has indexes built (avoid duplicate work) if hasattr(hybrid_retrieve, "_index_built") and hybrid_retrieve._index_built and not force_reindex: if INDEXES["built"]: return INDEXES["info"] if INDEXES["built"] and not force_reindex: # Ensure hybrid_retrieve module variables are also set if not hasattr(hybrid_retrieve, "_index_built") or not hybrid_retrieve._index_built: # Indexes exist but hybrid_retrieve wasn't initialized, reload them try: chunks, bm25, tokenized, corpus_texts, faiss_data = build_or_load_indexes(force_reindex=False) hybrid_retrieve._chunks = chunks hybrid_retrieve._bm25 = bm25 hybrid_retrieve._tokenized = tokenized hybrid_retrieve._corpus = corpus_texts hybrid_retrieve._faiss = faiss_data hybrid_retrieve._index_built = True logger.info("Initialized hybrid_retrieve module variables from existing indexes") except Exception: logger.exception("Failed to initialize hybrid_retrieve variables") return INDEXES["info"] try: chunks, bm25, tokenized, corpus_texts, faiss_data = build_or_load_indexes(force_reindex=force_reindex) # Set module-level variables in hybrid_retrieve to avoid rebuilding hybrid_retrieve._chunks = chunks hybrid_retrieve._bm25 = bm25 hybrid_retrieve._tokenized = tokenized hybrid_retrieve._corpus = corpus_texts hybrid_retrieve._faiss = faiss_data hybrid_retrieve._index_built = True INDEXES["built"] = True INDEXES["info"] = {"chunks_len": len(chunks) if chunks else 0, "corpus_len": len(corpus_texts) if corpus_texts else 0} logger.info("Indexes built/loaded: %d chunks, %d corpus texts", INDEXES["info"]["chunks_len"], INDEXES["info"]["corpus_len"]) return INDEXES["info"] except Exception: logger.exception("Index build/load failed") raise # ===== Token limiter helper ===== enc = tiktoken.get_encoding("cl100k_base") def trim_to_token_limit(texts, limit=4000): """Join text chunks until token limit is reached.""" joined = "" for t in texts: if len(enc.encode(joined + t)) > limit: break joined += t + "\n" return joined def extract_history_for_frontend(user_id: str, limit: int = 500): return build_gradio_history(user_id) # ---------- Routes ---------- @app.get("/health") def health(): """Basic health check.""" return {"status": "ok", "indexes_built": INDEXES["built"]} @app.post("/reindex") def reindex(force: Optional[bool] = False): """ Force rebuild of indexes. This calls the same build_or_load_indexes used by your retriever module. Use ?force=true to force. """ try: info = ensure_indexes(force_reindex=bool(force)) return {"status": "ok", "info": info} except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to build indexes: {e}") @app.post("/upload") async def upload_file(file: UploadFile = File(...), category: Optional[str] = Form("syllabus")): """ Upload PDF/PPTX into DATA_DIR (same dir used by chatbot_retriever.load_all_docs). After upload you may call /reindex to include the file. """ from chatbot_retriever import DATA_DIR # keep using same constant os.makedirs(DATA_DIR, exist_ok=True) dest_path = os.path.join(DATA_DIR, file.filename) try: with open(dest_path, "wb") as f: content = await file.read() f.write(content) return {"status": "ok", "filename": file.filename, "saved_to": dest_path} except Exception as e: logger.exception("upload failed") raise HTTPException(status_code=500, detail=str(e)) @app.get("/docs_list") def docs_list(): """List files in DATA_DIR (documents available to retriever).""" from chatbot_retriever import DATA_DIR if not os.path.isdir(DATA_DIR): return {"files": []} files = [f for f in os.listdir(DATA_DIR) if os.path.isfile(os.path.join(DATA_DIR, f))] return {"files": files} @app.get("/retrieve", response_model=RetrieveResponse) def retrieve(query: str, subject: Optional[str] = None, top_k: Optional[int] = None): """ Directly call the hybrid retriever for a query. Returns context + meta. """ try: # ensure indexes built (but don't force) ensure_indexes(force_reindex=False) res = hybrid_retrieve(query=query, subject=subject, top_k=(top_k or None)) return {"query": query, "context": res.get("context"), "meta": res.get("meta", [])} except Exception as e: logger.exception("retrieve failed") raise HTTPException(status_code=500, detail=str(e)) @app.get("/history/{user_id}") def get_history(user_id: str, limit: Optional[int] = 500): """Return persisted history for a user (in same format your frontend expects).""" try: hist = extract_history_for_frontend(user_id) if limit: hist = hist[-int(limit):] return {"user_id": user_id, "history": hist} except Exception as e: logger.exception("history fetch failed") raise HTTPException(status_code=500, detail=str(e)) @app.post("/memory/clear") def clear_memory(user_id: str): """Clear stored memory for user.""" try: deleted = clear_user_memory(user_id) return {"status": "ok", "deleted_rows": deleted} except Exception as e: logger.exception("clear failed") raise HTTPException(status_code=500, detail=str(e)) @app.post("/chat", response_model=ChatResponse) def chat(req: ChatRequest): """ Main chat endpoint. - saves user message - fetches last messages from sqlite memory - runs retriever to get context - builds the system prompt + last 3 user messages - calls the LLM via call_llm (same wrapper imported from chatbot_graph) - saves assistant reply and returns it + updated history """ uid = (req.user_id or os.getenv("DEFAULT_USER", "vinayak")).strip() or "vinayak" if not req.message: raise HTTPException(status_code=400, detail="message is required") try: # 1) persist user message save_message(uid, "user", req.message) # 2) get rows (chronological order) for retriever rows = get_last_messages(uid, limit=200) # 3) ensure indexes exist (non-force) try: ensure_indexes(force_reindex=False) except Exception: logger.warning("Indexes not built or failed. retriever may return no context.") # 4) run retrieve_node_from_rows to get context (keeps same logic as your retriever glue) try: retrieved = retrieve_node_from_rows(rows) context = retrieved.get("context") if context: logger.info("Retrieved context: %d characters from documents", len(context)) else: logger.warning("Retriever returned empty context for query: %s", req.message) except Exception as e: logger.exception("retriever call failed: %s", e) context = None # 5) build system prompt content # ===== Combine retrieval context + last 2 user turns ===== MAX_TOKENS_CONTEXT = 3000 NUM_RECENT_TURNS = 2 # last 2 user + assistant pairs # Get last few messages (both user + assistant) recent_pairs = rows[-(NUM_RECENT_TURNS * 2):] recent_chat = "\n".join([f"{r[0].upper()}: {r[1]}" for r in recent_pairs]) # Trim context to token-safe limit context_texts = context.split("\n\n") if context else [] trimmed_context = trim_to_token_limit(context_texts, limit=MAX_TOKENS_CONTEXT) # Final system prompt system_content = SYSTEM_PROMPT if trimmed_context: system_content += "\n\n===== RETRIEVED CONTEXT =====\n" + trimmed_context logger.debug("Added context to system prompt: %d characters", len(trimmed_context)) else: logger.warning("No context to add to system prompt for query: %s", req.message) # build prompt messages as list of simple dicts (call_llm expects same message format as in chatbot_graph) # chatbot_graph.call_llm expects langchain messages (SystemMessage/HumanMessage) — we built that in original file. # create messages as minimal objects that call_llm can accept (we rely on original call_llm). from langchain_core.messages import SystemMessage, HumanMessage # re-use same message classes prompt_msgs = [SystemMessage(content=system_content)] # collect last 3 user messages last_users = [r[1] for r in rows if r[0] == "user"][-1:] if not last_users: last_users = [req.message] for u in last_users: prompt_msgs.append(HumanMessage(content=u)) # 6) call LLM try: raw = call_llm(prompt_msgs) answer = _extract_answer_from_response(raw) or "" except Exception as e: logger.exception("LLM call failed") # If LLM client not configured (ChatGroq missing or no API KEY), return helpful message detail = str(e) answer = f"LLM call failed: {detail}" # 7) persist assistant reply try: save_message(uid, "assistant", answer) except Exception: logger.exception("Failed to persist assistant message") # 8) build history to return history = extract_history_for_frontend(uid) return { "user_id": uid, "message": req.message, "assistant": answer, "history": history, } except HTTPException: raise except Exception as e: logger.exception("chat failed: %s", e) raise HTTPException(status_code=500, detail=str(e)) # Mount static files for frontend FRONTEND_DIR = os.path.join(os.path.dirname(__file__), "frontend", "dist") if os.path.exists(FRONTEND_DIR): app.mount("/assets", StaticFiles(directory=os.path.join(FRONTEND_DIR, "assets")), name="assets") @app.get("/{full_path:path}") async def serve_frontend(full_path: str): """Serve the React frontend for all non-API routes""" if full_path and not full_path.startswith("api"): file_path = os.path.join(FRONTEND_DIR, full_path) if os.path.exists(file_path) and os.path.isfile(file_path): return FileResponse(file_path) return FileResponse(os.path.join(FRONTEND_DIR, "index.html")) # Run with: uvicorn main_api:app --reload --host 127.0.0.1 --port 8000 if __name__ == "__main__": uvicorn.run("main_api:app", host="127.0.0.1", port=8000, reload=True)