tavily is just generating problems for me

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+from typing import Dict, List
 from smolagents import CodeAgent, tool, InferenceClientModel
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
@@ -53,18 +54,13 @@ def modulus(a: int, b: int) -> int:
     return a % b
 
 @tool
-def web_search(query: str) -> str:
+def web_search(query: str, max_results: int) -> Dict[str, str]:
     """Search Tavily for a query and return maximum 3 results.
-    
     Args:
-        query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(input=query)
-    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 {"web_results": formatted_search_docs}
+        query: The search query.
+        max_results: The maximum number of results to return."""
+    search_docs = TavilySearchResults(max_results=max_results).invoke(input=query)
+    return {"web_results": search_docs}
 
 @tool
 def arxiv_search(query: str) -> str: