Spaces:
Sleeping
Sleeping
| 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 | |
| # ----------------------------- | |
| async def root(): | |
| return """ | |
| <!doctype html> | |
| <html> | |
| <head> | |
| <meta charset="utf-8" /> | |
| <title>Tessai LLM Bridge</title> | |
| <style> | |
| body { | |
| font-family: system-ui, sans-serif; | |
| padding: 2rem; | |
| } | |
| code { | |
| background: #f4f4f4; | |
| padding: 0.1rem 0.3rem; | |
| border-radius: 3px; | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <h1>Tessai LLM Bridge</h1> | |
| <p>FastAPI is running inside this Hugging Face Space.</p> | |
| <ul> | |
| <li><code>GET /health</code> – health and metrics</li> | |
| <li><code>POST /v1/chat</code> – main chat endpoint</li> | |
| <li><code>GET /admin</code> – simple meter board</li> | |
| </ul> | |
| </body> | |
| </html> | |
| """ | |
| 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 <token>" | |
| 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 | |
| # ----------------------------- | |
| 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, | |
| ) | |
| async def health(): | |
| return JSONResponse( | |
| { | |
| "status": "ok", | |
| "total_requests": metrics["total_requests"], | |
| "total_tokens": metrics["total_tokens"], | |
| "active_sessions_5m": count_active_sessions(300), | |
| } | |
| ) | |
| 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"<tr>" | |
| f"<td>{sid}</td>" | |
| f"<td>{s['requests']}</td>" | |
| f"<td>{s['tokens']}</td>" | |
| f"<td>{last_seen}</td>" | |
| f"</tr>" | |
| ) | |
| rows_html = "\n".join(rows) if rows else "<tr><td colspan='4'>No sessions yet</td></tr>" | |
| html = f""" | |
| <!doctype html> | |
| <html> | |
| <head> | |
| <title>Tessai LLM Admin</title> | |
| <meta charset="utf-8" /> | |
| <style> | |
| body {{ | |
| font-family: system-ui, sans-serif; | |
| margin: 20px; | |
| }} | |
| h1, h2 {{ | |
| margin-bottom: 0.2rem; | |
| }} | |
| .metrics {{ | |
| display: flex; | |
| gap: 1.5rem; | |
| margin-bottom: 1.5rem; | |
| }} | |
| .metric-card {{ | |
| padding: 1rem 1.5rem; | |
| border-radius: 8px; | |
| border: 1px solid #ddd; | |
| box-shadow: 0 1px 3px rgba(0,0,0,0.05); | |
| }} | |
| table {{ | |
| border-collapse: collapse; | |
| width: 100%; | |
| }} | |
| th, td {{ | |
| border: 1px solid #ddd; | |
| padding: 8px; | |
| font-size: 0.9rem; | |
| }} | |
| th {{ | |
| background: #f4f4f4; | |
| text-align: left; | |
| }} | |
| </style> | |
| </head> | |
| <body> | |
| <h1>Tessai LLM Bridge</h1> | |
| <p>Simple meter board for current traffic.</p> | |
| <div class="metrics"> | |
| <div class="metric-card"> | |
| <h2>Total requests</h2> | |
| <p>{metrics["total_requests"]}</p> | |
| </div> | |
| <div class="metric-card"> | |
| <h2>Total tokens (approx)</h2> | |
| <p>{metrics["total_tokens"]}</p> | |
| </div> | |
| <div class="metric-card"> | |
| <h2>Active sessions (last 5 min)</h2> | |
| <p>{active_5m}</p> | |
| </div> | |
| </div> | |
| <h2>Sessions</h2> | |
| <table> | |
| <thead> | |
| <tr> | |
| <th>Session ID</th> | |
| <th>Requests</th> | |
| <th>Tokens</th> | |
| <th>Last seen</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| {rows_html} | |
| </tbody> | |
| </table> | |
| </body> | |
| </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) | |