import json from pathlib import Path from typing import Optional import httpx from fastapi import APIRouter, HTTPException, UploadFile, File from fastapi.responses import JSONResponse from pydantic import BaseModel from config import ( BASE, ANAM_API_KEY, DEFAULT_AVATAR_ID, DEFAULT_VOICE_ID, DEFAULT_LLM_ID, DEFAULT_LANGUAGE_CODE, RAG_SYSTEM_PROMPT, SUPPORTED_LANGUAGES, ) from helpers import auth_headers, check # ── Persistent state ────────────────────────────────────────────────────────── # Stored in rag_state.json so IDs survive server restarts. # Resets require deleting that file or calling POST /api/rag/reset. RAG_STATE_FILE = Path(__file__).parent / "rag_state.json" def _load_state() -> dict: if RAG_STATE_FILE.exists(): try: return json.loads(RAG_STATE_FILE.read_text()) except Exception: pass return {"group_id": None, "tool_id": None, "persona_id": None} def _save_state() -> None: RAG_STATE_FILE.write_text(json.dumps(rag_state, indent=2)) rag_state: dict = _load_state() async def ensure_rag_setup() -> None: """ Lazily create, in order: 1. A knowledge group to hold uploaded documents. 2. A SERVER_RAG tool that searches that group. 3. A persona with the tool attached (so sessions can use personaId). IDs are persisted to rag_state.json after each step. """ async with httpx.AsyncClient(timeout=30) as c: if not rag_state["group_id"]: r = await c.post( f"{BASE}/knowledge/groups", headers=auth_headers(), json={"name": "Demo Knowledge Base", "description": "Uploaded documents for RAG demo"}, ) check(r) rag_state["group_id"] = r.json()["id"] _save_state() if not rag_state["tool_id"]: r = await c.post( f"{BASE}/tools", headers=auth_headers(), json={ "name": "search_knowledge_base", "description": ( "Search the knowledge base for relevant information to answer user questions. " "Use this whenever the user asks about content from their documents." ), "type": "SERVER_RAG", "config": {"documentFolderIds": [rag_state["group_id"]]}, }, ) check(r) rag_state["tool_id"] = r.json()["id"] _save_state() if not rag_state["persona_id"]: persona_payload = { "name": "RAG Demo Assistant", "avatarId": DEFAULT_AVATAR_ID, "voiceId": DEFAULT_VOICE_ID, "llmId": DEFAULT_LLM_ID, "systemPrompt": RAG_SYSTEM_PROMPT, "toolIds": [rag_state["tool_id"]], } if DEFAULT_LANGUAGE_CODE: persona_payload["languageCode"] = DEFAULT_LANGUAGE_CODE r = await c.post( f"{BASE}/personas", headers=auth_headers(), json=persona_payload, ) check(r) rag_state["persona_id"] = r.json()["id"] _save_state() # ── RAG router ──────────────────────────────────────────────────────────────── router = APIRouter(prefix="/api/rag") @router.get("/status") async def rag_status(): return JSONResponse({ "group_id": rag_state.get("group_id"), "tool_id": rag_state.get("tool_id"), "persona_id": rag_state.get("persona_id"), "ready": bool(rag_state.get("persona_id")), }) @router.post("/upload") async def rag_upload(file: UploadFile = File(...)): """Upload a document to the knowledge group. Auto-creates group/tool/persona on first call.""" await ensure_rag_setup() content = await file.read() async with httpx.AsyncClient(timeout=120) as c: r = await c.post( f"{BASE}/knowledge/groups/{rag_state['group_id']}/documents", headers={"Authorization": f"Bearer {ANAM_API_KEY}"}, files={"file": (file.filename, content, file.content_type or "application/octet-stream")}, data={"chunkSize": "1000", "chunkOverlap": "200"}, ) check(r) return JSONResponse(r.json()) @router.get("/documents") async def rag_documents(): """List all documents in the knowledge group.""" if not rag_state.get("group_id"): return JSONResponse({"items": []}) async with httpx.AsyncClient(timeout=30) as c: r = await c.get( f"{BASE}/knowledge/groups/{rag_state['group_id']}/documents", headers=auth_headers(), ) check(r) return JSONResponse(r.json()) @router.delete("/documents/{doc_id}") async def rag_delete_document(doc_id: str): """Delete a single document from the knowledge group.""" async with httpx.AsyncClient(timeout=30) as c: r = await c.delete( f"{BASE}/knowledge/documents/{doc_id}", headers=auth_headers(), ) check(r) return JSONResponse({"ok": True}) class UpdatePersonaRequest(BaseModel): system_prompt: str # Optional — one of config.SUPPORTED_LANGUAGES. Omit to leave the # persona's current language untouched. language_code: Optional[str] = None @router.put("/persona") async def update_rag_persona(req: UpdatePersonaRequest): """Update the RAG persona's system prompt (and optionally language) on the fly. No need to reconnect — Anam applies persona edits to the next turn. """ if not rag_state.get("persona_id"): raise HTTPException(status_code=404, detail="RAG persona not set up yet — upload a document first.") if req.language_code and req.language_code not in SUPPORTED_LANGUAGES: raise HTTPException( status_code=400, detail=f"Unsupported language_code '{req.language_code}'. Supported: {sorted(SUPPORTED_LANGUAGES)}", ) payload: dict = {"systemPrompt": req.system_prompt} if req.language_code: payload["languageCode"] = req.language_code async with httpx.AsyncClient(timeout=30) as c: r = await c.put( f"{BASE}/personas/{rag_state['persona_id']}", headers=auth_headers(), json=payload, ) check(r) return JSONResponse({"ok": True}) @router.get("/persona") async def get_rag_persona(): """Return the current RAG persona's system prompt and language.""" if not rag_state.get("persona_id"): return JSONResponse({"system_prompt": RAG_SYSTEM_PROMPT, "language_code": DEFAULT_LANGUAGE_CODE}) async with httpx.AsyncClient(timeout=30) as c: r = await c.get(f"{BASE}/personas/{rag_state['persona_id']}", headers=auth_headers()) check(r) data = r.json() return JSONResponse({ "system_prompt": data.get("systemPrompt", RAG_SYSTEM_PROMPT), "language_code": data.get("languageCode", DEFAULT_LANGUAGE_CODE), }) @router.post("/reset") async def rag_reset(): """ Clear local RAG state (group/tool/persona IDs). Does NOT delete resources from your Anam account — do that in the Anam dashboard if needed. """ rag_state["group_id"] = None rag_state["tool_id"] = None rag_state["persona_id"] = None _save_state() return JSONResponse({"ok": True})