Update agent.py

agent.py

@@ -3,12 +3,13 @@ import hashlib
 import json
 import logging
 from smolagents import CodeAgent, tool
+from smolagents.models import Model  # <-- base class for models
 from huggingface_hub import InferenceClient
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Cache for answers (persists between runs)
+# Cache for answers
 CACHE_FILE = "answer_cache.json"
 if os.path.exists(CACHE_FILE):
     with open(CACHE_FILE) as f:
@@ -36,7 +37,6 @@ def calculator(expression: str) -> str:
     if not all(c in allowed_chars for c in expression):
         return "Error: Expression contains disallowed characters."
     try:
-        # Restricted eval – only math allowed
         result = eval(expression, {"__builtins__": {}}, {})
         return str(result)
     except Exception as e:
@@ -69,27 +69,37 @@ def web_search(query: str) -> str:
     except Exception as e:
         return f"Search error: {e}"
 
-# ---------- Custom model that wraps HF InferenceClient ----------
-class CustomHFModel:
-    def __init__(self, model_id="HuggingFaceH4/zephyr-7b-beta"):
+# ---------- Custom model that inherits from smolagents.Model ----------
+class CustomHFModel(Model):
+    def __init__(self, model_id="HuggingFaceH4/zephyr-7b-beta", **kwargs):
+        super().__init__(**kwargs)
         self.client = InferenceClient(model=model_id, token=os.getenv("HF_TOKEN"))
         self.model_id = model_id
 
-    def __call__(self, messages, **kwargs):
+    def generate(
+        self,
+        messages,
+        stop_sequences=None,
+        grammar=None,
+        max_tokens=500,
+        temperature=0.7,
+        **kwargs
+    ):
         """
-        Expected by smolagents: takes a list of messages
+        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=500,
-            temperature=0.7,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            stop=stop_sequences,
             **kwargs
         )
         return response.choices[0].message.content
 
-# ---------- Assemble the agent (once, at import) ----------
+# ---------- Assemble the agent ----------
 tools = [calculator]
 try:
     import duckduckgo_search
@@ -104,7 +114,7 @@ agent = CodeAgent(tools=tools, model=model)
 # ---------- The class expected by app.py ----------
 class CustomAgent:
     def __call__(self, question: str) -> str:
-        """This method is called for each question."""
+        """Called for each question."""
         q_hash = hashlib.md5(question.encode()).hexdigest()
         if q_hash in answer_cache:
             logger.info(f"Cache hit for question: {question[:50]}...")