api_server.py

import uuid
import asyncio
from pathlib import Path

from fastapi import FastAPI
from fastapi.responses import StreamingResponse, FileResponse
from pydantic import BaseModel

import chatbot_core as core

app = FastAPI(title="HBL Internal Assistant API")

# ---------------------------------------------------------------------------
# Serve the chat UI at the root URL — anyone on the network who visits
# http://<server-ip>:8000/ gets the interface, no separate hosting needed.
# ---------------------------------------------------------------------------
UI_PATH = Path(__file__).parent / "index.html"

@app.get("/")
async def serve_ui():
    return FileResponse(UI_PATH)

# ---------------------------------------------------------------------------
# Request / response schemas
# ---------------------------------------------------------------------------
class ChatRequest(BaseModel):
    session_id: str | None = None   # if omitted, a new session is created
    message: str
    max_new_tokens: int | None = None
    # Note: this endpoint streams plain text chunks, not JSON — session_id for a
    # NEW session is generated server-side and won't be visible to the client
    # unless you also call /session/new first and pass it in explicitly.

# ---------------------------------------------------------------------------
# Streaming endpoint — sends words as they're generated (word-by-word, like
# the TextStreamer version in the terminal script, but over HTTP). This is
# the only chat endpoint — non-streaming was removed since waiting on a full
# generation before showing anything made inference feel slower than it is.
# ---------------------------------------------------------------------------
@app.post("/chat")
async def chat(req: ChatRequest):
    session_id = req.session_id or str(uuid.uuid4())
    history = core.get_history(session_id)

    if core.contains_code(req.message) or core.contains_math(req.message):
        async def blocked():
            yield "I can only help with HBL-related questions or professional writing — not code or math."
        return StreamingResponse(blocked(), media_type="text/plain")

    if core.is_writing_task(req.message):
        retrieved = []
    else:
        retrieved = core.retrieve(req.message)

    context = "\n\n".join(f"[{r['source_url']}]\n{r['text']}" for r in retrieved) if retrieved else ""
    system_prompt = core.UNIFIED_SYSTEM_PROMPT.format(context=context)

    async def token_generator():
        full_response = ""
        # generation_lock is acquired here, for the whole streamed generation —
        # this is what serializes concurrent users on the single GPU model.
        async with core.generation_lock:
            async for chunk in core.stream_llm_with_history(
                system_prompt, history, req.message, req.max_new_tokens
            ):
                full_response += chunk
                yield chunk

        cleaned = core.strip_or_block(full_response.strip())
        history.append({"role": "user", "content": req.message})
        history.append({"role": "assistant", "content": cleaned})

    return StreamingResponse(token_generator(), media_type="text/plain")

# ---------------------------------------------------------------------------
# Session management
# ---------------------------------------------------------------------------
@app.post("/session/new")
async def new_session():
    session_id = str(uuid.uuid4())
    core.get_history(session_id)  # initializes an empty history entry
    return {"session_id": session_id}

@app.post("/session/{session_id}/reset")
async def reset_session(session_id: str):
    core.reset_history(session_id)
    return {"status": "ok", "session_id": session_id}

@app.get("/health")
async def health():
    return {"status": "ok", "device": core.DEVICE}