Update hf_client.py

hf_client.py

@@ -1,45 +1,64 @@
 import os
-import openai
 from huggingface_hub import InferenceClient
 from tavily import TavilyClient
 
-# === Environment Setup ===
+# === Required Tokens ===
 
 HF_TOKEN = os.getenv('HF_TOKEN')
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
 
 if not HF_TOKEN:
     raise RuntimeError(
         "HF_TOKEN environment variable is not set. Please set it to your Hugging Face API token."
     )
 
-# === Dynamic Inference Client ===
+
+# === Groq Adapter ===
+
+class GroqChatClient:
+    def __init__(self, api_key: str, model: str = "mixtral-8x7b-32768"):
+        import openai
+        openai.api_key = api_key
+        openai.api_base = "https://api.groq.com/openai/v1"
+        self.openai = openai
+        self.model = model
+
+    class Chat:
+        def __init__(self, openai_instance, model):
+            self.completions = self.Completions(openai_instance, model)
+
+        class Completions:
+            def __init__(self, openai_instance, model):
+                self.client = openai_instance
+                self.model = model
+
+            def create(self, messages, temperature=0.7, max_tokens=1024):
+                return self.client.ChatCompletion.create(
+                    model=self.model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens
+                )
+
+    @property
+    def chat(self):
+        return self.Chat(self.openai, self.model)
+
+
+# === Unified Client Factory ===
 
 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.
+    Returns a unified client interface.
+    - If model_id is 'moonshotai/Kimi-K2-Instruct', use Groq via OpenAI adapter
+    - Otherwise, use Hugging Face's InferenceClient
     """
     if model_id == "moonshotai/Kimi-K2-Instruct":
         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"
+            raise RuntimeError("GROQ_API_KEY is not set. Required for Groq models.")
+        return GroqChatClient(api_key=GROQ_API_KEY)
 
-        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,
@@ -47,9 +66,9 @@ def get_inference_client(model_id: str, provider: str = "auto"):
         bill_to="huggingface"
     )
 
-# === Tavily Client ===
 
-TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
+# === Tavily Search Client ===
+
 tavily_client = None
 if TAVILY_API_KEY:
     try: