File size: 3,830 Bytes
3af8afc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | 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} |