navidrome-music-sync/app/wave/llm_reranker.py

import json
import re
import time
from datetime import datetime
import urllib.request
import urllib.error
from typing import List, Dict, Any, Optional
from app.config import settings
from app.providers.base import TrackMetadata

class WaveLlmReranker:
    def __init__(self, ollama_url: Optional[str] = None, model: Optional[str] = None):
        self.ollama_url = (ollama_url or settings.OLLAMA_URL).rstrip("/")
        self.model = model or settings.OLLAMA_MODEL

    def get_time_of_day_context(self) -> Dict[str, str]:
        """Detect current time of day and musical context."""
        hour = datetime.now().hour
        if 5 <= hour < 12:
            return {
                "period": "утро",
                "label": "🌅 Утренний чилл",
                "desc": "мягкое плавное пробуждение, мелодичный чилл, акустическое или лёгкое звучание"
            }
        elif 12 <= hour < 18:
            return {
                "period": "день",
                "label": "☀️ Дневной фокус",
                "desc": "бодрый дневной ритм, концентрация, динамичный и сбалансированный флоу"
            }
        elif 18 <= hour < 23:
            return {
                "period": "вечер",
                "label": "🌆 Вечерний вайб",
                "desc": "атмосферный глубокий звук, эмоциональные треки, плотный бас и хип-хоп/рок"
            }
        else:
            return {
                "period": "ночь",
                "label": "🌙 Ночной андеграунд",
                "desc": "ночной драйв, меланхолия, глубокий бас, lo-fi и атмосферный андеграунд"
            }

    def resolve_mood_context(self, mood: str) -> Dict[str, str]:
        """Map user-selected mood string to descriptive context."""
        mood_lower = (mood or "auto").lower()
        if mood_lower == "energy":
            return {
                "period": "драйв",
                "label": "⚡ Энергия & Спорт",
                "desc": "взрывной ритм, максимальная динамика, бодрящий кач и быстрый темп"
            }
        elif mood_lower == "chill":
            return {
                "period": "релакс",
                "label": "☕ Спокойствие & Чилл",
                "desc": "расслабляющая музыка, мягкий ненавязчивый ритм, плавные переходы"
            }
        elif mood_lower == "dark":
            return {
                "period": "андеграунд",
                "label": "🏴 Мрачный вайб",
                "desc": "мрачный звук, глубокий саб-бас, андеграундный рэп и сырой плотный саунд"
            }
        # Default: auto by time of day
        return self.get_time_of_day_context()

    def rerank_candidates(
        self,
        candidates: List[TrackMetadata],
        anchor_profile: Dict[str, Any],
        mood: str = "auto",
        target_count: int = 25
    ) -> List[TrackMetadata]:
        """
        Rerank and select candidates using Ollama qwen2.5 model.
        Falls back smoothly to heuristic ranking if Ollama is unreachable or times out.
        """
        if not candidates:
            return []

        if len(candidates) <= target_count:
            return candidates

        context = self.resolve_mood_context(mood)
        
        # Build compact candidate list representation
        # Format: 1. Artist - Title
        candidate_lines = []
        for idx, c in enumerate(candidates, 1):
            candidate_lines.append(f"{idx}. {c.artist} - {c.title}")

        candidates_text = "\n".join(candidate_lines)

        # Extract anchor artist names for context
        top_artists = [a.get("artist") for a in anchor_profile.get("top_artists", [])[:5] if a.get("artist")]
        artists_str = ", ".join(top_artists) if top_artists else "Разнообразный русский и мировой рэп/рок"

        prompt = (
            f"Ты — опытный музыкальный куратор радио.\n"
            f"Контекст: {context['label']} ({context['desc']}).\n"
            f"Любимые артисты пользователя: {artists_str}.\n\n"
            f"Список кандидатов:\n{candidates_text}\n\n"
            f"Задача: выбери ровно {target_count} номеров треков из списка, которые лучше всего подходят под этот вайб "
            f"и обеспечат гармоничную последовательность прослушивания.\n"
            f"Ответь ИСКЛЮЧИТЕЛЬНО в формате JSON списка номеров, например: [1, 5, 8, 12, ...]. Без лишнего текста."
        )

        print(f"[WaveLLM] Sourcing {target_count} tracks via {self.model} for mood '{context['label']}'...")

        try:
            req_data = json.dumps({
                "model": self.model,
                "prompt": prompt,
                "stream": False,
                "keep_alive": 0,
                "options": {
                    "temperature": 0.25,
                    "num_predict": 150
                }
            }).encode("utf-8")

            req = urllib.request.Request(
                f"{self.ollama_url}/api/generate",
                data=req_data,
                headers={"Content-Type": "application/json"}
            )

            # 45 second timeout for local CPU inference
            with urllib.request.urlopen(req, timeout=45) as resp:
                resp_json = json.loads(resp.read().decode("utf-8"))
                raw_response = resp_json.get("response", "").strip()

            print(f"[WaveLLM] Model output: {raw_response[:100]}...")

            # Extract JSON list of indices from response
            match = re.search(r"\[([0-9,\s]+)\]", raw_response)
            if match:
                selected_indices = [int(x.strip()) for x in match.group(1).split(",") if x.strip().isdigit()]
                
                # Filter valid 1-based indices
                valid_tracks = []
                seen_indices = set()
                for idx in selected_indices:
                    if 1 <= idx <= len(candidates) and idx not in seen_indices:
                        seen_indices.add(idx)
                        valid_tracks.append(candidates[idx - 1])

                if len(valid_tracks) >= min(5, target_count):
                    # If model returned fewer than target_count, fill remaining from unselected
                    if len(valid_tracks) < target_count:
                        for idx, c in enumerate(candidates, 1):
                            if idx not in seen_indices:
                                valid_tracks.append(c)
                                if len(valid_tracks) >= target_count:
                                    break
                    print(f"[WaveLLM] Successfully reranked {len(valid_tracks)} tracks using {self.model}")
                    return valid_tracks[:target_count]

        except Exception as e:
            print(f"[WaveLLM] Notice: Ollama reranker fallback ({e}). Applying heuristic selection.")

        # Deadair Fallback Principle: Heuristic selection
        # Interleave tracks from diverse artists
        artist_buckets: Dict[str, List[TrackMetadata]] = {}
        for c in candidates:
            artist_buckets.setdefault(c.artist, []).append(c)

        heuristic_tracks: List[TrackMetadata] = []
        while len(heuristic_tracks) < target_count and any(artist_buckets.values()):
            for artist, tracks in list(artist_buckets.items()):
                if tracks and len(heuristic_tracks) < target_count:
                    heuristic_tracks.append(tracks.pop(0))

        return heuristic_tracks[:target_count]