router_chat.py

from fastapi import APIRouter
from pydantic import BaseModel, Field
from typing import List, Dict
from rag import get_answer
from faq_store import FAQ_ENTRIES

router = APIRouter()

# Category ordering and ID-prefix mapping
_CATEGORY_ORDER = ["App Features", "Lab Values", "Medications", "General Medical"]
_PREFIX_TO_CATEGORY = {
    "faq-app-": "App Features",
    "faq-lab-": "Lab Values",
    "faq-med-": "Medications",
    "faq-gen-": "General Medical",
}


class ChatRequest(BaseModel):
    question: str = Field(..., min_length=1, max_length=500)
    history: List[Dict[str, str]] = []


class Citation(BaseModel):
    id: str
    source: str


class ChatResponse(BaseModel):
    answer: str
    citations: List[Citation]


@router.post("/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
    answer, citations = await get_answer(request.question, request.history)
    return ChatResponse(
        answer=answer,
        citations=[Citation(**c) for c in citations],
    )


@router.get("/faqs")
async def get_faqs():
    """Return all FAQ entries grouped by category, in canonical order."""
    buckets: dict[str, list] = {cat: [] for cat in _CATEGORY_ORDER}
    for fid, entry in FAQ_ENTRIES.items():
        for prefix, category in _PREFIX_TO_CATEGORY.items():
            if fid.startswith(prefix):
                buckets[category].append({
                    "id": entry["id"],
                    "question": entry["question"],
                    "answer": entry["answer"],
                    "source": entry["source"],
                })
                break
    return [
        {"category": cat, "entries": entries}
        for cat, entries in buckets.items()
        if entries
    ]