Update cool_agent.py

cool_agent.py

@@ -17,23 +17,31 @@ from smolagents import (
     CodeAgent,
     LiteLLMModel,
     ToolCallingAgent,
+    Tool,
 )
 import threading
-from langchain_core.tools import tool
 
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
 
-    Args:
-        query: The search query."""
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"wiki_results": formatted_search_docs}
+class WikiSearchTool(Tool):
+    name = "wiki_search"
+    description = "Search Wikipedia for a query and return maximum 2 results."
+    inputs = {
+        "query": {"type": "string", "description": "The search query to perform."},
+    }
+    output_type = "dict"
+
+    def __init__(self):
+        super().__init__(self)
+
+    def forward(self, query: str) -> dict:
+        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+        formatted_search_docs = "\n\n---\n\n".join(
+            [
+                f'<Document source="{doc.metadata["source"]}" page="'
+                f'{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+                for doc in search_docs
+            ])
+        return {"wiki_results": formatted_search_docs}
 
 
 AUTHORIZED_IMPORTS = [
@@ -68,10 +76,10 @@ AUTHORIZED_IMPORTS = [
 
 append_answer_lock = threading.Lock()
 
-
 custom_role_conversions = {"tool-call": "assistant", "tool-response": "user"}
 
-user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
+user_agent = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
+              "Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0")
 
 BROWSER_CONFIG = {
     "viewport_size": 1024 * 5,
@@ -107,7 +115,7 @@ def create_agent(model_id="gpt-4o-mini"):
         FindNextTool(browser),
         ArchiveSearchTool(browser),
         TextInspectorTool(model, text_limit),
-        wiki_search,
+        wiki_search(),
     ]
     text_webbrowser_agent = ToolCallingAgent(
         model=model,
@@ -118,17 +126,26 @@ def create_agent(model_id="gpt-4o-mini"):
         name="search_agent",
         description="""A team member that will search the internet to answer your question.
     Ask him for all your questions that require browsing the web. 
-    Note that this agent is using a powerful language model and it can do the search and analyse the results. 
-    You should ask question in a way to let the language model to perform the best, i.e. provide as much context as possible and ask in a clear way.
-    Provide him as much context as possible, in particular if you need to search on a specific timeframe!
-    And don't hesitate to provide him with a complex search task, like finding a difference between two webpages.
-    Your request must be a real sentence, not a google search! Like "Find me this information (...)" rather than a few keywords.
+    Note that this agent is using a powerful language model and it can do the search and analyse 
+    the results. 
+    You should ask question in a way to let the language model to perform the best, i.e. provide 
+    as much context as possible and ask in a clear way.
+    Provide him as much context as possible, in particular if you need to search on a specific 
+    timeframe!
+    And don't hesitate to provide him with a complex search task, like finding a difference 
+    between two webpages.
+    Your request must be a real sentence, not a google search! Like "Find me this information (
+    ...)" rather than a few keywords.
     """,
         provide_run_summary=True,
     )
-    text_webbrowser_agent.prompt_templates["managed_agent"]["task"] += """You can navigate to .txt online files.
-    If a non-html page is in another format, especially .pdf or a Youtube video, use tool 'inspect_file_as_text' to inspect it.
-    Additionally, if after some searching you find out that you need more information to answer the question, you can use `final_answer` with your request for clarification as argument to request for more information."""
+    text_webbrowser_agent.prompt_templates["managed_agent"]["task"] += """You can navigate to 
+    .txt online files.
+    If a non-html page is in another format, especially .pdf or a Youtube video, use tool 
+    'inspect_file_as_text' to inspect it.
+    Additionally, if after some searching you find out that you need more information to answer 
+    the question, you can use `final_answer` with your request for clarification as argument to 
+    request for more information."""
 
     manager_agent = CodeAgent(
         model=model,
@@ -140,4 +157,4 @@ def create_agent(model_id="gpt-4o-mini"):
         managed_agents=[text_webbrowser_agent],
     )
 
-    return manager_agent
+    return manager_agent