refactor: rebuild orchestration around a track selection
This commit is contained in:
Justin Visser
parent
cead39edbc
commit
e8d20158e3
2 changed files with 124 additions and 96 deletions
|
|
@ -2,9 +2,8 @@
|
|||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import AsyncGenerator
|
||||
from collections.abc import AsyncGenerator, Iterator
|
||||
from contextlib import aclosing
|
||||
from dataclasses import dataclass
|
||||
|
||||
import structlog
|
||||
|
||||
|
|
@ -35,14 +34,58 @@ RERANK_FALLBACK_MESSAGE = "Ranking output was invalid, so grounded results are s
|
|||
RERANK_FALLBACK_JUSTIFICATION = "Selected as a grounded match for your request."
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RankingState:
|
||||
selected_ids: set[str]
|
||||
selected_tracks: list[Track]
|
||||
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 track_count(self) -> int:
|
||||
return len(self.selected_tracks)
|
||||
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:
|
||||
|
|
@ -86,19 +129,7 @@ class RecommendationPipeline:
|
|||
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
|
||||
|
||||
pool = await self._grounded_pool(session_id, catalog, intent, taste)
|
||||
if len(pool) < self.settings.grounding_floor:
|
||||
yield PipelineErrorEvent(
|
||||
code="insufficient_grounding",
|
||||
|
|
@ -106,103 +137,96 @@ class RecommendationPipeline:
|
|||
)
|
||||
return
|
||||
|
||||
state = _RankingState(selected_ids=set(), selected_tracks=[])
|
||||
async for event in self._stream_ranking(intent, pool, taste.text, history, state):
|
||||
selection = _TrackSelection(pool, self.settings.rerank_count)
|
||||
async for event in self._ranked_events(intent, taste.text, history, selection):
|
||||
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,
|
||||
)
|
||||
_log_completion(selection, taste)
|
||||
yield PipelineDoneEvent(
|
||||
track_count=state.track_count,
|
||||
track_count=len(selection.selected),
|
||||
total_ms=round((time.monotonic() - started_at) * 1000),
|
||||
)
|
||||
|
||||
async def _stream_ranking(
|
||||
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,
|
||||
pool: tuple[Track, ...],
|
||||
taste_summary: str,
|
||||
history: tuple[ConversationTurn, ...],
|
||||
state: _RankingState,
|
||||
selection: _TrackSelection,
|
||||
) -> AsyncGenerator[PipelineTrackEvent | PipelineWarningEvent]:
|
||||
"""Stream one rerank, retry once on invalid output, then fall back."""
|
||||
correction: str | None = None
|
||||
for _attempt_number in range(2):
|
||||
remaining_count = self.settings.rerank_count - state.track_count
|
||||
if remaining_count <= 0:
|
||||
for _ in range(2):
|
||||
if selection.is_full:
|
||||
return
|
||||
try:
|
||||
async for event in self._validated_rerank(
|
||||
intent,
|
||||
pool,
|
||||
taste_summary,
|
||||
history,
|
||||
remaining_count,
|
||||
correction,
|
||||
state,
|
||||
async for event in self._rerank_once(
|
||||
intent, taste_summary, history, selection, correction
|
||||
):
|
||||
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}."
|
||||
f"Validation failed: {error}."
|
||||
f" Already emitted track ids: {selection.describe_selected_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,
|
||||
)
|
||||
for event in selection.fill_from_pool(RERANK_FALLBACK_JUSTIFICATION):
|
||||
yield event
|
||||
|
||||
async def _validated_rerank(
|
||||
async def _rerank_once(
|
||||
self,
|
||||
intent: Intent,
|
||||
pool: tuple[Track, ...],
|
||||
taste_summary: str,
|
||||
history: tuple[ConversationTurn, ...],
|
||||
selection_count: int,
|
||||
selection: _TrackSelection,
|
||||
correction: str | None,
|
||||
state: _RankingState,
|
||||
) -> AsyncGenerator[PipelineTrackEvent]:
|
||||
tracks_by_id = {track.id: track for track in pool}
|
||||
stream = self.recommender.stream_rerank(
|
||||
intent,
|
||||
pool,
|
||||
selection.pool,
|
||||
taste_summary,
|
||||
history,
|
||||
selection_count,
|
||||
selection.remaining_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,
|
||||
)
|
||||
async for item in selections:
|
||||
yield selection.select(item.track_id, item.justification)
|
||||
|
||||
|
||||
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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue