Create agent.py

agent.py

@@ -0,0 +1,96 @@
+import os
+import time
+from smolagents import (
+    ToolCallingAgent,
+    CodeAgent,
+    MCPClient,
+    WikipediaSearchTool,
+    InferenceClientModel,
+)
+
+max_steps_internet_researcher = 10
+max_steps_manager = 10
+max_retries = 3
+cooldown_on_error = 10
+
+
+qwuen_72b = "Qwen/Qwen2.5-72B-Instruct"
+model_235b_thinking = "qwen/qwen3-235b-a22b-thinking-2507"
+model_235b_instruct = "qwen/Qwen3-235B-A22B-Instruct-2507"
+model_235b = "qwen/qwen3-235b-a22b-fp8"
+model_480b = "qwen/qwen3-coder-480b-a35b-instruct"
+
+
+def get_hf_model():
+    hf_token = os.getenv("HF_TOKEN")
+    return InferenceClientModel(model_235b_instruct, token=hf_token, timeout=300)
+
+
+def get_tavily_mcp_client():
+    tavily_token = os.getenv("TAVILY_TOKEN")
+
+    # context manager + Streamable HTTP transport:
+    return MCPClient(
+        {
+            "url": f"https://mcp.tavily.com/mcp/?tavilyApiKey={tavily_token}",
+            "transport": "streamable-http",
+        }
+    )
+
+
+class Agent:
+    def __init__(self):
+
+        model = get_hf_model()
+        mcpClient = get_tavily_mcp_client()
+
+        tools = mcpClient.get_tools()
+        tools.append(
+            WikipediaSearchTool(
+                user_agent="Research ([email protected])",
+                language="en",
+                content_type="summary",
+                extract_format="WIKI",
+            )
+        )
+
+        internet_reseacher = ToolCallingAgent(
+            tools=tools,
+            model=model,
+            max_steps=max_steps_internet_researcher,
+            name="web_search_agent",
+            description="Runs web searches for you.",
+        )
+
+        agent = CodeAgent(
+            tools=[],
+            managed_agents=[internet_reseacher],
+            max_steps=max_steps_manager,
+            model=model,
+        )
+
+        self.agent = agent
+
+        print("BasicAgent initialized.")
+
+    def __call__(self, question: str) -> str:
+        prompt = f"You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with ONLY the final answer. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. the question is: {question}"
+        print(f"Agent received question: {prompt}...")
+
+        for _ in range(0, max_retries):
+            try:
+                answer = self.agent.run(prompt)
+
+                if isinstance(answer, str):
+                    if "AGENT ERROR" in answer:
+                        print("agent error: ", answer)
+                        time.sleep(cooldown_on_error)
+                        continue
+                print(f"Agent returning fixed answer: {answer}")
+
+                return answer
+            except Exception as e:
+                print("Other error: ", e)
+                time.sleep(cooldown_on_error)
+
+        return "ERROR: no more retries"