Update agent.py

agent.py

@@ -12,7 +12,7 @@ from typing import Any
 from dotenv import load_dotenv
 from google.generativeai import types, configure
 
-from smolagents import InferenceClientModel, LiteLLMModel, CodeAgent, ToolCallingAgent, Tool, DuckDuckGoSearchTool
+from smolagents import InferenceClientModel, LiteLLMModel, CodeAgent, ToolCallingAgent, Tool, DuckDuckGoSearchTool, HfApiModel
 
 # Logging
 #logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
@@ -90,8 +90,14 @@ class WikiContentFetcher(Tool):
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
-        model = HF_MODEL_NAME
-        client = InferenceClientModel()
+        #model = HF_MODEL_NAME
+        model = HfApiModel(
+            max_tokens=2096,
+            temperature=0.5,
+            model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
+            #model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
+            custom_role_conversions=None,
+            )
         tools = [
             DuckDuckGoSearchTool(),
             WikiTitleFinder(),
@@ -101,7 +107,7 @@ class BasicAgent:
             TextTransformer(),
         ]
         self.agent = CodeAgent(
-            model=HF_MODEL_NAME,
+            model=model,
             tools=tools,
             add_base_tools=False,
             max_steps=10,
@@ -161,7 +167,13 @@ class BasicAgent:
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         result = self.agent.run(question)
-        final_str = str(result).strip()
+
+        if isinstance(result, dict) and "output" in result:
+            final_str = str(result["output"]).strip()
+        elif hasattr(result, "output"):
+            final_str = str(result.output).strip()
+        else:
+            final_str = str(result).strip()
 
         return final_str