Update app.py

app.py

@@ -2,7 +2,11 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from pydantic import BaseModel
+from collections import defaultdict
+from threading import Lock
 import uvicorn
+import uuid
+import time
 
 app = FastAPI()
 
@@ -31,20 +35,53 @@ SYSTEM_PROMPT = (
     "Keep your replies short, casual, and to the point."
 )
 
+# --- In-memory session history store ---------------------------------------
+# NOTE: this resets whenever the process restarts. For real persistence across
+# restarts/multiple server instances, swap this dict for Firestore/Supabase
+# (you already use both elsewhere in Bella AI) — the shape below (a list of
+# {"role", "content"} dicts per session_id) maps directly onto either.
+MAX_TURNS = 12          # how many user+assistant turn-pairs to keep per session
+SESSION_TTL_SECONDS = 60 * 60 * 6  # drop sessions idle for 6+ hours
+
+_histories: dict[str, list[dict]] = defaultdict(list)
+_last_seen: dict[str, float] = {}
+_lock = Lock()
+
+
+def _prune_stale_sessions():
+    now = time.time()
+    stale = [sid for sid, ts in _last_seen.items() if now - ts > SESSION_TTL_SECONDS]
+    for sid in stale:
+        _histories.pop(sid, None)
+        _last_seen.pop(sid, None)
+
 
 class ChatRequest(BaseModel):
     message: str
+    session_id: str | None = None  # client can pass one to resume a conversation
+
+
+class ChatResponse(BaseModel):
+    response: str
+    session_id: str
 
 
-@app.post("/chat")
+@app.post("/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
-    # Priming turn anchors the persona much more strongly for small models
-    messages = [
-        {"role": "system", "content": SYSTEM_PROMPT},
-        {"role": "user", "content": "Who are you?"},
-        {"role": "assistant", "content": "I'm Bella! What can I help you with?"},
-        {"role": "user", "content": request.message},
-    ]
+    with _lock:
+        _prune_stale_sessions()
+
+        session_id = request.session_id or str(uuid.uuid4())
+        history = _histories[session_id]
+
+        # Build the message list: system + priming turn + saved history + new message
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": "Who are you?"},
+            {"role": "assistant", "content": "I'm Bella! What can I help you with?"},
+            *history,
+            {"role": "user", "content": request.message},
+        ]
 
     prompt = tokenizer.apply_chat_template(
         messages,
@@ -75,8 +112,26 @@ async def chat(request: ChatRequest):
     if not response:
         response = "Hmm, I didn't quite catch that — can you rephrase?"
 
-    return {"response": response}
+    with _lock:
+        history = _histories[session_id]
+        history.append({"role": "user", "content": request.message})
+        history.append({"role": "assistant", "content": response})
+        # keep only the last MAX_TURNS turn-pairs
+        if len(history) > MAX_TURNS * 2:
+            _histories[session_id] = history[-MAX_TURNS * 2:]
+        _last_seen[session_id] = time.time()
+
+    return {"response": response, "session_id": session_id}
+
+
+@app.post("/chat/reset")
+async def reset(session_id: str):
+    with _lock:
+        _histories.pop(session_id, None)
+        _last_seen.pop(session_id, None)
+    return {"status": "ok"}
 
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)
+