Update hf_client.py

hf_client.py

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