NujacsaintS's picture
Update app.py
53410da verified
Raw
History Blame Contribute Delete
8.83 kB
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 """
<!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 (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"<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)