Update agent.py

agent.py

@@ -97,6 +97,7 @@ 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.
@@ -111,6 +112,22 @@ def wiki_search(query: str) -> str:
         ])
     return {"wiki_results": formatted_search_docs}
 
+
+    
+@tool
+def wikidata_query(query: str) -> str:
+    """
+    Run a SPARQL query on Wikidata and return results.
+    """
+    endpoint_url = "https://query.wikidata.org/sparql"
+    headers = {
+        "Accept": "application/sparql-results+json"
+    }
+    response = requests.get(endpoint_url, headers=headers, params={"query": query})
+    data = response.json()
+    return json.dumps(data, indent=2)
+
+
 @tool
 def web_search(query: str) -> str:
     """Search Tavily for a query and return maximum 3 results.
@@ -260,6 +277,7 @@ tool_map = {
     "get_youtube_transcript": get_youtube_transcript,
     "extract_video_id": extract_video_id,
     "analyze_attachment": analyze_attachment,
+    "wikidata_query": wikidata_query
     
 }