from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel
from typing import Dict
import time
import os
import httpx
app = FastAPI(title="Tessai LLM Bridge", version="0.1.0")
# -----------------------------
# Models
# -----------------------------
@app.get("/", response_class=HTMLResponse)
async def root():
return """
Tessai LLM Bridge
Tessai LLM Bridge
FastAPI is running inside this Hugging Face Space.
GET /health – health and metricsPOST /v1/chat – main chat endpointGET /admin – simple meter board
"""
class ChatRequest(BaseModel):
session_id: str
message: str
context: Dict[str, str] | None = None
class ChatResponse(BaseModel):
session_id: str
reply: str
tokens_used: int
# -----------------------------
# Simple in-memory metrics
# -----------------------------
metrics = {
"total_requests": 0,
"total_tokens": 0,
"sessions": {} # session_id -> {"last_seen": float, "requests": int, "tokens": int}
}
def record_request(session_id: str, tokens_used: int) -> None:
now = time.time()
metrics["total_requests"] += 1
metrics["total_tokens"] += tokens_used
if session_id not in metrics["sessions"]:
metrics["sessions"][session_id] = {
"last_seen": now,
"requests": 0,
"tokens": 0,
}
s = metrics["sessions"][session_id]
s["last_seen"] = now
s["requests"] += 1
s["tokens"] += tokens_used
def count_active_sessions(window_seconds: int = 300) -> int:
now = time.time()
return sum(
1
for s in metrics["sessions"].values()
if now - s["last_seen"] <= window_seconds
)
# -----------------------------
# LLM integration
# -----------------------------
# Configure via environment variables in Hugging Face:
# TESSAI_LLM_URL = full URL of the model endpoint
# TESSAI_LLM_KEY = auth token (if needed)
LLM_URL = os.getenv("TESSAI_LLM_URL", "").strip()
LLM_KEY = os.getenv("TESSAI_LLM_KEY", "").strip()
async def call_llm(message: str, context: Dict[str, str] | None = None) -> tuple[str, int]:
"""
Replace this with whatever target you want.
This implementation is ready for a typical Hugging Face Inference API
text-generation endpoint. If LLM_URL is not set, it falls back to a local echo.
"""
# Fallback so the bridge still works even if you haven't configured the model yet
if not LLM_URL:
reply = f"[stub] Echo: {message}"
tokens_used = len(message)
return reply, tokens_used
# Basic prompt construction. You can get fancier later.
prompt = message
if context:
# Simple way to inject context: a header line then the message.
ctx_str = "; ".join(f"{k}={v}" for k, v in context.items())
prompt = f"[context: {ctx_str}]\n\n{message}"
headers = {}
if LLM_KEY:
# Hugging Face style: "Bearer "
headers["Authorization"] = f"Bearer {LLM_KEY}"
payload = {
"inputs": prompt,
# Optional: adjust parameters to your model
"parameters": {
"max_new_tokens": 256,
"temperature": 0.3,
"do_sample": False,
}
}
async with httpx.AsyncClient(timeout=60.0) as client:
resp = await client.post(LLM_URL, headers=headers, json=payload)
resp.raise_for_status()
data = resp.json()
# Hugging Face text-generation outputs are usually a list of dicts with "generated_text"
# Adjust here if your model returns a different shape.
if isinstance(data, list) and data and "generated_text" in data[0]:
full_text = data[0]["generated_text"]
# Heuristic: reply is the part after the original prompt
reply = full_text[len(prompt) :].strip() or full_text.strip()
else:
# Fallback: just stringify whatever we got
reply = str(data)
# Approx token count: you can replace this with a real tokenizer later
tokens_used = len(reply.split())
return reply, tokens_used
# -----------------------------
# API endpoints
# -----------------------------
@app.post("/v1/chat", response_model=ChatResponse)
async def chat_endpoint(payload: ChatRequest):
reply, tokens_used = await call_llm(payload.message, payload.context)
record_request(payload.session_id, tokens_used)
return ChatResponse(
session_id=payload.session_id,
reply=reply,
tokens_used=tokens_used,
)
@app.get("/health")
async def health():
return JSONResponse(
{
"status": "ok",
"total_requests": metrics["total_requests"],
"total_tokens": metrics["total_tokens"],
"active_sessions_5m": count_active_sessions(300),
}
)
@app.get("/admin", response_class=HTMLResponse)
async def admin_dashboard():
active_5m = count_active_sessions(300)
rows = []
for sid, s in metrics["sessions"].items():
last_seen = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(s["last_seen"]))
rows.append(
f""
f"| {sid} | "
f"{s['requests']} | "
f"{s['tokens']} | "
f"{last_seen} | "
f"
"
)
rows_html = "\n".join(rows) if rows else "| No sessions yet |
"
html = f"""
Tessai LLM Admin
Tessai LLM Bridge
Simple meter board for current traffic.
Total requests
{metrics["total_requests"]}
Total tokens (approx)
{metrics["total_tokens"]}
Active sessions (last 5 min)
{active_5m}
Sessions
| Session ID |
Requests |
Tokens |
Last seen |
{rows_html}
"""
return HTMLResponse(content=html)
# -----------------------------
# Local dev entrypoint
# -----------------------------
if __name__ == "__main__":
import uvicorn
uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)