s

agent.py

@@ -69,19 +69,58 @@ def modulus(a: int, b: int) -> int:
     """
     return a % b
 
+# @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}
 @tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 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}
+    try:
+        # Initialize the tool
+        tavily_tool = TavilySearchResults(max_results=3)
+        # Invoke it correctly
+        search_results = tavily_tool.invoke(input=query) # <--- CORRECTED LINE
+
+        # The result of TavilySearchResults.invoke is usually a list of strings or a single string.
+        # Let's check its type and format accordingly.
+        # Typically, TavilySearchResults directly returns a list of Document objects
+        # or a list of dictionaries if you've configured it differently.
+        # For the default, it's often a list of strings or a single concatenated string.
+        # If it returns a list of Document objects (which is common for loaders/retrievers):
+        if isinstance(search_results, list) and all(hasattr(doc, 'metadata') and hasattr(doc, 'page_content') for doc in search_results):
+            formatted_search_docs = "\n\n---\n\n".join(
+                [
+                    f'<Document source="{doc.metadata.get("source", "N/A")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+                    for doc in search_results
+                ])
+        # If it returns a list of strings (less likely for Tavily in recent versions, but good to check)
+        elif isinstance(search_results, list) and all(isinstance(item, str) for item in search_results):
+            formatted_search_docs = "\n\n---\n\n".join(search_results)
+        # If it returns a single string
+        elif isinstance(search_results, str):
+            formatted_search_docs = search_results
+        else:
+            # Fallback or handle unexpected format
+            print(f"Unexpected Tavily search result format: {type(search_results)}")
+            formatted_search_docs = str(search_results)
+
+        return {"web_results": formatted_search_docs}
+    except Exception as e:
+        print(f"Error during Tavily search for query '{query}': {e}")
+        return {"web_results": f"Error performing web search: {e}"}
 
 @tool
 def web_search(query: str) -> str:
@@ -154,6 +193,11 @@ hf_token = os.environ.get('HF_TOKEN')
 if not hf_token:
     raise ValueError("Hugging Face API token (HF_TOKEN) not found in environment variables.")
 
+tavili_key = os.environ.get('TAVILY_API_KEY')
+if not tavili_key:
+    raise ValueError("Hugging Face API token (HF_TOKEN) not found in environment variables.")
+
+
 # Build graph function
 def build_graph(provider: str = "huggingface"):