feat: orchestrate grounded recommendations
This commit is contained in:
Justin Visser
parent
751391e6a2
commit
cead39edbc
8 changed files with 1098 additions and 7 deletions
255
backend/app/pipeline/orchestrator.py
Normal file
255
backend/app/pipeline/orchestrator.py
Normal file
|
|
@ -0,0 +1,255 @@
|
|||
"""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),
|
||||
)
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue