Update src/llm_client.py

src/llm_client.py

@@ -1,34 +1,76 @@
 import os
 from google import genai
 from google.genai import types
+from .prompts import get_system_prompt  # <--- Essential import
 
 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
     system_instruction = get_system_prompt(mode)
     
+    # --- 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."
+                return "Error: GOOGLE_API_KEY not found in Secrets."
             
-            # NO genai.configure() here. 
-            # The API key goes directly into the Client.
+            # DEBUG PRINT: Verify we are running the new code
+            print("DEBUG: Initializing new Gemini Client...")
+
+            # NEW SYNTAX: Instantiate a Client object directly
+            # (No genai.configure() calls allowed here!)
             client = genai.Client(api_key=google_key)
             
+            # Construct the prompt
             full_prompt = f"USER QUERY: {query}\n\nDOCUMENT CONTEXT:\n{context[:30000]}"
             
-            # Using Gemini 2.0 Flash
+            # NEW SYNTAX: Call generate_content via the 'models' attribute
             response = client.models.generate_content(
-                model='gemini-2.0-flash',
+                model='gemini-2.0-flash', 
                 contents=full_prompt,
                 config=types.GenerateContentConfig(
                     system_instruction=system_instruction,
                     max_output_tokens=1000,
-                    temperature=0.3,
-                ),
+                    temperature=0.3
+                )
             )
             
             return response.text
             
         except Exception as e:
+            print(f"DEBUG ERROR: {e}")
             return f"Gemini Error: {str(e)}"
+
+    # --- OPTION B: GRANITE / LOCAL SPACE (Private Option) ---
+    else:
+        import requests # Imported locally to keep dependencies clean if needed
+        hf_token = os.environ.get("HF_TOKEN")
+        if not hf_token:
+            return "Error: HF_TOKEN is missing."
+
+        api_url = "https://navydevildoc-private-granite.hf.space/generate"
+        
+        payload = {
+            "text": f"USER QUESTION: {query}\n\nDOCUMENT CONTEXT:\n{context[:6000]}",
+            "persona": system_instruction, 
+            "model": "granite4:latest", 
+            "max_tokens": 1024
+        }
+        
+        headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "Content-Type": "application/json"
+        }
+        
+        try:
+            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:
+                return f"Error {response.status_code}: {response.text}"
+        except Exception as e:
+            return f"Connection Error: {str(e)}"