Update app.py

app.py

@@ -1,68 +1,113 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import JSONResponse, HTMLResponse
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from tinydb import TinyDB, Query
+import os
+import logging
 import uvicorn
 
-app = FastAPI()
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Load model
-tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5")
-model = AutoModelForCausalLM.from_pretrained("microsoft/phi-1_5")
-model.eval()
+# Base path from Hugging Face Spaces
+BASE_PATH = os.getenv("SPACE_APP_PATH", "").rstrip("/")
+logger.info(f"Using base path: '{BASE_PATH}'")
 
-# Memory setup
-db = TinyDB("memory.json")
-User = Query()
+# Initialize app
+app = FastAPI(title="Trigger AI", description="Lightning fast chatbot", version="1.0")
+
+# Load lightweight fast model (phi-1.5)
+try:
+    logger.info("Loading tokenizer and model...")
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5")
+    model = AutoModelForCausalLM.from_pretrained("microsoft/phi-1_5")
+    logger.info("Model loaded.")
+except Exception as e:
+    logger.error(f"Model load error: {e}")
+    raise RuntimeError("Model failed to load")
+
+# In-memory chat memory per user_id
+chat_memory = {}
+
+@app.middleware("http")
+async def strip_base_path(request: Request, call_next):
+    path = request.scope["path"]
+    if BASE_PATH and path.startswith(BASE_PATH):
+        request.scope["path"] = path[len(BASE_PATH):]
+    return await call_next(request)
 
 @app.get("/")
-async def chat(query: str = "", number: str = "", prompt: str = ""):
-    if not query or not number:
-        return {"response": "Missing query or number."}
-
-    record = db.get(User.number == number)
-    if not record:
-        record = {
-            "number": number,
-            "history": [],
-            "prompt": "You are Trigger, a clever, emotional AI boy created by 'I am him'. Speak like a witty, chill human friend."
-        }
-        db.insert(record)
-
-    if prompt:
-        record["prompt"] = prompt
-        db.update(record, User.number == number)
-
-    history = record["history"]
-    system_prompt = record.get("prompt", "")
-
-    # Append new input
-    history.append(f"User: {query}")
-    trimmed = history[-6:]
-
-    # Construct final prompt
-    full_prompt = system_prompt + "\n" + "\n".join(trimmed) + "\nTrigger:"
-    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True)
-
-    with torch.no_grad():
+async def root():
+    return {
+        "message": "✅ Trigger AI is active",
+        "try": f"{BASE_PATH}/ai?query=Hello&user_id=233XXXXXXXXX"
+    }
+
+@app.get("/ai")
+async def ai(request: Request):
+    query = request.query_params.get("query", "").strip()
+    user_id = request.query_params.get("user_id", "").strip()
+
+    if not query or not user_id:
+        raise HTTPException(status_code=400, detail="Missing 'query' or 'user_id'")
+
+    try:
+        # Tokenize input
+        input_ids = tokenizer.encode(query, return_tensors="pt")
+
+        # Load history
+        history = chat_memory.get(user_id, [])
+        full_input = torch.cat(history + [input_ids], dim=-1) if history else input_ids
+
+        # Generate response
         output = model.generate(
-            **inputs,
+            full_input,
             max_new_tokens=100,
-            temperature=0.9,
-            top_p=0.9,
             do_sample=True,
+            top_k=40,
+            top_p=0.9,
+            temperature=0.8,
             pad_token_id=tokenizer.eos_token_id
         )
 
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    reply = decoded.split("Trigger:")[-1].strip()
+        # Decode result
+        response = tokenizer.decode(output[:, full_input.shape[-1]:][0], skip_special_tokens=True)
+
+        # Save memory
+        chat_memory[user_id] = [full_input, output]
+
+        return {"reply": response}
+
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/reset")
+async def reset(user_id: str = "default"):
+    if user_id in chat_memory:
+        del chat_memory[user_id]
+    return {"status": "cleared", "user_id": user_id}
 
-    # Save memory
-    history.append(f"Trigger: {reply}")
-    db.update({"number": number, "history": history, "prompt": system_prompt}, User.number == number)
+@app.get("/health")
+async def health():
+    return {
+        "status": "🟢 online",
+        "users": len(chat_memory),
+        "model": "phi-1.5",
+        "base_path": BASE_PATH
+    }
 
-    return {"response": reply}
+@app.get("/test", response_class=HTMLResponse)
+async def test():
+    return f"""
+    <html>
+        <body>
+            <h2>Trigger AI Test</h2>
+            <a href="{BASE_PATH}/ai?query=Hello&user_id=tester">Talk to Trigger</a>
+        </body>
+    </html>
+    """
 
-# Only runs locally, HF handles it differently
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)