Update src/services/ai_modal_engine.py

src/services/ai_modal_engine.py

@@ -1,23 +1,24 @@
-import google.generativeai as genai
+from google import genai
+from google.genai import types
 import traceback
 from ..config import db, get_user_keys
 
 class AiModalEngine:
+    @staticmethod
+    def _get_client(api_key):
+        return genai.Client(api_key=api_key)
+
     @staticmethod
     def initialize_firebase_session(uid, context):
-        """
-        Bootstraps Auditor session.
-        Args:
-            uid: User ID.
-            context: Raw CSV string provided by frontend.
-        """
         try:
             keys = get_user_keys(uid)
             api_key = keys.get('gemini_key')
             if not api_key: 
                 return "Error: No API Key found in Settings."
 
-            genai.configure(api_key=api_key) 
+            client = AiModalEngine._get_client(api_key)
+            
+            # Parsona
             instruction = f"""
             PARSONA:
             You are the QuantVAT AI Trading Journal Auditor, a senior Risk Manager and Trading Psychologist with 50 years trading experience like a Market Wizard. 
@@ -37,22 +38,25 @@ class AiModalEngine:
             {context}
             """
 
-            model = genai.GenerativeModel(
-                model_name='gemini-3-flash-preview',
-                system_instruction=instruction
-            )
-            
-            chat = model.start_chat(history=[])
-            # Priming prompt to force immediate analysis
             prompt = "Analyze my execution performance based on the CSV data above. End with: 'I have analyzed your data. Ready for audit.'"
-            response = chat.send_message(prompt)
             
-            #  Persist initial state to Firestore for session continuity
+            response = client.models.generate_content(
+                model='gemini-3-flash-preview', 
+                contents=prompt,
+                config=types.GenerateContentConfig(
+                    system_instruction=instruction
+                )
+            )
+            
             history = [
-                {"role": "user", "parts": [prompt]}, 
-                {"role": "model", "parts": [response.text]}
+                {"role": "user", "parts": [{"text": prompt}]}, 
+                {"role": "model", "parts": [{"text": response.text}]}
             ]
-            db.collection('users').document(uid).set({"ai_history": history}, merge=True)
+            
+            db.collection('users').document(uid).set({
+                "ai_history": history,
+                "ai_context": context
+            }, merge=True)
             
             return response.text
         except Exception as e:
@@ -61,32 +65,44 @@ class AiModalEngine:
 
     @staticmethod
     def continue_firebase_chat(uid, prompt):
-        """
-        Resumes active session using state retrieved from Firestore.
-        """
         try:
             user_doc = db.collection('users').document(uid).get()
             data = user_doc.to_dict() if user_doc.exists else {}
             history = data.get("ai_history", [])
+            context = data.get("ai_context", "") 
             
             api_key = get_user_keys(uid).get('gemini_key')
             if not api_key: 
                 return "Error: API Key missing."
 
-            genai.configure(api_key=api_key)
+            client = AiModalEngine._get_client(api_key)
             
-            #  Re-assert persona for stateless continuity
-            model = genai.GenerativeModel(
-                model_name='gemini-3-flash-preview',
-                system_instruction="PARSONA: QuantVAT AI Trading Journal Auditor. Senior Risk Manager."
-            )
+            # Robust mapping
+            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=text_content)]
+                ))
             
-            chat = model.start_chat(history=history)
-            response = chat.send_message(prompt)
+            contents.append(types.Content(role="user", parts=[types.Part.from_text(text=prompt)]))
+
+            # Re-injects Persona
+            instruction = f"PARSONA: QuantVAT AI Trading Journal Auditor. Senior Risk Manager.\nDATA:\n{context}"
+
+            response = client.models.generate_content(
+                model='gemini-3-flash-preview',
+                contents=contents,
+                config=types.GenerateContentConfig(
+                    system_instruction=instruction
+                )
+            )
             
-            # Append new turn and sync to DB
-            history.append({"role": "user", "parts": [prompt]})
-            history.append({"role": "model", "parts": [response.text]})
+            # 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)
             
             return response.text