cyanite.py

import logging
import os
import json
import requests

CYANITE_API_URL = "https://api.cyanite.ai/graphql"
CYANITE_ACCESS_TOKEN = os.getenv("CYANITE_ACCESS_TOKEN")

def free_text_search(search_text, num_tracks=5):
    headers = {
        "Authorization": f"Bearer {CYANITE_ACCESS_TOKEN}",
        "Content-Type": "application/json"
    }

    query = '''
    query FreeTextSearch($searchText: String!, $numTracks: Int!) {
      freeTextSearch(
        first: $numTracks
        target: { library: {} }
        searchText: $searchText
      ) {
        ... on FreeTextSearchError {
          message
          code
        }
        ... on FreeTextSearchConnection {
          edges {
            cursor
            node {
              id
              title
            }
          }
        }
      }
    }
    '''

    variables = {
        "searchText": search_text,
        "numTracks": num_tracks
    }
    import time

    start_time = time.time()

    response = requests.post(
      CYANITE_API_URL,
      headers=headers,
      json={'query': query, 'variables': variables}
    )

    end_time = time.time()
    time_taken = end_time - start_time
    logging.warning(f"Cyanite API: Time taken: {time_taken} seconds")

    if response.status_code == 200:
        songs = extract_songs_from_response(response.json())
        if songs:
            return songs
        else:
            raise Exception("No songs found")
    else:
        raise Exception(f"Query failed with status code {response.status_code}")

def extract_songs_from_response(response_json):
    try:
        edges = response_json['data']['freeTextSearch']['edges']
        if not edges:
            return None  # No songs found
        songs = [{"id": edge["node"]["id"], "title": edge["node"]["title"]} for edge in edges]
        return songs
    except KeyError:
        raise Exception("Invalid response format")