Update agent.py

agent.py

@@ -60,6 +60,7 @@ from functools import reduce
 import operator
 
 
+
 load_dotenv()
 
 
@@ -154,18 +155,24 @@ def wikidata_query(query: str) -> str:
 
 @tool
 def web_search(query: str) -> str:
-    """Search Tavily for a query and return maximum 3 results.
+    """Search Tavily for a query and return up to 3 results."""
+    tavily_key = os.getenv("TAVILY_API_KEY")
+    
+    if not tavily_key:
+        return "Error: Tavily API key not set."
+
+    search_tool = TavilySearchResults(tavily_api_key=tavily_key, max_results=3)
+    search_docs = search_tool.invoke(query=query)
     
-    Args:
-        query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(query=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 formatted_search_docs
 
+
 @tool
 def arvix_search(query: str) -> str:
     """Search Arxiv for a query and return maximum 3 result.