Update llm_client.py

llm_client.py

@@ -64,6 +64,7 @@ def setup_llama_binaries():
 
 # --- Custom LangChain LLM Wrapper for Hybrid Approach ---
 class HybridLLM(LLM):
+    groq_client: Any = None
     api_url: str = ""
     local_server_url: str = "http://localhost:8080"
     
@@ -72,7 +73,18 @@ class HybridLLM(LLM):
         return "hybrid_llm"
 
     def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs: Any) -> str:
-        # 1. Try Colab API first
+        # 1. Try Groq API (Highest Priority)
+        if self.groq_client:
+            try:
+                print("⚡ Using Groq API...")
+                # ChatGroq returns a message object, we need the content string
+                response = self.groq_client.invoke(prompt)
+                return response.content
+            except Exception as e:
+                print(f"⚠️ Groq API Failed: {e}")
+                # Continue to next fallback
+
+        # 2. Try Colab API
         if self.api_url:
             try:
                 print(f"🌐 Calling Colab API: {self.api_url}")
@@ -88,13 +100,13 @@ class HybridLLM(LLM):
             except Exception as e:
                 print(f"⚠️ API Connection Failed: {e}")
         
-        # 2. Fallback to Local Server
+        # 3. Fallback to Local Server
         print("💻 Using Local llama-server Fallback...")
         try:
             # OpenAI-compatible completion endpoint
             payload = {
                 "prompt": prompt,
-                "n_predict": 512,
+                "n_predict": 1024,
                 "temperature": 0.3,
                 "stop": (stop or []) + ["<|im_end|>", "Input:", "Context:"]
             }
@@ -114,7 +126,11 @@ class HybridLLM(LLM):
 
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
-        return {"api_url": self.api_url, "local_server_url": self.local_server_url}
+        return {
+            "groq_enabled": self.groq_client is not None,
+            "api_url": self.api_url, 
+            "local_server_url": self.local_server_url
+        }
 
 class LLMClient:
     def __init__(self, vector_store=None):
@@ -125,28 +141,28 @@ class LLMClient:
         self.api_url = os.environ.get("COLAB_API_URL", "")
         self.server_process = None
         self.server_port = 8080
+        self.groq_client = None
 
-        # 1. Setup Groq (Highest Priority)
-        self.groq_api_key = os.environ.get("GROQ_API_KEY")
-        self.groq_model = "qwen-2.5-32b" # Default, user can override via env if needed
+        # 1. Setup Groq Client
+        groq_api_key = os.environ.get("GROQ_API_KEY")
+        groq_model = "qwen/qwen3-32b" 
 
-        if self.groq_api_key:
+        if groq_api_key:
             try:
                 from langchain_groq import ChatGroq
-                print(f"⚡ Initializing Groq API with model: {self.groq_model}...")
-                self.llm = ChatGroq(
+                print(f"⚡ Initializing Groq Client ({groq_model})...")
+                self.groq_client = ChatGroq(
                     temperature=0.3,
-                    model=self.groq_model,
-                    groq_api_key=self.groq_api_key,
+                    model=groq_model,
+                    groq_api_key=groq_api_key,
                     max_tokens=1024,
                     model_kwargs={"stop": ["<|im_end|>", "Input:", "Context:"]}
                 )
-                print("✅ Groq API ready! Skipping local setup.")
-                return # Exit early, no need for local server
+                print("✅ Groq Client ready.")
             except Exception as e:
                 print(f"⚠️ Groq Init Failed: {e}")
         
-        # 2. Setup Local Fallback (Only if Groq fails/missing)
+        # 2. Setup Local Fallback (Always setup as requested)
         try:
             # Setup Binary
             self.server_bin, self.lib_path = setup_llama_binaries()
@@ -168,8 +184,9 @@ class LLMClient:
         except Exception as e:
             print(f"⚠️ Could not setup local fallback: {e}")
 
-        # Create Hybrid LangChain Wrapper (Fallback)
+        # Create Hybrid LangChain Wrapper
         self.llm = HybridLLM(
+            groq_client=self.groq_client,
             api_url=self.api_url,
             local_server_url=f"http://localhost:{self.server_port}"
         )