""" Post-call hook for the adaptive router. On each successful or failed completion, build a Turn from the request/response and push it through `AdaptiveRouter.record_turn`. The router then updates the in-memory bandit cell + session state and queues writes for the proxy flusher. All work happens after the response has been returned to the caller. Any exception is swallowed — signal recording must never break a request. """ from __future__ import annotations import hashlib import json from typing import Any, Dict, List, Optional from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.adaptive_router.config import ( ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, ADAPTIVE_ROUTER_RESPONSE_HEADER, SIGNAL_GATE_MIN_MESSAGES, ) from litellm.router_strategy.adaptive_router.signals import Turn # Identity fields hashed into a derived session key so the same conversation # from the same caller produces a stable key, while different keys/teams/users # stay segregated even if they happen to send identical first messages. _IDENTITY_FIELDS = ( "user_api_key_hash", "user_api_key_team_id", "user_api_key_user_id", "user_api_key_end_user_id", ) def _resolve_session_key(kwargs: Dict[str, Any]) -> Optional[str]: """Pick a stable per-conversation key for owner-cache attribution. Order: 1. Honor a client-supplied session id (`litellm_session_id` on either `litellm_params` or `litellm_params.metadata`, or `session_id` on metadata) — backward compat for callers already wired up. 2. Otherwise derive a sha256 over (identity fields, first SIGNAL_GATE_MIN_MESSAGES messages) so the key is stable across turns and only materialises once there is enough context for the bandit to act on (matching the gate in the signal-processing path). Returns None if the conversation is shorter than SIGNAL_GATE_MIN_MESSAGES. """ litellm_params = kwargs.get("litellm_params") or {} sid = litellm_params.get("litellm_session_id") if sid: return str(sid) metadata = litellm_params.get("metadata") or {} if isinstance(metadata, dict): sid = metadata.get("session_id") or metadata.get("litellm_session_id") if sid: return str(sid) messages = kwargs.get("messages") or [] if len(messages) < SIGNAL_GATE_MIN_MESSAGES: # Don't attribute until we have enough turns to match the signal gate — # ensures the hash is stable (same N messages every time) and avoids # crediting the bandit for conversations that are too short to signal. return None identity = ":".join( str(metadata.get(f) or "") if isinstance(metadata, dict) else "" for f in _IDENTITY_FIELDS ) anchor = messages[:SIGNAL_GATE_MIN_MESSAGES] payload = ( identity + "|" + json.dumps( [{"role": m.get("role"), "content": m.get("content")} for m in anchor], sort_keys=True, default=str, ) ) return hashlib.sha256(payload.encode("utf-8")).hexdigest() def _last_user_content(messages: Optional[List[Dict[str, Any]]]) -> Optional[str]: if not messages: return None for msg in reversed(messages): if msg.get("role") == "user": content = msg.get("content") if isinstance(content, str): return content if isinstance(content, list): # OpenAI vision-style content: pick first text part. for part in content: if isinstance(part, dict) and part.get("type") == "text": return part.get("text") return None return None def _recent_tool_results( messages: Optional[List[Dict[str, Any]]], ) -> List[Dict[str, Any]]: """Extract the current turn's tool result payloads from the request messages. Tool results are `role == "tool"` messages that sit at the tail of the conversation — i.e. after the most recent assistant message with `tool_calls`, waiting for the model to produce a user-facing reply. Walk backwards from the end and collect the contiguous run of tool messages; stop at the first non-tool message. Each result is normalized to `{content, is_error}` — the only fields `signals._detect_failure` / `_detect_exhaustion` actually read. """ if not messages: return [] results: List[Dict[str, Any]] = [] for msg in reversed(messages): if not isinstance(msg, dict): break if msg.get("role") != "tool": break content = msg.get("content") # Some providers (Anthropic-style) carry an explicit error flag; OpenAI # tool results don't, so fall back to an empty/missing content heuristic # inside `_detect_failure`. is_error = bool(msg.get("is_error")) results.append({"content": content, "is_error": is_error}) results.reverse() return results def _assistant_content_and_tool_calls(response_obj: Any) -> tuple: """Return (assistant_text, tool_calls_list) extracted from a ModelResponse-ish object.""" if response_obj is None: return None, [] try: choices = getattr(response_obj, "choices", None) or response_obj.get("choices") except Exception: return None, [] if not choices: return None, [] msg = choices[0] msg = getattr(msg, "message", None) or ( msg.get("message") if isinstance(msg, dict) else None ) if msg is None: return None, [] content = getattr(msg, "content", None) if content is None and isinstance(msg, dict): content = msg.get("content") raw_tool_calls = getattr(msg, "tool_calls", None) if raw_tool_calls is None and isinstance(msg, dict): raw_tool_calls = msg.get("tool_calls") tool_calls: List[Dict[str, Any]] = [] for tc in raw_tool_calls or []: if isinstance(tc, dict): tool_calls.append(tc) else: try: tool_calls.append(tc.model_dump()) except Exception: tool_calls.append({"name": getattr(tc, "name", ""), "arguments": ""}) return content, tool_calls class AdaptiveRouterPostCallHook(CustomLogger): """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" def __init__(self, adaptive_router: AdaptiveRouter) -> None: self.adaptive_router = adaptive_router async def async_post_call_response_headers_hook( self, data: Dict[str, Any], user_api_key_dict: Any, response: Any, request_headers: Optional[Dict[str, str]] = None, litellm_call_info: Optional[Dict[str, Any]] = None, ) -> Optional[Dict[str, str]]: """ Surface the chosen logical model as the `x-litellm-adaptive-router-model` response header for both streaming and non-streaming responses. `async_post_call_success_hook` fires after the stream is fully consumed, so writing to `_hidden_params["additional_headers"]` there is too late for streaming — the StreamingResponse headers are already frozen. This hook is called during header construction (before StreamingResponse is built), so the header is included for both paths. """ metadata = data.get("metadata") or {} chosen = ( metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) if isinstance(metadata, dict) else None ) if not chosen: return None return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._record(kwargs, response_obj, response_status=200) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): status = kwargs.get("response_status") if status is None: exc = kwargs.get("exception") status = getattr(exc, "status_code", 500) if exc is not None else 500 await self._record(kwargs, response_obj, response_status=int(status)) async def _record( self, kwargs: Dict[str, Any], response_obj: Any, response_status: int, ) -> None: try: messages = kwargs.get("messages") or [] if len(messages) < SIGNAL_GATE_MIN_MESSAGES: # Too few turns for any signal to be meaningful — skip. return session_key = _resolve_session_key(kwargs) if not session_key: return # The bandit cells are keyed by the *logical* model name from # `available_models` (e.g. "smart"/"fast"). `kwargs["model"]` at # post-call time is the physical upstream model # (e.g. "anthropic/claude-opus-4-7"), so it cannot be used directly. # The pre-routing hook stashes the logical pick under this key. litellm_params = kwargs.get("litellm_params") or {} metadata = litellm_params.get("metadata") or {} current_model = ( metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) if isinstance(metadata, dict) else None ) if not current_model: return if not self.adaptive_router.claim_or_check_owner( session_key, current_model ): # A different model owns this conversation — skip attribution. return user_text = _last_user_content(messages) assistant_text, tool_calls = _assistant_content_and_tool_calls(response_obj) tool_results = _recent_tool_results(messages) request_type = classify_prompt(user_text or "") turn = Turn( user_content=user_text, assistant_content=( assistant_text if isinstance(assistant_text, str) else None ), tool_calls=tool_calls, tool_results=tool_results, response_status=response_status, ) await self.adaptive_router.record_turn( session_id=session_key, model_name=current_model, request_type=request_type, turn=turn, ) except Exception as e: verbose_router_logger.exception( "AdaptiveRouterPostCallHook: failed to record turn: %s", e )