Populate LRU caching allowing recent chat retrieval. Add memory handler

README.md

@@ -7,7 +7,7 @@ sdk: docker
 sdk_version: latest
 pinned: false
 license: apache-2.0
-short_description: Medical Chatbot, with FAISS, Gemini Flash, and MongoDB
+short_description: MedicalChatbot - FAISS RAG, Gemini Flash, MongoDB vDB, LRU cache
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,13 +1,15 @@
+# app.py
 import os
 import faiss
 import numpy as np
 import time
 import uvicorn
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 from pymongo import MongoClient
 from google import genai
 from sentence_transformers import SentenceTransformer
+from memory import MemoryManager
 
 # ✅ Enable Logging for Debugging
 import logging
@@ -20,9 +22,10 @@ for name in [
     "google", "google.auth",
 ]:
     logging.getLogger(name).setLevel(logging.WARNING)
-logging.basicConfig(level=logging.DEBUG, format="%(asctime)s — %(name)s — %(levelname)s — %(message)s", force=True)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(name)s — %(levelname)s — %(message)s", force=True) # Change INFO to DEBUG for full-ctx JSON loader
 logger = logging.getLogger("medical-chatbot")
 logger.setLevel(logging.DEBUG)
+
 # Debug Start
 logger.info("🚀 Starting Medical Chatbot API...")
 print("🚀 Starting Medical Chatbot API...")
@@ -59,6 +62,8 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # ✅ Initialize FastAPI app
 app = FastAPI(title="Medical Chatbot API")
+memory = MemoryManager()
+
 from fastapi.middleware.cors import CORSMiddleware # Bypassing CORS origin
 # Define the origins
 origins = [
@@ -153,44 +158,50 @@ class RAGMedicalChatbot:
         self.model_name = model_name
         self.retrieve = retrieve_function
 
-    def chat(self, user_query, lang="EN"):
+    def chat(self, user_id: str, user_query: str, lang: str = "EN") -> str:
+        # 1. Fetch knowledge
         retrieved_info = self.retrieve(user_query)
         knowledge_base = "\n".join(retrieved_info)
 
-        # ✅ Construct Prompt
-        prompt = f"""
-        You are a medical chatbot, designed to answer medical questions.
-        
-        Please format your answer using markdown. 
-        **Bold for titles**, *italic for emphasis*, and clear headings.
-
-        **Medical knowledge (trained with 256,916 data entries):**
-        {knowledge_base}
-
-        **Question:** {user_query}
-
-        **Language Required:** {lang}
-        """
-        completion = gemini_flash_completion(prompt, model=self.model_name, temperature=0.7)
-        return completion.strip()
+        # 2. Fetch recent context (last 3 chats)
+        context = memory.get_context(user_id, num_turns=3)
+
+        # 3. Build prompt parts
+        parts = ["You are a medical chatbot, designed to answer medical questions."]
+        parts.append("Please format your answer using MarkDown.")
+        parts.append("**Bold for titles**, *italic for emphasis*, and clear headings.")
+        # Historical chat retrieval case
+        if context:
+            parts.append(f"Previous conversation:\n{context}")
+        parts.append(f"Medical knowledge (256,916 medical scenario): {knowledge_base}")
+        parts.append(f"Question: {user_query}")
+        parts.append(f"Language: {lang}")
+        prompt = "\n\n".join(parts)
+        response = gemini_flash_completion(prompt, model=self.model_name, temperature=0.7)
+        # Add to STM caching by id+req+res
+        if user_id:
+            memory.add_exchange(user_id, user_query, response)
+        return response.strip()
 
 # ✅ Initialize Chatbot
 chatbot = RAGMedicalChatbot(model_name="gemini-2.5-flash-preview-04-17", retrieve_function=retrieve_medical_info)
 
 # ✅ Chat Endpoint
 @app.post("/chat")
-async def chat_endpoint(data: dict):
-    user_query = data.get("query", "")
-    lang = data.get("lang", "EN")
-    if not user_query:
-        return JSONResponse(content={"response": "No query provided."})
-    # Output parameter
-    start_time = time.time()
-    response_text = chatbot.chat(user_query, lang)
-    end_time = time.time()
-    response_text += f"\n\n(Response time: {end_time - start_time:.2f} seconds)"
-    # Send JSON response
-    return JSONResponse(content={"response": response_text})
+async def chat_endpoint(req: Request):
+    body = await req.json()
+    user_id = body.get("user_id", "anonymous")
+    query   = body.get("query", "").strip()
+    lang    = body.get("lang", "EN")
+    # Error
+    if not query:
+        return JSONResponse({"response": "No query provided."})
+    start = time.time()
+    answer = chatbot.chat(user_id, query, lang)
+    elapsed = time.time() - start
+    # Final
+    return JSONResponse({"response": f"{answer}\n\n(Response time: {elapsed:.2f}s)"})
+
 
 # ✅ Run Uvicorn
 if __name__ == "__main__":

memory.py

@@ -0,0 +1,43 @@
+# memory.py
+
+from collections import defaultdict, deque
+from typing import List, Tuple
+
+class MemoryManager:
+    """
+    In-memory LRU cache of short-term history:
+      • max_users: total distinct user_ids cached
+      • history_per_user: number of exchanges to keep per user
+    """
+    def __init__(self, max_users: int = 1000, history_per_user: int = 10):
+        self.memory     = defaultdict(lambda: deque(maxlen=history_per_user))
+        self.user_queue = deque(maxlen=max_users)
+
+    def add_exchange(self, user_id: str, query: str, response: str):
+        """
+        Append a (query, response) pair to the user's history.
+        Evict oldest users when capacity is reached.
+        """
+        if user_id not in self.memory:
+            # If cache full, drop oldest user and all their history
+            if len(self.user_queue) >= self.user_queue.maxlen:
+                oldest = self.user_queue.popleft()
+                del self.memory[oldest]
+            self.user_queue.append(user_id)
+
+        self.memory[user_id].append((query.strip(), response.strip()))
+
+    def get_context(self, user_id: str, num_turns: int = 3) -> str:
+        """
+        Return the last `num_turns` as a single formatted string:
+          User: question
+          Bot:    answer
+        """
+        turns = list(self.memory.get(user_id, []))[-num_turns:]
+        return "\n".join(f"User: {q}\nBot: {r}" for q, r in turns)
+
+    def reset(self, user_id: str):
+        """Clear all history for a given user_id."""
+        if user_id in self.memory:
+            del self.memory[user_id]
+            self.user_queue.remove(user_id)