Update cool_agent.py

cool_agent.py

@@ -1,6 +1,7 @@
-import os 
+import os
 from linkup_utils import LinkupSearchTool
 from text_inspector_tool import TextInspectorTool
+from langchain_community.document_loaders import WikipediaLoader
 from text_web_browser import (
     ArchiveSearchTool,
     FinderTool,
@@ -18,6 +19,21 @@ from smolagents import (
     ToolCallingAgent,
 )
 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}
 
 
 AUTHORIZED_IMPORTS = [
@@ -91,6 +107,7 @@ def create_agent(model_id="gpt-4o-mini"):
         FindNextTool(browser),
         ArchiveSearchTool(browser),
         TextInspectorTool(model, text_limit),
+        wiki_search(),
     ]
     text_webbrowser_agent = ToolCallingAgent(
         model=model,
@@ -123,4 +140,4 @@ def create_agent(model_id="gpt-4o-mini"):
         managed_agents=[text_webbrowser_agent],
     )
 
-    return manager_agent
+    return manager_agent