Update app.py

app.py

@@ -15,6 +15,7 @@ app = Flask(__name__)
 LAMBDA_URL = os.getenv("LAMBDA_URL", "https://your-lambda-function-url")
 GEMINI_KEY = os.getenv("GEMINI_API_KEY", "")
 FLUSH_INTERVAL = 30  # seconds between DB backups per user
+MAX_HISTORY_TURNS = 10  # Maximum conversation turns to keep in context
 
 client = genai.Client(api_key=GEMINI_KEY)
 user_memory = {}  # { user_id: { "history": [], "last_sync": timestamp } }
@@ -26,10 +27,12 @@ def flush_loop():
         for uid, data in list(user_memory.items()):
             if now - data.get("last_sync", 0) >= FLUSH_INTERVAL and data["history"]:
                 try:
-                    payload = {"user_id": uid, "history": data["history"]}
+                    # Only sync the most recent MAX_HISTORY_TURNS entries
+                    history_to_sync = data["history"][-MAX_HISTORY_TURNS:]
+                    payload = {"user_id": uid, "history": history_to_sync}
                     requests.post(LAMBDA_URL, json=payload, timeout=5)
                     user_memory[uid]["last_sync"] = now
-                    app.logger.info(f"Synced memory for {uid}")
+                    app.logger.info(f"Synced memory for {uid} ({len(history_to_sync)} turns)")
                 except Exception as e:
                     app.logger.warning(f"Failed sync for {uid}: {e}")
         time.sleep(5)
@@ -86,24 +89,51 @@ HTML = """
 </body></html>
 """
 
-# --- Gemini Generation ---
-def generate_from_gemini(prompt, image_bytes=None):
+# --- Gemini Generation with History ---
+def generate_from_gemini(prompt, image_bytes=None, history=None):
     start_time = time.time()
-    parts = []
-    if prompt: parts.append(types.Part.from_text(text=prompt))
-    if image_bytes: parts.append(types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg"))
-
-    contents = [types.Content(role="user", parts=parts)]
+    
+    # Build contents list with history
+    contents = []
+    
+    # Add historical messages (limit to recent turns to avoid token limits)
+    if history:
+        recent_history = history[-MAX_HISTORY_TURNS:]
+        for entry in recent_history:
+            # Add user message
+            user_parts = [types.Part.from_text(text=entry["prompt"])]
+            contents.append(types.Content(role="user", parts=user_parts))
+            
+            # Add model response
+            model_parts = [types.Part.from_text(text=entry["response"])]
+            contents.append(types.Content(role="model", parts=model_parts))
+    
+    # Add current user message
+    current_parts = []
+    if prompt: 
+        current_parts.append(types.Part.from_text(text=prompt))
+    if image_bytes: 
+        current_parts.append(types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg"))
+    
+    contents.append(types.Content(role="user", parts=current_parts))
+    
     cfg = types.GenerateContentConfig(response_mime_type="text/plain")
 
     model_start = time.time()
-    res = client.models.generate_content(model="gemini-2.5-flash-lite", contents=contents, config=cfg)
+    res = client.models.generate_content(
+        model="gemini-2.5-flash-lite", 
+        contents=contents, 
+        config=cfg
+    )
     model_end = time.time()
 
-    return {"text": res.text, "timing": {
-        "total_ms": int((time.time() - start_time)*1000),
-        "model_ms": int((model_end - model_start)*1000)
-    }}
+    return {
+        "text": res.text, 
+        "timing": {
+            "total_ms": int((time.time() - start_time)*1000),
+            "model_ms": int((model_end - model_start)*1000)
+        }
+    }
 
 # --- History Management ---
 def get_user_history(uid):
@@ -111,8 +141,13 @@ def get_user_history(uid):
         try:
             resp = requests.get(f"{LAMBDA_URL}?user_id={uid}", timeout=5)
             resp.raise_for_status()
-            user_memory[uid] = {"history": resp.json().get("history", []), "last_sync": 0}
-            app.logger.info(f"Loaded history for {uid}")
+            loaded_history = resp.json().get("history", [])
+            # Only keep the most recent MAX_HISTORY_TURNS when loading
+            user_memory[uid] = {
+                "history": loaded_history[-MAX_HISTORY_TURNS:], 
+                "last_sync": 0
+            }
+            app.logger.info(f"Loaded history for {uid} ({len(user_memory[uid]['history'])} turns)")
         except Exception as e:
             app.logger.warning(f"Failed to load history for {uid}: {e}")
             user_memory[uid] = {"history": [], "last_sync": 0}
@@ -121,6 +156,11 @@ def get_user_history(uid):
 def update_user_history(uid, prompt, response):
     entry = {"prompt": prompt, "response": response, "timestamp": time.time()}
     user_memory.setdefault(uid, {"history": [], "last_sync": 0})["history"].append(entry)
+    
+    # Trim history to MAX_HISTORY_TURNS to prevent unbounded growth
+    if len(user_memory[uid]["history"]) > MAX_HISTORY_TURNS:
+        user_memory[uid]["history"] = user_memory[uid]["history"][-MAX_HISTORY_TURNS:]
+        app.logger.debug(f"Trimmed history for {uid} to {MAX_HISTORY_TURNS} turns")
 
 # --- Routes ---
 @app.route("/")
@@ -140,9 +180,15 @@ def gen():
         return jsonify({"error": "No prompt or image provided"}), 400
 
     try:
-        _ = get_user_history(uid)
-        result = generate_from_gemini(prompt, img_bytes)
+        # Load user's conversation history
+        history = get_user_history(uid)
+        
+        # Generate response with history context
+        result = generate_from_gemini(prompt, img_bytes, history=history)
+        
+        # Update history with new exchange
         update_user_history(uid, prompt, result["text"])
+        
         return jsonify({"result": result["text"], "timing": result["timing"]})
     except Exception as e:
         app.logger.exception("Generation failed")