import os import uuid import asyncio import shutil import tempfile import time from pathlib import Path from typing import Dict from fastapi import FastAPI, UploadFile, File, Form, HTTPException from fastapi.responses import StreamingResponse, JSONResponse, FileResponse from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from dotenv import load_dotenv from sentence_transformers import SentenceTransformer from src.data_loader import load_single_file, SUPPORTED_EXTENSIONS from src.vectorstore import FaissVectorStore from src.search import RAGSearch load_dotenv() app = FastAPI(title="RAG Chat") ALLOWED_ORIGINS = os.getenv( "ALLOWED_ORIGINS", "http://localhost:5173,http://localhost:3000" ).split(",") app.add_middleware( CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) shared_model: SentenceTransformer = None UPLOAD_BASE = Path(tempfile.gettempdir()) / "rag_uploads" STORE_BASE = Path(tempfile.gettempdir()) / "rag_stores" ASSETS_BASE = Path(tempfile.gettempdir()) / "rag_assets" MAX_FILE_SIZE = 10 * 1024 * 1024 MAX_FILES_PER_SESSION = 5 SESSION_TTL_SECONDS = 24 * 60 * 60 sessions: Dict[str, dict] = {} session_locks: Dict[str, asyncio.Lock] = {} @app.on_event("startup") async def startup(): global shared_model print("[INFO] Loading shared embedding model...") shared_model = SentenceTransformer("all-MiniLM-L6-v2") print("[INFO] Shared embedding model loaded.") asyncio.create_task(cleanup_old_sessions()) async def cleanup_old_sessions(): """Background task: delete sessions older than 24 hours, runs hourly.""" while True: await asyncio.sleep(3600) now = time.time() expired = [ sid for sid, sdata in sessions.items() if now - sdata.get("created_at", now) > SESSION_TTL_SECONDS ] for sid in expired: print(f"[INFO] Cleaning up expired session: {sid}") session = sessions.pop(sid, None) session_locks.pop(sid, None) if session: upload_dir = session.get("upload_dir") if upload_dir and upload_dir.exists(): shutil.rmtree(upload_dir, ignore_errors=True) store_dir = STORE_BASE / sid if store_dir.exists(): shutil.rmtree(store_dir, ignore_errors=True) assets_dir = ASSETS_BASE / sid if assets_dir.exists(): shutil.rmtree(assets_dir, ignore_errors=True) if expired: print(f"[INFO] Cleaned up {len(expired)} expired sessions.") def get_or_create_session(session_id: str) -> dict: if session_id not in sessions: session_dir = UPLOAD_BASE / session_id store_dir = STORE_BASE / session_id assets_dir = ASSETS_BASE / session_id session_dir.mkdir(parents=True, exist_ok=True) store_dir.mkdir(parents=True, exist_ok=True) assets_dir.mkdir(parents=True, exist_ok=True) vs = FaissVectorStore(persist_dir=str(store_dir), shared_model=shared_model) search = RAGSearch(vectorstore=vs) sessions[session_id] = { "vectorstore": vs, "search": search, "upload_dir": session_dir, "assets_dir": assets_dir, "files": [], "created_at": time.time(), } session_locks[session_id] = asyncio.Lock() return sessions[session_id] @app.post("/api/session") async def create_session(): session_id = uuid.uuid4().hex[:12] get_or_create_session(session_id) return {"session_id": session_id} @app.post("/api/upload") async def upload_file( file: UploadFile = File(...), session_id: str = Form(...), ): if not session_id: raise HTTPException(400, "session_id required") ext = Path(file.filename).suffix.lower() if ext not in SUPPORTED_EXTENSIONS: raise HTTPException(400, f"Unsupported file type: {ext}. Supported: {', '.join(SUPPORTED_EXTENSIONS)}") session = get_or_create_session(session_id) if len(session["files"]) >= MAX_FILES_PER_SESSION: raise HTTPException(400, f"Maximum {MAX_FILES_PER_SESSION} files per session. Please start a new session.") content = await file.read() if len(content) > MAX_FILE_SIZE: raise HTTPException( 400, f"File too large ({len(content) / (1024*1024):.1f} MB). " f"Maximum allowed is {MAX_FILE_SIZE / (1024*1024):.0f} MB." ) dest = session["upload_dir"] / file.filename with open(dest, "wb") as f: f.write(content) lock = session_locks[session_id] async with lock: try: assets_dir = str(session["assets_dir"]) result = await asyncio.to_thread(load_single_file, str(dest), assets_dir) text_docs = result["documents"] multimodal_chunks = result["multimodal_chunks"] if text_docs: await asyncio.to_thread(session["vectorstore"].add_documents, text_docs) if multimodal_chunks: await asyncio.to_thread(session["vectorstore"].add_multimodal_chunks, multimodal_chunks) total_chunks = len(text_docs) + len(multimodal_chunks) session["files"].append({ "name": file.filename, "size": len(content), "type": ext.lstrip("."), "documents": total_chunks, }) except Exception as e: raise HTTPException(500, f"Processing failed: {str(e)}") return { "status": "success", "filename": file.filename, "documents_loaded": total_chunks, "total_files": len(session["files"]), } @app.get("/api/chat") async def chat_stream(query: str, session_id: str, top_k: int = 5): """SSE endpoint — GET so browser EventSource works natively.""" if not session_id or session_id not in sessions: raise HTTPException(400, "Invalid session") session = sessions[session_id] search = session["search"] if session["vectorstore"].index is None: raise HTTPException(400, "No documents uploaded yet. Please upload files first.") async def event_generator(): async for event in search.stream_answer(query, top_k=top_k): yield event return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering":"no", }, ) @app.get("/api/health") async def health(): """Railway healthcheck endpoint.""" return {"status": "ok", "model_loaded": shared_model is not None} @app.get("/api/files") async def list_files(session_id: str): if session_id not in sessions: return {"files": []} return {"files": sessions[session_id]["files"]} @app.get("/api/assets/{filename:path}") async def serve_asset(filename: str, session_id: str = ""): """Serve multimodal asset files (table screenshots, image thumbnails).""" if session_id and session_id in sessions: asset_path = sessions[session_id]["assets_dir"] / filename if asset_path.exists(): return FileResponse(str(asset_path)) for sid, sdata in sessions.items(): asset_path = sdata["assets_dir"] / filename if asset_path.exists(): return FileResponse(str(asset_path)) raise HTTPException(404, "Asset not found") # ─── Serve React Frontend (production) ──────────────────────────────────────── # When SERVE_FRONTEND=true (set by railway.toml), FastAPI serves the built # React app from frontend/dist/. All non-API routes fall through to index.html # so React Router works correctly. FRONTEND_DIST = Path(__file__).parent / "frontend" / "dist" if os.getenv("SERVE_FRONTEND", "false").lower() == "true" and FRONTEND_DIST.exists(): # Mount static assets (JS/CSS/images) at /assets app.mount("/assets", StaticFiles(directory=str(FRONTEND_DIST / "assets")), name="assets") @app.get("/{full_path:path}", include_in_schema=False) async def serve_react(full_path: str): """Catch-all: serve index.html so React Router handles client-side nav.""" index = FRONTEND_DIST / "index.html" if index.exists(): return FileResponse(str(index)) raise HTTPException(404, "Frontend not built. Run: cd frontend && npm run build") if __name__ == "__main__": import uvicorn uvicorn.run("app:app", host="0.0.0.0", port=int(os.getenv("PORT", 8000)), reload=False)