Update tools.py

tools.py

@@ -20,6 +20,10 @@ import re
 from datetime import datetime, timedelta
 from langchain_core.tools import tool
 import logging
+from tavily import TavilyClient
+import os
+
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 
 # Dictionary of known GAIA-style entities → canonical Wikipedia URLs
 WIKIPEDIA_PAGES = {
@@ -32,6 +36,26 @@ WIKIPEDIA_PAGES = {
     # ➕ add more GAIA topics here
 }
 
+@tool
+def extract_number_from_snippets(snippets: list[str]) -> Optional[int]:
+    for s in snippets:
+        match = re.search(r"\b\d{1,4}\b", s)
+        if match:
+            return int(match.group())
+    return None
+
+
+@tool
+def tavily_search(query: str, k: int = 5) -> list[str]:
+    """
+    Perform a web search using the Tavily API and return up to k relevant snippets.
+    """
+    try:
+        response = client.search(query=query, search_depth="advanced", max_results=k)
+        return [r["content"] for r in response.get("results", [])]
+    except Exception as e:
+        return [f"❌ Error during Tavily search: {str(e)}"]
+
 @tool
 def get_article_nominator_from_fac_page(title: str) -> str:
     """
@@ -196,6 +220,9 @@ def route_question(question: str) -> str:
     """
     q = question.lower()
 
+    if "who" in q or "what" in q or "how many" in q or "when" in q:
+        return "tavily_search"
+
     if "wikipedia" in q and any(k in q for k in ["how many", "list", "albums", "awards", "release"]):
         return "extract_structured_facts_from_url"
 
@@ -417,6 +444,8 @@ def filter_vegetables(items: list[str]) -> list[str]:
 
 # List of all tools
 all_tools = [
+    extract_number_from_snippets,
+    tavily_search,
     route_question,
     resolve_wikipedia_url,
     handle_question,