"""Compose taste, intent, grounding, and reranking into streamed events.""" import asyncio import time from collections.abc import AsyncGenerator from contextlib import aclosing from dataclasses import dataclass import structlog from app.config import Settings from app.domain.models import ( CompressedTasteProfile, ConversationTurn, Intent, PreviousRecommendation, TasteProfile, Track, ) from app.domain.profile import compress_taste_profile from app.observability.timing import increment_cache_hits from app.pipeline.event import ( PipelineDoneEvent, PipelineErrorEvent, PipelineEvent, PipelineMetadataEvent, PipelineTrackEvent, PipelineWarningEvent, ) from app.pipeline.grounding import Grounder from app.ports.protocols import MusicCatalog, Recommender, RecommenderOutputError RERANK_FALLBACK_CODE = "rerank_fallback" RERANK_FALLBACK_MESSAGE = "Ranking output was invalid, so grounded results are shown instead." RERANK_FALLBACK_JUSTIFICATION = "Selected as a grounded match for your request." @dataclass class _RankingState: selected_ids: set[str] selected_tracks: list[Track] @property def track_count(self) -> int: return len(self.selected_tracks) class RecommendationPipeline: """Orchestrate the code-defined recommendation stages.""" def __init__( self, recommender: Recommender, settings: Settings, grounder: Grounder | None = None, ) -> None: """Create process-local caches around the provided service ports.""" self.recommender = recommender self.settings = settings self.grounder = grounder or Grounder(settings) self.taste_cache = _TasteProfileCache(settings) self.last_pools: dict[str, tuple[Track, ...]] = {} async def stream( self, session_id: str, request_id: str, catalog: MusicCatalog, query: str, history: tuple[ConversationTurn, ...], previous_recommendations: tuple[PreviousRecommendation, ...], ) -> AsyncGenerator[PipelineEvent]: """Yield ordered events for one recommendation request.""" started_at = time.monotonic() taste = await self.taste_cache.get(session_id, catalog) intent = await self.recommender.create_intent( query, history, previous_recommendations, taste.text, self.settings.candidate_count, ) yield PipelineMetadataEvent( request_id=request_id, intent_summary=intent.intent_summary, candidate_count=len(intent.candidates), ) pool = self.last_pools.get(session_id) if intent.is_refinement else None if pool is None: result = await self.grounder.ground( catalog, intent.candidates, taste.known_track_ids, intent.familiarity, self.settings.rerank_count + self.settings.rerank_pool_buffer, ) pool = result.tracks if pool: self.last_pools[session_id] = pool if len(pool) < self.settings.grounding_floor: yield PipelineErrorEvent( code="insufficient_grounding", message="Not enough requested tracks could be verified safely.", ) return state = _RankingState(selected_ids=set(), selected_tracks=[]) async for event in self._stream_ranking(intent, pool, taste.text, history, state): yield event new_track_count = sum( track.id not in taste.known_track_ids for track in state.selected_tracks ) structlog.get_logger().info( "recommendations_complete", track_count=state.track_count, new_track_count=new_track_count, ) yield PipelineDoneEvent( track_count=state.track_count, total_ms=round((time.monotonic() - started_at) * 1000), ) async def _stream_ranking( self, intent: Intent, pool: tuple[Track, ...], taste_summary: str, history: tuple[ConversationTurn, ...], state: _RankingState, ) -> AsyncGenerator[PipelineTrackEvent | PipelineWarningEvent]: correction: str | None = None for _attempt_number in range(2): remaining_count = self.settings.rerank_count - state.track_count if remaining_count <= 0: return try: async for event in self._validated_rerank( intent, pool, taste_summary, history, remaining_count, correction, state, ): yield event return except RecommenderOutputError as error: emitted_ids = ", ".join(sorted(state.selected_ids)) or "none" correction = ( f"Validation failed: {error}. Already emitted track ids: {emitted_ids}." ) yield PipelineWarningEvent(code=RERANK_FALLBACK_CODE, message=RERANK_FALLBACK_MESSAGE) for track in pool: if state.track_count >= self.settings.rerank_count: break if track.id in state.selected_ids: continue state.selected_ids.add(track.id) state.selected_tracks.append(track) yield PipelineTrackEvent( rank=state.track_count, track=track, justification=RERANK_FALLBACK_JUSTIFICATION, ) async def _validated_rerank( self, intent: Intent, pool: tuple[Track, ...], taste_summary: str, history: tuple[ConversationTurn, ...], selection_count: int, correction: str | None, state: _RankingState, ) -> AsyncGenerator[PipelineTrackEvent]: tracks_by_id = {track.id: track for track in pool} stream = self.recommender.stream_rerank( intent, pool, taste_summary, history, selection_count, correction, ) async with aclosing(stream) as selections: async for selection in selections: if state.track_count >= self.settings.rerank_count: raise RecommenderOutputError("Rerank returned too many track ids") track = tracks_by_id.get(selection.track_id) if track is None: raise RecommenderOutputError("Rerank selected an out-of-pool track id") if track.id in state.selected_ids: raise RecommenderOutputError("Rerank selected a duplicate track id") state.selected_ids.add(track.id) state.selected_tracks.append(track) yield PipelineTrackEvent( rank=state.track_count, track=track, justification=selection.justification, ) class _TasteProfileCache: def __init__(self, settings: Settings) -> None: self.settings = settings self._entries: dict[str, tuple[float, CompressedTasteProfile]] = {} self._locks: dict[str, asyncio.Lock] = {} async def get(self, session_id: str, catalog: MusicCatalog) -> CompressedTasteProfile: cached = self._fresh_entry(session_id) if cached is not None: increment_cache_hits() return cached lock = self._locks.setdefault(session_id, asyncio.Lock()) async with lock: cached = self._fresh_entry(session_id) if cached is not None: increment_cache_hits() return cached compressed = await self._fetch(catalog) self._entries[session_id] = (time.monotonic(), compressed) return compressed def _fresh_entry(self, session_id: str) -> CompressedTasteProfile | None: entry = self._entries.get(session_id) if entry is None: return None created_at, profile = entry if time.monotonic() - created_at >= self.settings.taste_profile_ttl_seconds: del self._entries[session_id] return None return profile async def _fetch(self, catalog: MusicCatalog) -> CompressedTasteProfile: short_artists, long_artists, short_tracks, long_tracks, saved_tracks = await asyncio.gather( catalog.fetch_top_artists("short_term", self.settings.top_items_limit), catalog.fetch_top_artists("long_term", self.settings.top_items_limit), catalog.fetch_top_tracks("short_term", self.settings.top_items_limit), catalog.fetch_top_tracks("long_term", self.settings.top_items_limit), catalog.fetch_saved_tracks(self.settings.saved_tracks_limit), ) return compress_taste_profile( TasteProfile( short_term_artists=tuple(short_artists), long_term_artists=tuple(long_artists), short_term_tracks=tuple(short_tracks), long_term_tracks=tuple(long_tracks), saved_tracks=tuple(saved_tracks), ) )