Vaibhav Kathait
Add en/it/fr/es/de language support; default to our own agent LLM instead of Anam's
7f73287 | 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") | |
| 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")), | |
| }) | |
| 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()) | |
| 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()) | |
| 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 | |
| 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}) | |
| 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), | |
| }) | |
| 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}) | |