Update src/services/ai_modal_engine.py

src/services/ai_modal_engine.py

@@ -6,11 +6,14 @@ from ..config import db, get_user_keys
 class AiModalEngine:
     @staticmethod
     def _get_client(api_key):
-        # Factory method for thread-safe client instantiation per request
+        """Initializes scoped GenAI client per request."""
         return genai.Client(api_key=api_key)
 
     @staticmethod
     def initialize_firebase_session(uid, context):
+        """
+        Bootstraps Auditor session using gemini-3-flash-preview.
+        """
         try:
             keys = get_user_keys(uid)
             api_key = keys.get('gemini_key')
@@ -19,10 +22,10 @@ class AiModalEngine:
 
             client = AiModalEngine._get_client(api_key)
             
-            # Context injection: CSV data is embedded in the persona constraints
+            # 100% Parity with original Persona, Mandate, and CSV injection 
             instruction = f"""
             PARSONA:
-            You are the QuantVAT AI Trading Journal Auditor, a senior Risk Manager and Trading Psychologist. 
+            You are the QuantVAT AI Trading Journal Auditor, a senior Risk Manager and Trading Psychologist with 50 years trading experience like a Market Wizard. 
             Speak with veteran authority. Tone is blunt but constructive.
 
             MANDATE:
@@ -39,10 +42,9 @@ class AiModalEngine:
             {context}
             """
 
-            # Priming prompt forces the model to read the ledger immediately
+            # Original priming prompt 
             prompt = "Analyze my execution performance based on the CSV data above. End with: 'I have analyzed your data. Ready for audit.'"
             
-            # Stateless generation call with system instruction config
             response = client.models.generate_content(
                 model='gemini-3-flash-preview', 
                 contents=prompt,
@@ -51,33 +53,33 @@ class AiModalEngine:
                 )
             )
             
-            # Initialize history structure matching Firestore schema
+            # Persist initial state matching original history schema 
             history = [
                 {"role": "user", "parts": [{"text": prompt}]}, 
                 {"role": "model", "parts": [{"text": response.text}]}
             ]
             
-            # Persist history AND raw context (CSV) to survive stateless resets
+            # Critical: Store ai_context to ensure 'continue' has data memory 
             db.collection('users').document(uid).set({
                 "ai_history": history,
-                "ai_context": context 
+                "ai_context": context
             }, merge=True)
             
             return response.text
         except Exception as e:
-            # Catch-all for API or DB failures to prevent crash
             print(f"AI Init Error: {traceback.format_exc()}")
             return f"System Error: {str(e)}"
 
     @staticmethod
     def continue_firebase_chat(uid, prompt):
+        """
+        Resumes session using history and re-injecting the original CSV context.
+        """
         try:
             user_doc = db.collection('users').document(uid).get()
             data = user_doc.to_dict() if user_doc.exists else {}
-            
-            # Retrieve artifacts required to rebuild the "brain"
             history = data.get("ai_history", [])
-            context = data.get("ai_context", "") 
+            context = data.get("ai_context", "") # Retrieve the original CSV 
             
             api_key = get_user_keys(uid).get('gemini_key')
             if not api_key: 
@@ -85,21 +87,20 @@ class AiModalEngine:
 
             client = AiModalEngine._get_client(api_key)
             
-            contents = [
-                types.Content(
+            # Robust mapping: handles both old string parts and new dict parts 
+            contents = []
+            for h in history:
+                p = h['parts'][0]
+                text_content = p['text'] if isinstance(p, dict) else str(p)
+                contents.append(types.Content(
                     role=h['role'], 
-                    parts=[types.Part.from_text(text=p['parts'][0]['text'])]
-                ) for h in history
-            ]
+                    parts=[types.Part.from_text(text=text_content)]
+                ))
             
-            # Append current user prompt to the stack
             contents.append(types.Content(role="user", parts=[types.Part.from_text(text=prompt)]))
 
-            instruction = f"""
-            PARSONA: QuantVAT AI Trading Journal Auditor. Senior Risk Manager.
-            ORIGINAL LEDGER CONTEXT:
-            {context}
-            """
+            # Re-injects Persona + CSV to prevent memory loss in follow-ups 
+            instruction = f"PARSONA: QuantVAT AI Trading Journal Auditor. Senior Risk Manager.\nDATA:\n{context}"
 
             response = client.models.generate_content(
                 model='gemini-3-flash-preview',
@@ -109,7 +110,7 @@ class AiModalEngine:
                 )
             )
             
-            # Update history with new turn and sync to DB
+            # Append new turn and sync to Firestore 
             history.append({"role": "user", "parts": [{"text": prompt}]})
             history.append({"role": "model", "parts": [{"text": response.text}]})
             db.collection('users').document(uid).set({"ai_history": history}, merge=True)