Update agent.py

agent.py

@@ -3,8 +3,6 @@ from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
 
 
-
-
 class add(Tool):
     name = "add"
     description = """
@@ -128,7 +126,7 @@ class modulo(Tool):
     def forward(self, a: int, b: int):
         return a % b
 
-class WikipediaSearchTool(tool):
+class WikipediaSearchTool(Tool):
     name = "wikipedia_search_tool"
     
     description = """
@@ -155,7 +153,7 @@ class WikipediaSearchTool(tool):
         ])
         return {"wikipedia_results": condensed_docs}
 
-class TavilySearchTool(tool):
+class TavilySearchTool(Tool):
     name = "tavily_search_tool"
     
     description = """
@@ -183,20 +181,32 @@ class TavilySearchTool(tool):
         return {"web_search_results": condensed_docs}
 
 
-@tool
-def arvix_search(query: str) -> str:
-    """Search Arxiv for a query and return maximum 3 result.
+class ArvixSearchTool(Tool):
+    name = "arvix_search_tool"
+    
+    description = """
+        Search arxiv for a query and return maximum 3 result.
     
-    Args:
-        query: The search query."""
-    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    formatted_search_docs = "\n\n---\n\n".join(
+        Args:
+            query: The search query.
+        """
+    
+    inputs = {
+        "query":{
+            "type":"string",
+            "description":"the search query"
+        }
+    }
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        documents = ArxivLoader(query=query, load_max_docs=3).load()
+        condensed_docs = "\n\n---\n\n".join(
         [
             f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
             for doc in search_docs
         ])
-    return {"arvix_results": formatted_search_docs}
-
+    return {"arvix_search_results": formatted_search_docs}
 
 
 tools=[
@@ -205,20 +215,33 @@ tools=[
     multiply(),
     divide(),
     modulo(),
-    DuckDuckGoSearchTool(),
-    WikipediaSearchTool()
+    WikipediaSearchTool(),
+    ArvixSearchTool(),
+    TavilySearchTool()
 ]
 
 model = LiteLLMModel(
-  model_id='ollama_chat/qwen2.5:14b'
+    model_id='ollama_chat/gemma3:27b',
+    api_base="http://127.0.0.1:11434",
+    num_ctx=8192
 )
 
+def delay_execution_10(pagent, **kwargs) -> bool:
+    """
+    Delays the execution for 10 seconds.
+    """
+    time.sleep(10)
+    return True
+
 def create_agent():
     agent = CodeAgent(
         model = model,
         tools = tools,
         max_steps=10,
-        verbosity_level=2
+        verbosity_level=2,
+        additional_authorized_imports=['*'],
+        planning_interval=5,  
+        step_callbacks=[delay_execution_10]
     )
 
     return agent