Spaces:

NujacsaintS
/

CorporateChatAPI

Runtime error

File size: 8,832 Bytes

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)