Update src/llm_client.py

src/llm_client.py

@@ -1,72 +1,34 @@
-import requests
 import os
 from google import genai
-from google.genai import types # The new configuration handler
-from .prompts import get_system_prompt
+from google.genai import types
 
 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 in Secrets."
+                return "Error: GOOGLE_API_KEY not found."
             
-            # NEW SYNTAX: Instantiate a Client object directly
+            # NO genai.configure() here. 
+            # The API key goes directly into the Client.
             client = genai.Client(api_key=google_key)
             
-            # Construct the prompt
             full_prompt = f"USER QUERY: {query}\n\nDOCUMENT CONTEXT:\n{context[:30000]}"
             
-            # NEW SYNTAX: Call generate_content via the 'models' attribute
+            # Using Gemini 2.0 Flash
             response = client.models.generate_content(
-                model='gemini-2.0-flash', # or 'gemini-2.0-flash' if available to you
+                model='gemini-2.0-flash',
                 contents=full_prompt,
                 config=types.GenerateContentConfig(
                     system_instruction=system_instruction,
                     max_output_tokens=1000,
-                    temperature=0.3 # Low temperature for factual accuracy
-                )
+                    temperature=0.3,
+                ),
             )
             
-            # 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 (Unchanged) ---
-    else:
-        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)}"