mirror of
https://github.com/run-llama/llama-hub.git
synced 2025-08-13 03:01:46 +00:00
68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
"""Spotify reader."""
|
|
|
|
from typing import List, Optional
|
|
|
|
from llama_index.readers.base import BaseReader
|
|
from llama_index.readers.schema.base import Document
|
|
|
|
|
|
class SpotifyReader(BaseReader):
|
|
"""Spotify Reader.
|
|
|
|
Read a user's saved albums, tracks, or playlists from Spotify.
|
|
|
|
"""
|
|
|
|
def load_data(self, collection: Optional[str] = "albums") -> List[Document]:
|
|
"""Load data from a user's Spotify account.
|
|
|
|
Args:
|
|
collections (Optional[str]): "albums", "tracks", or "playlists"
|
|
"""
|
|
|
|
import spotipy
|
|
from spotipy.oauth2 import SpotifyOAuth
|
|
|
|
scope = "user-library-read"
|
|
sp = spotipy.Spotify(auth_manager=SpotifyOAuth(scope=scope))
|
|
|
|
results = []
|
|
|
|
if collection == "albums":
|
|
response = sp.current_user_saved_albums()
|
|
items = response["items"]
|
|
for item in items:
|
|
album = item["album"]
|
|
album_name = album["name"]
|
|
artist_name = album["artists"][0]["name"]
|
|
album_string = f"Album {album_name} by Artist {artist_name}\n"
|
|
results.append(Document(album_string))
|
|
elif collection == "tracks":
|
|
response = sp.current_user_saved_tracks()
|
|
items = response["items"]
|
|
for item in items:
|
|
track = item["track"]
|
|
track_name = track["name"]
|
|
artist_name = track["artists"][0]["name"]
|
|
artist_string = f"Track {track_name} by Artist {artist_name}\n"
|
|
results.append(Document(artist_string))
|
|
elif collection == "playlists":
|
|
response = sp.current_user_playlists()
|
|
items = response["items"]
|
|
for item in items:
|
|
playlist_name = item["name"]
|
|
owner_name = item["owner"]["display_name"]
|
|
playlist_string = f"Playlist {playlist_name} created by {owner_name}\n"
|
|
results.append(Document(playlist_string))
|
|
else:
|
|
raise ValueError(
|
|
"Invalid collection parameter value. Allowed values are 'albums', 'tracks', or 'playlists'."
|
|
)
|
|
|
|
return results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
reader = SpotifyReader()
|
|
print(reader.load_data())
|