server.py

"""
ChatJio — FastAPI server.
Exposes the RAG pipeline over HTTP so any frontend can connect.

Run:
    python server.py

Endpoints:
    POST /chat   { "query": "...", "confirmed": false }
                 -> { "needs_confirmation": true, "message": "..." }   (when no/partial match)
                 -> { "needs_confirmation": false, "answer": "...", "sources": [...] }  (when answered)
    GET  /health -> { "status": "ok" }
"""

import logging
import os
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

logging.basicConfig(level=logging.WARNING, format="%(levelname)s | %(name)s | %(message)s")

from retrieval.retriever import retrieve
from generation.generator import Generator, GenerationError

app = FastAPI(title="ChatJio API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

_generator = None


def get_generator() -> Generator:
    global _generator
    if _generator is None:
        _generator = Generator()
    return _generator


class ChatRequest(BaseModel):
    query: str
    confirmed: bool = False


class ChatResponse(BaseModel):
    needs_confirmation: bool = False
    message: str = ""
    query: str = ""
    answer: str = ""
    sources: list[str] = []
    match_status: str = ""


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/chat", response_model=ChatResponse)
def chat(request: ChatRequest):
    query = request.query.strip()
    if not query:
        raise HTTPException(status_code=400, detail="query must not be empty")

    try:
        chunks, match_status = retrieve(query)
        generator = get_generator()

        if match_status == "good":
            result = generator.generate(query, chunks)
            return ChatResponse(
                needs_confirmation=False,
                query=result["query"],
                answer="Found it in the DB.\n\n" + result["answer"],
                sources=list(dict.fromkeys(result["sources"])),
                match_status=match_status,
            )

        # no_match or partial — ask user first unless already confirmed
        if not request.confirmed:
            return ChatResponse(
                needs_confirmation=True,
                message="Sorry, could not find anything on this topic in the stored DB. Do you want me to fetch it outside the DB?",
                match_status=match_status,
            )

        # User confirmed — use LLM
        result = generator.generate(query, [])
        return ChatResponse(
            needs_confirmation=False,
            query=result["query"],
            answer="No results found in DB, seeking external help.\n\n" + result["answer"],
            sources=[],
            match_status=match_status,
        )

    except GenerationError as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)