from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel
from typing import Dict
import time
import os
import httpx
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
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 (local model)
# -----------------------------
# This uses a small local model so we don't depend on HF router / inference URLs.
# You can later swap "gpt2" for your own model or your notebook code.
MODEL_NAME = os.getenv("TESSAI_LOCAL_MODEL", "gpt2")
print(f"Loading local model: {MODEL_NAME}")
_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
_model.eval() # inference mode
@torch.inference_mode()
def generate_local_reply(message: str, context: Dict[str, str] | None = None) -> tuple[str, int]:
# Simple prompt format; you can replace with your notebook's prompt engineering
if context:
ctx_str = "; ".join(f"{k}={v}" for k, v in context.items())
prompt = f"[context: {ctx_str}]\n\nUser: {message}\nAssistant:"
else:
prompt = f"User: {message}\nAssistant:"
inputs = _tokenizer(prompt, return_tensors="pt")
outputs = _model.generate(
**inputs,
max_new_tokens=128,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=_tokenizer.eos_token_id,
)
full_text = _tokenizer.decode(outputs[0], skip_special_tokens=True)
# Heuristic: the reply is whatever came after the prompt
reply = full_text[len(prompt):].strip() or full_text.strip()
# Approx tokens used in the reply
reply_tokens = _tokenizer.encode(reply)
tokens_used = len(reply_tokens)
return reply, tokens_used
import openai, os
openai.api_key = os.getenv("OPENAI_API_KEY")
async def call_llm(message: str, context: Dict[str, str] | None = None):
reply = f"Echo: {len(message)} boing flip"
tokens_used = len(message)
return reply, tokens_used
"""
prompt = f"User: {message}\nAssistant:"
try:
completion = openai.chat.completions.create(
model="gpt-4o-mini", # or gpt-5-nano, or gpt-5.1-chat-latest
messages=[
{"role": "system", "content": "You are Tessai, an analytical change-management agent."},
{"role": "user", "content": message},
],
max_tokens=300
)
reply = completion.choices[0].message["content"]
tokens_used = completion.usage.total_tokens
return reply, tokens_used
except Exception as e:
return f"[openai error] {e}", 0
"""
from openai import OpenAI
client = OpenAI()
# For estimation you’d use a tokenizer helper, *not* an API call.
from tiktoken import get_encoding
enc = get_encoding("o200k_base") # for GPT-4.x / 5.x style models
def estimate_tokens_for_messages(messages):
text = ""
for msg in messages:
# Simplest concat; you can apply more exact rules later
text += f"{msg['role']}: {msg['content']}\n"
return len(enc.encode(text))
# -----------------------------
# 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)