Update src/llm_client.py

src/llm_client.py

@@ -1,6 +1,7 @@
 import requests
 import os
-import google.genai as genai
+from google import genai
+from google.genai import types # The new configuration handler
 from .prompts import get_system_prompt
 
 def ask_llm(query, context, mode="Executive Summary", model_provider="Gemini"):
@@ -8,34 +9,40 @@ def ask_llm(query, context, mode="Executive Summary", model_provider="Gemini"):
     Switchable Brain: Defaults to Gemini (Powerful), falls back to Granite (Private).
     """
     
-    # 1. Get the Persona/Prompt from our new library
-    # This ensures "Instructor Mode" works for BOTH models
+    # 1. Get the Persona/Prompt
     system_instruction = get_system_prompt(mode)
     
-    # --- OPTION A: GOOGLE GEMINI (The "Big Gun") ---
+    # --- OPTION A: GOOGLE GEMINI (New SDK) ---
     if model_provider == "Gemini":
         try:
             google_key = os.environ.get("GOOGLE_API_KEY")
-            
             if not google_key:
                 return "Error: GOOGLE_API_KEY not found in Secrets."
             
-            genai.configure(api_key=google_key)
-            
-            model = genai.GenerativeModel(
-                model_name='gemini-1.5-flash',
-                system_instruction=system_instruction
-            )
+            # NEW SYNTAX: Instantiate a Client object directly
+            client = genai.Client(api_key=google_key)
             
+            # Construct the prompt
             full_prompt = f"USER QUERY: {query}\n\nDOCUMENT CONTEXT:\n{context[:30000]}"
             
-            response = model.generate_content(full_prompt)
+            # NEW SYNTAX: Call generate_content via the 'models' attribute
+            response = client.models.generate_content(
+                model='gemini-1.5-flash', # or 'gemini-2.0-flash' if available to you
+                contents=full_prompt,
+                config=types.GenerateContentConfig(
+                    system_instruction=system_instruction,
+                    max_output_tokens=1000,
+                    temperature=0.3 # Low temperature for factual accuracy
+                )
+            )
+            
+            # The response object structure is slightly different now
             return response.text
             
         except Exception as e:
             return f"Gemini Error: {str(e)}"
 
-    # --- OPTION B: GRANITE / LOCAL SPACE (The "Private Option") ---
+    # --- OPTION B: GRANITE / LOCAL SPACE (Unchanged) ---
     else:
         hf_token = os.environ.get("HF_TOKEN")
         if not hf_token:
@@ -45,10 +52,9 @@ def ask_llm(query, context, mode="Executive Summary", model_provider="Gemini"):
         
         payload = {
             "text": f"USER QUESTION: {query}\n\nDOCUMENT CONTEXT:\n{context[:6000]}",
-            # CRITICAL FIX: We now inject the dynamic persona here
             "persona": system_instruction, 
-            "model": "granite4:latest", # Or gemma3:latest
-            "max_tokens": 5000
+            "model": "granite4:latest", 
+            "max_tokens": 1024
         }
         
         headers = {
@@ -57,8 +63,7 @@ def ask_llm(query, context, mode="Executive Summary", model_provider="Gemini"):
         }
         
         try:
-            # We increased timeout because local models can be slower
-            response = requests.post(api_url, json=payload, headers=headers, timeout=600)
+            response = requests.post(api_url, json=payload, headers=headers, timeout=120)
             if response.status_code == 200:
                 return response.json().get("response", "Error: Empty response.")
             else: