Update agent.py

agent.py

@@ -3,12 +3,19 @@ import hashlib
 import json
 import logging
 from smolagents import CodeAgent, tool
-from smolagents.models import Model  # <-- base class for models
+from smolagents.models import Model
 from huggingface_hub import InferenceClient
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Check HF_TOKEN at startup
+hf_token = os.getenv("HF_TOKEN")
+if hf_token:
+    logger.info("HF_TOKEN is set (first 4 chars: %s...)", hf_token[:4])
+else:
+    logger.warning("HF_TOKEN is NOT set - inference may fail")
+
 # Cache for answers
 CACHE_FILE = "answer_cache.json"
 if os.path.exists(CACHE_FILE):
@@ -69,12 +76,18 @@ def web_search(query: str) -> str:
     except Exception as e:
         return f"Search error: {e}"
 
-# ---------- Custom model that inherits from smolagents.Model ----------
+# ---------- Custom model with explicit Hugging Face endpoint ----------
 class CustomHFModel(Model):
-    def __init__(self, model_id="HuggingFaceH4/zephyr-7b-beta", **kwargs):
+    def __init__(self, model_id="mistralai/Mistral-7B-Instruct-v0.2", **kwargs):
         super().__init__(**kwargs)
-        self.client = InferenceClient(model=model_id, token=os.getenv("HF_TOKEN"))
+        # Force Hugging Face inference endpoint
+        self.client = InferenceClient(
+            model=model_id,
+            token=os.getenv("HF_TOKEN"),
+            base_url="https://api-inference.huggingface.co"  # Explicit HF URL
+        )
         self.model_id = model_id
+        logger.info(f"CustomHFModel initialized with model: {model_id}")
 
     def generate(
         self,
@@ -85,19 +98,18 @@ class CustomHFModel(Model):
         temperature=0.7,
         **kwargs
     ):
-        """
-        Required by smolagents: takes a list of messages
-        (e.g., [{"role": "user", "content": "..."}])
-        and returns the assistant's reply as a string.
-        """
-        response = self.client.chat_completion(
-            messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            stop=stop_sequences,
-            **kwargs
-        )
-        return response.choices[0].message.content
+        try:
+            response = self.client.chat_completion(
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stop=stop_sequences,
+                **kwargs
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            # Wrap error with model info for debugging
+            raise Exception(f"Model {self.model_id} inference failed: {e}")
 
 # ---------- Assemble the agent ----------
 tools = [calculator]
@@ -108,7 +120,7 @@ try:
 except ImportError:
     logger.warning("duckduckgo-search not installed, web_search disabled.")
 
-model = CustomHFModel()  # you can change the model_id if desired
+model = CustomHFModel()  # You can change model_id here if needed
 agent = CodeAgent(tools=tools, model=model)
 
 # ---------- The class expected by app.py ----------