sobhandutta/agentic-personal-assistant

1"""2api.py — FastAPI app: WebSocket streaming chat + static file serving.3 4Runs alongside the existing Gradio app (app.py / main.py) — they are5independent processes sharing the same Orchestrator code.6 7Run:8    uvicorn api:app --host 0.0.0.0 --port 8000 --reload9 10Then open http://localhost:8000 for the custom UI.11Gradio remains available separately via: python main.py (port 7860)12"""13 14import logging15import traceback16import sys17from pathlib import Path18 19_LOG_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"20logging.basicConfig(format=_LOG_FORMAT, level=logging.INFO)21log = logging.getLogger("api")22 23# ── Build vector store on startup (same as app.py / main.py) ─────────────────24sys.path.insert(0, str(Path(__file__).parent))25 26def _build_vector_store():27    """Build the vector store only if it doesn't already have data.28    Rebuilding on every startup is dangerous: the collection is deleted first,29    so if the OpenAI embedding call fails, the knowledge base is gone."""30    try:31        from chromadb import PersistentClient32        from pathlib import Path as _Path33        _store_path = str(_Path(__file__).parent / "vector_store")34        _collection_name = "sobhan_knowledge_base"35        chroma = PersistentClient(path=_store_path)36        existing = [c.name for c in chroma.list_collections()]37        if _collection_name in existing:38            count = chroma.get_collection(_collection_name).count()39            if count > 0:40                log.info("Vector store already has %d chunks — skipping rebuild.", count)41                return42    except Exception:43        pass  # can't check, fall through to build attempt44 45    log.info("Building RAG vector store from knowledge_base/...")46    try:47        from data.ingest_kb import load_documents, create_chunks, embed_and_store48        documents = load_documents()49        if not documents:50            log.warning("No markdown files found in knowledge_base/ — skipping.")51            return52        chunks = create_chunks(documents)53        embed_and_store(chunks)54        log.info("Vector store ready.")55    except Exception:56        log.error("Failed to build vector store:\n%s", traceback.format_exc())57        log.warning("KnowledgeBaseAgent will be unavailable.")58 59_build_vector_store()60 61# ── FastAPI app ───────────────────────────────────────────────────────────────62from fastapi import FastAPI, WebSocket, WebSocketDisconnect63from fastapi.staticfiles import StaticFiles64from fastapi.responses import FileResponse65 66from orchestrator import Orchestrator67 68app = FastAPI(title="Personal Assistant API")69 70# Single shared Orchestrator instance — run_stream() uses only local state71# per invocation, so concurrent WebSocket connections are safe.72_orc = Orchestrator()73 74STATIC_DIR = Path(__file__).parent / "static"75 76# ── WebSocket endpoint ────────────────────────────────────────────────────────77 78@app.websocket("/ws/chat")79async def ws_chat(websocket: WebSocket):80    await websocket.accept()81    log.info("WebSocket client connected: %s", websocket.client)82    try:83        while True:84            data = await websocket.receive_json()85            user_message = data.get("message", "").strip()86            history = data.get("history", [])87 88            if not user_message:89                await websocket.send_json({"type": "error", "text": "Empty message."})90                continue91 92            log.info("WS message: %s", user_message[:120])93            async for chunk in _orc.run_stream(user_message, history):94                await websocket.send_json(chunk)95 96    except WebSocketDisconnect:97        log.info("WebSocket client disconnected: %s", websocket.client)98    except Exception:99        log.exception("Unexpected error in ws_chat()")100        try:101            await websocket.send_json({"type": "error", "text": "Server error."})102        except Exception:103            pass104 105# ── Static file serving ───────────────────────────────────────────────────────106# Root returns index.html directly; /static/ serves CSS and JS.107# Routes are registered before StaticFiles mount so /ws/chat is matched first.108 109@app.get("/")110async def index():111    return FileResponse(STATIC_DIR / "index.html")112 113app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")114