Update hf_client.py

hf_client.py

@@ -1,29 +1,54 @@
-
-
-# hf_client.py
-
 import os
-from huggingface_hub import InferenceClient, HfApi
+import openai
+from huggingface_hub import InferenceClient
 from tavily import TavilyClient
 
-# HF Inference Client
+# === Environment Setup ===
+
 HF_TOKEN = os.getenv('HF_TOKEN')
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+
 if not HF_TOKEN:
     raise RuntimeError(
         "HF_TOKEN environment variable is not set. Please set it to your Hugging Face API token."
     )
 
-def get_inference_client(model_id: str, provider: str = "auto") -> InferenceClient:
-    """Return an InferenceClient with the appropriate provider."""
+# === Dynamic Inference Client ===
+
+def get_inference_client(model_id: str, provider: str = "auto"):
+    """
+    Return an inference client depending on model ID.
+    Uses Groq's native API for specific models, otherwise HuggingFace InferenceClient.
+    """
     if model_id == "moonshotai/Kimi-K2-Instruct":
-        provider = "groq"
+        if not GROQ_API_KEY:
+            raise RuntimeError("GROQ_API_KEY is required for Groq models.")
+        
+        # Configure OpenAI client for Groq
+        openai.api_key = GROQ_API_KEY
+        openai.api_base = "https://api.groq.com/openai/v1"
+
+        def chat(messages, temperature=0.7, max_tokens=1024):
+            response = openai.ChatCompletion.create(
+                model="mixtral-8x7b-32768",  # You can map the model here
+                messages=messages,
+                temperature=temperature,
+                max_tokens=max_tokens
+            )
+            return response["choices"][0]["message"]["content"]
+        
+        return chat  # Return callable interface
+
+    # Fallback to Hugging Face
     return InferenceClient(
+        model=model_id,
         provider=provider,
         api_key=HF_TOKEN,
         bill_to="huggingface"
     )
 
-# Tavily Search Client
+# === Tavily Client ===
+
 TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
 tavily_client = None
 if TAVILY_API_KEY: