funkwhale/api/funkwhale_api/radios/lb_recommendations.py

147 wiersze
4.4 KiB
Python

import logging
import time
import troi
import troi.core
from django.core.cache import cache
from django.core.exceptions import ValidationError
from django.db.models import Q
from requests.exceptions import ConnectTimeout
from funkwhale_api.music import models as music_models
from funkwhale_api.typesense import utils
logger = logging.getLogger(__name__)
patches = troi.utils.discover_patches()
SUPPORTED_PATCHES = patches.keys()
def run(config, **kwargs):
"""Validate the received config and run the queryset generation"""
candidates = kwargs.pop("candidates", music_models.Track.objects.all())
validate(config)
return TroiPatch().get_queryset(config, candidates)
def validate(config):
patch = config.get("patch")
if patch not in SUPPORTED_PATCHES:
raise ValidationError(
'Invalid patch "{}". Supported patches: {}'.format(
config["patch"], SUPPORTED_PATCHES
)
)
return True
def build_radio_queryset(patch, radio_qs):
"""Take a troi patch, match the missing mbid and then build a radio queryset"""
start_time = time.time()
try:
recommendations = patch.generate_playlist()
except ConnectTimeout:
raise ValueError(
"Timed out while connecting to ListenBrainz. No candidates could be retrieved for the radio."
)
end_time_rec = time.time()
logger.info("Troi fetch took :" + str(end_time_rec - start_time))
if not recommendations:
raise ValueError("No candidates found by troi")
recommended_mbids = [
recommended_recording.mbid
for recommended_recording in recommendations.playlists[0].recordings
]
logger.info("Searching for MusicBrainz ID in Funkwhale database")
qs_recommended = (
music_models.Track.objects.all()
.filter(mbid__in=recommended_mbids)
.order_by("mbid", "pk")
.distinct("mbid")
)
qs_recommended_mbid = [str(i.mbid) for i in qs_recommended]
recommended_mbids_not_qs = [
mbid for mbid in recommended_mbids if mbid not in qs_recommended_mbid
]
cached_match = cache.get_many(recommended_mbids_not_qs)
cached_match_mbid = [str(i) for i in cached_match.keys()]
if qs_recommended and cached_match_mbid:
logger.info("MusicBrainz IDs found in Funkwhale database and redis")
qs_recommended_mbid.extend(cached_match_mbid)
mbids_found = qs_recommended_mbid
elif qs_recommended and not cached_match_mbid:
logger.info("MusicBrainz IDs found in Funkwhale database")
mbids_found = qs_recommended_mbid
elif not qs_recommended and cached_match_mbid:
logger.info("MusicBrainz IDs found in redis cache")
mbids_found = cached_match_mbid
else:
logger.info(
"Couldn't find any matches in Funkwhale database. Trying to match all"
)
mbids_found = []
recommended_recordings_not_found = [
i for i in recommendations.playlists[0].recordings if i.mbid not in mbids_found
]
logger.info("Matching missing MusicBrainz ID to Funkwhale track")
start_time_resolv = time.time()
utils.resolve_recordings_to_fw_track(recommended_recordings_not_found)
end_time_resolv = time.time()
logger.info(
"Resolving "
+ str(len(recommended_recordings_not_found))
+ " tracks in "
+ str(end_time_resolv - start_time_resolv)
)
cached_match = cache.get_many(recommended_mbids)
if not mbids_found and not cached_match:
raise ValueError("No candidates found for troi radio")
mbids_found_pks = list(
music_models.Track.objects.all()
.filter(mbid__in=mbids_found)
.order_by("mbid", "pk")
.distinct("mbid")
.values_list("pk", flat=True)
)
mbids_found_pks_unique = [
i for i in mbids_found_pks if i not in cached_match.keys()
]
if mbids_found and cached_match:
return radio_qs.filter(
Q(pk__in=mbids_found_pks_unique) | Q(pk__in=cached_match.values())
)
if mbids_found and not cached_match:
return radio_qs.filter(pk__in=mbids_found_pks_unique)
if not mbids_found and cached_match:
return radio_qs.filter(pk__in=cached_match.values())
class TroiPatch:
code = "troi-patch"
label = "Troi Patch"
def get_queryset(self, config, qs):
patch_string = config.pop("patch")
patch = patches[patch_string]
return build_radio_queryset(patch(config), qs)