Update tools.py

tools.py

@@ -1,8 +1,6 @@
 from langchain_core.tools import tool
-import pandas as pd
-import os
-import re
 import requests
+from datetime import datetime
 
 
 @tool
@@ -52,25 +50,56 @@ def divide(a: int, b: int) -> int:
     return a / b
 
 @tool
-def wikidata_search_tool(sparql_query: str) -> str:
+def get_studio_album_count(artist_name: str, start_year: int, end_year: int): -> dict
     """
-       A tool that allows to search in WikiData for specific information
-       and returns the information as a json in string format
-
-        Args:
-            sparql_query: The query to search in wikidata. This must be a SPARQL query.
+    Search music APIs for relevant information related to music artists.
+    It returns a dictionary with the number of albums produced during that interval and the names.
+    
+    Args:
+        artist_name: Name of the artist to search
+        start_year: The start year for searching information
+        end_year: The end year for searching information
 
-        Example:
-            {"message": "How many albums did the Beatles produce between 1965 and 1968?"}
-            wikidata_search_tool("SELECT (COUNT(DISTINCT ?album) AS ?albumCount) WHERE {
-                                  ?album wdt:P31 wd:Q482994;       # instance of album
-                                         wdt:P175 wd:Q1299;        # performer: The Beatles
-                                         wdt:P577 ?publicationDate. # publication date
-                                  FILTER(YEAR(?publicationDate) >= 1965 && YEAR(?publicationDate) <= 1968)
-                                }")
+    Exampe:
+        {"message": "How many albums did The Beatles produce between 1965 and 1968 ?"}
+        get_studio_album_count("The Beatles", 1965, 1968)
     """
-    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()
+    # Step 1: Get artist MBID
+    search_url = f"https://musicbrainz.org/ws/2/artist/?query=artist:{artist_name}&fmt=json"
+    artist_resp = requests.get(search_url).json()
+    if not artist_resp["artists"]:
+        return {"error": f"Artist '{artist_name}' not found."}
+    
+    artist_mbid = artist_resp["artists"][0]["id"]
+
+    # Step 2: Get release-groups (albums)
+    album_url = f"https://musicbrainz.org/ws/2/release-group?artist={artist_mbid}&type=album&fmt=json&limit=100"
+    albums_resp = requests.get(album_url).json()
+    
+    studio_albums = []
+    for release_group in albums_resp.get("release-groups", []):
+        title = release_group.get("title", "")
+        first_date = release_group.get("first-release-date", "")
+        
+        # Filter by date
+        try:
+            year = int(first_date[:4])
+        except:
+            continue
+        if not (start_year <= year <= end_year):
+            continue
+        
+        # Filter out secondary types (e.g. Live, Compilation)
+        secondary_types = release_group.get("secondary-types", [])
+        if secondary_types:
+            continue  # Not a studio album
+
+        studio_albums.append({
+            "title": title,
+            "release_year": year,
+        })
+    
+    return {
+        "studio_album_count": len(studio_albums),
+        "studio_album_list": sorted(studio_albums, key=lambda x: x["release_year"]),
+    }