339 lines
12 KiB
Python
339 lines
12 KiB
Python
"""Anthropic implementation of structured intent and streamed reranking."""
|
|
|
|
import json
|
|
import re
|
|
from collections.abc import AsyncIterator
|
|
from typing import Literal, cast
|
|
|
|
from anthropic import AsyncAnthropic
|
|
from anthropic.lib._parse._transform import transform_schema
|
|
from anthropic.types import Message, TextBlockParam, Usage
|
|
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
|
|
|
from app.config import Settings
|
|
from app.domain.models import (
|
|
ConversationTurn,
|
|
Familiarity,
|
|
Intent,
|
|
PreviousRecommendation,
|
|
RerankSelection,
|
|
Track,
|
|
TrackCandidate,
|
|
)
|
|
from app.observability.timing import record_llm_tokens
|
|
from app.ports.protocols import RecommenderOutputError
|
|
from app.prompts import INTENT_SYSTEM_PROMPT, RERANK_SYSTEM_PROMPT
|
|
|
|
Effort = Literal["low", "medium", "high", "xhigh", "max"]
|
|
_RECOMMENDATION_ARRAY = re.compile(r'"recommendations"\s*:\s*\[')
|
|
|
|
|
|
class CandidateOutput(BaseModel):
|
|
"""One bounded candidate in the intent response."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
title: str = Field(min_length=1, max_length=200)
|
|
artist: str = Field(min_length=1, max_length=200)
|
|
|
|
|
|
class IntentOutput(BaseModel):
|
|
"""Bounded structured output for the intent call."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
mood: list[str] = Field(default_factory=list, max_length=6)
|
|
activity: str | None = Field(default=None, max_length=100)
|
|
era: list[str] = Field(default_factory=list, max_length=5)
|
|
languages: list[str] = Field(default_factory=list, max_length=8)
|
|
genres: list[str] = Field(default_factory=list, max_length=8)
|
|
familiarity: Familiarity
|
|
is_refinement: bool
|
|
intent_summary: str = Field(min_length=1, max_length=300, pattern=r"^[^\r\n]+$")
|
|
candidates: list[CandidateOutput] = Field(min_length=30, max_length=40)
|
|
|
|
|
|
class RerankSelectionOutput(BaseModel):
|
|
"""One validated object extracted from the rerank stream."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
track_id: str = Field(min_length=1, max_length=200)
|
|
justification: str = Field(min_length=1, max_length=300, pattern=r"^[^\r\n]+$")
|
|
|
|
|
|
class RerankOutput(BaseModel):
|
|
"""Complete bounded rerank output used for final validation."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
recommendations: list[RerankSelectionOutput] = Field(min_length=1, max_length=50)
|
|
|
|
|
|
class AnthropicRecommender:
|
|
"""Use two Anthropic calls for intent generation and grounded ranking."""
|
|
|
|
def __init__(self, client: AsyncAnthropic, settings: Settings) -> None:
|
|
"""Bind the shared asynchronous client and immutable settings."""
|
|
self.client = client
|
|
self.settings = settings
|
|
|
|
async def create_intent(
|
|
self,
|
|
query: str,
|
|
history: tuple[ConversationTurn, ...],
|
|
previous_recommendations: tuple[PreviousRecommendation, ...],
|
|
taste_summary: str,
|
|
candidate_count: int,
|
|
) -> Intent:
|
|
"""Interpret a request through Anthropic structured output."""
|
|
message = await self.client.messages.parse(
|
|
model=self.settings.llm_model,
|
|
max_tokens=self.settings.intent_max_tokens,
|
|
output_config={"effort": _parse_effort(self.settings.intent_effort)},
|
|
output_format=IntentOutput,
|
|
system=[_system_block(INTENT_SYSTEM_PROMPT)],
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": _render_intent_input(
|
|
query,
|
|
history,
|
|
previous_recommendations,
|
|
taste_summary,
|
|
candidate_count,
|
|
),
|
|
}
|
|
],
|
|
)
|
|
_validate_stop_reason(message)
|
|
_record_usage(message.usage)
|
|
parsed = message.parsed_output
|
|
if parsed is None:
|
|
raise RecommenderOutputError("Intent response contained no structured output")
|
|
validated = IntentOutput.model_validate(parsed.model_dump())
|
|
if len(validated.candidates) != candidate_count:
|
|
raise RecommenderOutputError("Intent response returned the wrong candidate count")
|
|
return _to_intent(validated)
|
|
|
|
async def stream_rerank(
|
|
self,
|
|
intent: Intent,
|
|
grounded_tracks: tuple[Track, ...],
|
|
taste_summary: str,
|
|
history: tuple[ConversationTurn, ...],
|
|
selection_count: int,
|
|
correction: str | None = None,
|
|
) -> AsyncIterator[RerankSelection]:
|
|
"""Yield each complete valid selection while the JSON is streaming."""
|
|
schema = transform_schema(RerankOutput.model_json_schema())
|
|
parser = _RecommendationObjectParser()
|
|
async with self.client.messages.stream(
|
|
model=self.settings.llm_model,
|
|
max_tokens=self.settings.rerank_max_tokens,
|
|
output_config={
|
|
"effort": _parse_effort(self.settings.rerank_effort),
|
|
"format": {"type": "json_schema", "schema": schema},
|
|
},
|
|
system=[_system_block(RERANK_SYSTEM_PROMPT)],
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": _render_rerank_input(
|
|
intent,
|
|
grounded_tracks,
|
|
taste_summary,
|
|
history,
|
|
selection_count,
|
|
correction,
|
|
),
|
|
}
|
|
],
|
|
) as stream:
|
|
async for text_delta in stream.text_stream:
|
|
for selection in parser.feed(text_delta):
|
|
yield RerankSelection(
|
|
track_id=selection.track_id,
|
|
justification=selection.justification,
|
|
)
|
|
final_message = await stream.get_final_message()
|
|
|
|
_validate_stop_reason(final_message)
|
|
_record_usage(final_message.usage)
|
|
try:
|
|
validated = RerankOutput.model_validate_json(parser.complete_text)
|
|
except ValidationError as error:
|
|
raise RecommenderOutputError("Rerank response failed final validation") from error
|
|
if len(validated.recommendations) > selection_count:
|
|
raise RecommenderOutputError("Rerank response returned too many selections")
|
|
|
|
|
|
class _RecommendationObjectParser:
|
|
def __init__(self) -> None:
|
|
self.complete_text = ""
|
|
self._scan_index = 0
|
|
self._object_start: int | None = None
|
|
self._object_depth = 0
|
|
self._is_in_string = False
|
|
self._is_escaped = False
|
|
self._has_found_array = False
|
|
|
|
def feed(self, text_delta: str) -> list[RerankSelectionOutput]:
|
|
self.complete_text += text_delta
|
|
if not self._has_found_array:
|
|
match = _RECOMMENDATION_ARRAY.search(self.complete_text)
|
|
if match is None:
|
|
return []
|
|
self._has_found_array = True
|
|
self._scan_index = match.end()
|
|
|
|
selections: list[RerankSelectionOutput] = []
|
|
while self._scan_index < len(self.complete_text):
|
|
character = self.complete_text[self._scan_index]
|
|
completed = self._scan_character(character)
|
|
self._scan_index += 1
|
|
if completed is not None:
|
|
selections.append(completed)
|
|
return selections
|
|
|
|
def _scan_character(self, character: str) -> RerankSelectionOutput | None:
|
|
if self._object_start is None:
|
|
if character == "{":
|
|
self._object_start = self._scan_index
|
|
self._object_depth = 1
|
|
return None
|
|
|
|
if self._is_in_string:
|
|
if self._is_escaped:
|
|
self._is_escaped = False
|
|
elif character == "\\":
|
|
self._is_escaped = True
|
|
elif character == '"':
|
|
self._is_in_string = False
|
|
return None
|
|
|
|
if character == '"':
|
|
self._is_in_string = True
|
|
elif character == "{":
|
|
self._object_depth += 1
|
|
elif character == "}":
|
|
self._object_depth -= 1
|
|
if self._object_depth == 0:
|
|
return self._finish_object()
|
|
return None
|
|
|
|
def _finish_object(self) -> RerankSelectionOutput:
|
|
assert self._object_start is not None
|
|
object_text = self.complete_text[self._object_start : self._scan_index + 1]
|
|
self._object_start = None
|
|
try:
|
|
return RerankSelectionOutput.model_validate_json(object_text)
|
|
except ValidationError as error:
|
|
raise RecommenderOutputError("Rerank item failed validation") from error
|
|
|
|
|
|
def _to_intent(output: IntentOutput) -> Intent:
|
|
return Intent(
|
|
mood=tuple(output.mood),
|
|
activity=output.activity,
|
|
era=tuple(output.era),
|
|
languages=tuple(output.languages),
|
|
genres=tuple(output.genres),
|
|
familiarity=output.familiarity,
|
|
is_refinement=output.is_refinement,
|
|
intent_summary=output.intent_summary,
|
|
candidates=tuple(
|
|
TrackCandidate(title=candidate.title, artist=candidate.artist)
|
|
for candidate in output.candidates
|
|
),
|
|
)
|
|
|
|
|
|
def _render_intent_input(
|
|
query: str,
|
|
history: tuple[ConversationTurn, ...],
|
|
previous_recommendations: tuple[PreviousRecommendation, ...],
|
|
taste_summary: str,
|
|
candidate_count: int,
|
|
) -> str:
|
|
payload = {
|
|
"query": query,
|
|
"history": [turn.__dict__ for turn in history],
|
|
"prior_recommendations": [
|
|
recommendation.__dict__ for recommendation in previous_recommendations
|
|
],
|
|
"taste_profile": taste_summary,
|
|
"required_candidate_count": candidate_count,
|
|
}
|
|
return json.dumps(payload, ensure_ascii=True, separators=(",", ":"))
|
|
|
|
|
|
def _render_rerank_input(
|
|
intent: Intent,
|
|
grounded_tracks: tuple[Track, ...],
|
|
taste_summary: str,
|
|
history: tuple[ConversationTurn, ...],
|
|
selection_count: int,
|
|
correction: str | None,
|
|
) -> str:
|
|
payload = {
|
|
"intent": {
|
|
"mood": intent.mood,
|
|
"activity": intent.activity,
|
|
"era": intent.era,
|
|
"languages": intent.languages,
|
|
"genres": intent.genres,
|
|
"familiarity": intent.familiarity,
|
|
"intent_summary": intent.intent_summary,
|
|
},
|
|
"grounded_pool": [
|
|
{"track_id": track.id, "title": track.title, "artists": track.artists}
|
|
for track in grounded_tracks
|
|
],
|
|
"taste_profile": taste_summary,
|
|
"history": [turn.__dict__ for turn in history],
|
|
"requested_selection_count": selection_count,
|
|
"correction": correction,
|
|
}
|
|
return json.dumps(payload, ensure_ascii=True, separators=(",", ":"))
|
|
|
|
|
|
def _system_block(prompt: str) -> TextBlockParam:
|
|
return {
|
|
"type": "text",
|
|
"text": prompt,
|
|
"cache_control": {"type": "ephemeral"},
|
|
}
|
|
|
|
|
|
def _parse_effort(value: str) -> Effort:
|
|
allowed = {"low", "medium", "high", "xhigh", "max"}
|
|
if value not in allowed:
|
|
raise ValueError(f"Unsupported Anthropic effort: {value}")
|
|
return cast(Effort, value)
|
|
|
|
|
|
def _validate_stop_reason(message: Message) -> None:
|
|
stop_reason_messages = {
|
|
None: "Anthropic response had no stop reason",
|
|
"max_tokens": "Anthropic response reached its token limit",
|
|
"stop_sequence": "Anthropic response hit an unexpected stop sequence",
|
|
"tool_use": "Anthropic response attempted tool use",
|
|
"pause_turn": "Anthropic response paused before completion",
|
|
"refusal": "Anthropic response was refused",
|
|
"model_context_window_exceeded": "Anthropic context window was exceeded",
|
|
}
|
|
if message.stop_reason == "end_turn":
|
|
return
|
|
raise RecommenderOutputError(
|
|
stop_reason_messages.get(message.stop_reason, "Anthropic response stopped unexpectedly")
|
|
)
|
|
|
|
|
|
def _record_usage(usage: Usage) -> None:
|
|
cache_creation_tokens = usage.cache_creation_input_tokens or 0
|
|
cache_read_tokens = usage.cache_read_input_tokens or 0
|
|
record_llm_tokens(
|
|
usage.input_tokens + cache_creation_tokens + cache_read_tokens,
|
|
usage.output_tokens,
|
|
)
|