Update app.py

app.py

@@ -5,6 +5,8 @@ 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")
 
@@ -95,76 +97,61 @@ def count_active_sessions(window_seconds: int = 300) -> int:
 
 
 # -----------------------------
-# LLM integration
+# LLM integration (local model)
 # -----------------------------
-# Configure via environment variables in Hugging Face:
-#   TESSAI_LLM_URL  = full URL of the model endpoint
-#   TESSAI_LLM_KEY  = auth token (if needed)
+# 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")
 
-LLM_URL = os.getenv("TESSAI_LLM_URL", "").strip()
-LLM_KEY = os.getenv("TESSAI_LLM_KEY", "").strip()
+print(f"Loading local model: {MODEL_NAME}")
+_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+_model.eval()  # inference mode
 
-
-from httpx import HTTPStatusError
-
-async def call_llm(message: str, context: Dict[str, str] | None = None) -> tuple[str, int]:
-    if not LLM_URL:
-        reply = f"[stub] Echo: {message}"
-        tokens_used = len(message)
-        return reply, tokens_used
-
-    prompt = message
+@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\n{message}"
-
-    headers = {}
-    if LLM_KEY:
-        headers["Authorization"] = f"Bearer {LLM_KEY}"
-
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 256,
-            "temperature": 0.3,
-            "do_sample": False,
-        },
-    }
+        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,
+    )
 
-    try:
-        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()
-    except HTTPStatusError as e:
-        # Surface *what HF said* back to the front end instead of 500
-        body_text = ""
-        try:
-            body_text = e.response.text
-        except Exception:
-            pass
-
-        reply = (
-            f"[bridge error] HF returned {e.response.status_code} for {LLM_URL}. "
-            f"Response body: {body_text}"
-        )
-        return reply, 0
-    except Exception as e:
-        reply = f"[bridge error] Unexpected exception calling HF: {e}"
-        return reply, 0
+    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()
 
-    # Normal happy path
-    if isinstance(data, list) and data and "generated_text" in data[0]:
-        full_text = data[0]["generated_text"]
-        reply = full_text[len(prompt):].strip() or full_text.strip()
-    else:
-        reply = str(data)
+    # Approx tokens used in the reply
+    reply_tokens = _tokenizer.encode(reply)
+    tokens_used = len(reply_tokens)
 
-    tokens_used = len(reply.split())
     return reply, tokens_used
 
 
+async def call_llm(message: str, context: Dict[str, str] | None = None) -> tuple[str, int]:
+    """
+    Main hook used by /v1/chat. Currently calls the local model.
+    Later you can swap this to call your custom notebook LLM or an external API.
+    """
+    try:
+        return generate_local_reply(message, context)
+    except Exception as e:
+        # Last resort: don't crash the bridge, just report an error message
+        err = f"[local-llm error] {type(e).__name__}: {e}"
+        print(err)
+        return err, 0
+
 
 # -----------------------------
 # API endpoints