301 lines
11 KiB
Python
301 lines
11 KiB
Python
"""Compose taste, intent, grounding, and reranking into streamed events."""
|
|
|
|
import asyncio
|
|
import time
|
|
from collections.abc import AsyncGenerator, Iterator
|
|
from contextlib import aclosing
|
|
|
|
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."
|
|
|
|
|
|
class _TrackSelection:
|
|
"""Rank tracks from one grounded pool, enforcing bound and uniqueness."""
|
|
|
|
def __init__(self, pool: tuple[Track, ...], limit: int) -> None:
|
|
"""Bind the only tracks that may ever be selected."""
|
|
self.pool = pool
|
|
self.limit = limit
|
|
self.selected: list[Track] = []
|
|
self._tracks_by_id = {track.id: track for track in pool}
|
|
self._selected_ids: set[str] = set()
|
|
|
|
@property
|
|
def is_full(self) -> bool:
|
|
"""Return whether the selection reached its limit."""
|
|
return len(self.selected) >= self.limit
|
|
|
|
@property
|
|
def remaining_count(self) -> int:
|
|
"""Return how many further selections are allowed."""
|
|
return self.limit - len(self.selected)
|
|
|
|
def describe_selected_ids(self) -> str:
|
|
"""Render the ids already emitted, for a correction instruction."""
|
|
return ", ".join(sorted(self._selected_ids)) or "none"
|
|
|
|
def select(self, track_id: str, justification: str) -> PipelineTrackEvent:
|
|
"""Accept one recommender selection or reject it as invalid output."""
|
|
if self.is_full:
|
|
raise RecommenderOutputError("Rerank returned too many track ids")
|
|
track = self._tracks_by_id.get(track_id)
|
|
if track is None:
|
|
raise RecommenderOutputError("Rerank selected an out-of-pool track id")
|
|
if track_id in self._selected_ids:
|
|
raise RecommenderOutputError("Rerank selected a duplicate track id")
|
|
return self._emit(track, justification)
|
|
|
|
def fill_from_pool(self, justification: str) -> Iterator[PipelineTrackEvent]:
|
|
"""Complete the selection in pool order after a failed rerank."""
|
|
for track in self.pool:
|
|
if self.is_full:
|
|
return
|
|
if track.id not in self._selected_ids:
|
|
yield self._emit(track, justification)
|
|
|
|
def _emit(self, track: Track, justification: str) -> PipelineTrackEvent:
|
|
self._selected_ids.add(track.id)
|
|
self.selected.append(track)
|
|
return PipelineTrackEvent(
|
|
rank=len(self.selected),
|
|
track=track,
|
|
justification=justification,
|
|
)
|
|
|
|
|
|
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 = await self._grounded_pool(session_id, catalog, intent, taste)
|
|
if len(pool) < self.settings.grounding_floor:
|
|
yield PipelineErrorEvent(
|
|
code="insufficient_grounding",
|
|
message="Not enough requested tracks could be verified safely.",
|
|
)
|
|
return
|
|
|
|
selection = _TrackSelection(pool, self.settings.rerank_count)
|
|
async for event in self._ranked_events(intent, taste.text, history, selection):
|
|
yield event
|
|
|
|
_log_completion(selection, taste)
|
|
yield PipelineDoneEvent(
|
|
track_count=len(selection.selected),
|
|
total_ms=round((time.monotonic() - started_at) * 1000),
|
|
)
|
|
|
|
async def _grounded_pool(
|
|
self,
|
|
session_id: str,
|
|
catalog: MusicCatalog,
|
|
intent: Intent,
|
|
taste: CompressedTasteProfile,
|
|
) -> tuple[Track, ...]:
|
|
"""Reuse the session's pool on refinement, otherwise ground anew."""
|
|
if intent.is_refinement:
|
|
cached_pool = self.last_pools.get(session_id)
|
|
if cached_pool:
|
|
return cached_pool
|
|
result = await self.grounder.ground(
|
|
catalog,
|
|
intent.candidates,
|
|
taste.known_track_ids,
|
|
intent.familiarity,
|
|
self.settings.rerank_count + self.settings.rerank_pool_buffer,
|
|
)
|
|
if result.tracks:
|
|
self.last_pools[session_id] = result.tracks
|
|
return result.tracks
|
|
|
|
async def _ranked_events(
|
|
self,
|
|
intent: Intent,
|
|
taste_summary: str,
|
|
history: tuple[ConversationTurn, ...],
|
|
selection: _TrackSelection,
|
|
) -> AsyncGenerator[PipelineTrackEvent | PipelineWarningEvent]:
|
|
"""Stream the rerank with one corrected retry, then fall back."""
|
|
try:
|
|
async for event in self._rerank_with_one_retry(
|
|
intent, taste_summary, history, selection
|
|
):
|
|
yield event
|
|
return
|
|
except RecommenderOutputError as error:
|
|
_log_rerank_failure(attempt=2, error=error)
|
|
|
|
yield PipelineWarningEvent(code=RERANK_FALLBACK_CODE, message=RERANK_FALLBACK_MESSAGE)
|
|
for event in selection.fill_from_pool(RERANK_FALLBACK_JUSTIFICATION):
|
|
yield event
|
|
|
|
async def _rerank_with_one_retry(
|
|
self,
|
|
intent: Intent,
|
|
taste_summary: str,
|
|
history: tuple[ConversationTurn, ...],
|
|
selection: _TrackSelection,
|
|
) -> AsyncGenerator[PipelineTrackEvent]:
|
|
"""Rerank once; on invalid output, retry once with a correction."""
|
|
try:
|
|
async for event in self._rerank_once(intent, taste_summary, history, selection, None):
|
|
yield event
|
|
return
|
|
except RecommenderOutputError as error:
|
|
_log_rerank_failure(attempt=1, error=error)
|
|
if selection.is_full:
|
|
return
|
|
correction = (
|
|
f"Validation failed: {error}."
|
|
f" Already emitted track ids: {selection.describe_selected_ids()}."
|
|
)
|
|
|
|
async for event in self._rerank_once(intent, taste_summary, history, selection, correction):
|
|
yield event
|
|
|
|
async def _rerank_once(
|
|
self,
|
|
intent: Intent,
|
|
taste_summary: str,
|
|
history: tuple[ConversationTurn, ...],
|
|
selection: _TrackSelection,
|
|
correction: str | None,
|
|
) -> AsyncGenerator[PipelineTrackEvent]:
|
|
stream = self.recommender.stream_rerank(
|
|
intent,
|
|
selection.pool,
|
|
taste_summary,
|
|
history,
|
|
selection.remaining_count,
|
|
correction,
|
|
)
|
|
async with aclosing(stream) as selections:
|
|
async for item in selections:
|
|
yield selection.select(item.track_id, item.justification)
|
|
|
|
|
|
def _log_rerank_failure(attempt: int, error: RecommenderOutputError) -> None:
|
|
"""Log one invalid rerank attempt with its validation reason."""
|
|
structlog.get_logger().warning("rerank_attempt_failed", attempt=attempt, error=str(error))
|
|
|
|
|
|
def _log_completion(selection: _TrackSelection, taste: CompressedTasteProfile) -> None:
|
|
"""Log how many recommendations were served and how many are new."""
|
|
new_track_count = sum(track.id not in taste.known_track_ids for track in selection.selected)
|
|
structlog.get_logger().info(
|
|
"recommendations_complete",
|
|
track_count=len(selection.selected),
|
|
new_track_count=new_track_count,
|
|
)
|
|
|
|
|
|
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),
|
|
)
|
|
)
|