Update tools.py

tools.py

@@ -2,7 +2,7 @@ from langchain_core.tools import tool
 import pandas as pd
 import os
 import re
-import wikipedia
+import requests
 
 
 @tool
@@ -52,38 +52,20 @@ def divide(a: int, b: int) -> int:
     return a / b
 
 @tool
-def wikipedia_search_tool(query: str, year_range: str = None) -> str:
+def wikidata_search_tool(sparql_query: str) -> str:
     """
-        Deep Wikipedia search with optional year filtering.
-        - Retrieves full text from multiple related pages.
-        - Filters for year mentions if year_range is specified (e.g. '1965-1968').
+       A tool that allows to search in WikiData for specific information
+       and returns the information as a json in string format
 
         Args:
-            query: The query to search in wikipedia
-            year_range: The range of years to search information for. Only use if necessary
+            sparql_query: The query to search in wikidata. This must be a SPARQL query.
 
         Example:
             {"message": "How many albums did the Beatles produce between 1965 and 1968?"}
-            wikipedia_search_tool("The Beatles albums", "1965-1968")
+            wikidata_search_tool("The Beatles albums")
     """
-    try:
-        results = wikipedia.search(query, results=5)
-        if not results:
-            return "No relevant pages found."
-
-        combined = ""
-        for title in results:
-            try:
-                page = wikipedia.page(title)
-                content = clean_and_truncate(page.content)
-                if year_range:
-                    y1, y2 = map(int, year_range.split('-'))
-                    content = extract_year_range(content, y1, y2)
-                if content:
-                    combined += f"\n\n## {title}\n{content}"
-            except:
-                continue
-
-        return combined.strip() if combined else "No content matched the year filter."
-    except Exception as e:
-        return f"Error: {str(e)}"
+    endpoint_url = "https://query.wikidata.org/sparql"
+    headers = {"Accept": "application/sparql-results+json"}
+    response = requests.get(endpoint_url, headers=headers, params={"query": sparql_query})
+    response.raise_for_status()
+    return response.json()