FEAT: Capture API response stop reason on MessagePiece metadata - #2340
FEAT: Capture API response stop reason on MessagePiece metadata#2340varunj-msft wants to merge 1 commit into
Conversation
Targets already recorded token usage but discarded why generation stopped, and dropped usage entirely on content-filtered responses. Capture the provider's stop reason alongside token usage: finish_reason for Chat Completions, Completions and LiteLLM; status plus incomplete_reason for the Responses API. A base no-op hook on OpenAITarget is called from _handle_content_filter_response, so a filtered response now records the tokens it consumed instead of returning bare metadata. These keys are reserved for the provider. construct_response_from_request merges the request's metadata into every response piece, so all of them are cleared before a capture writes back the subset its own API reports.
| choices = getattr(response, "choices", None) or [] | ||
| for index, piece in enumerate(pieces): | ||
| choice = choices[index] if index < len(choices) else None | ||
| set_response_metadata(pieces=[piece], values={"finish_reason": getattr(choice, "finish_reason", None)}) |
There was a problem hiding this comment.
NIT: add small inline comment explaining why this call must be per-piece.
| "TargetCapabilities", | ||
| "TargetIdentifier", | ||
| "TextDataTypeSerializer", | ||
| "TOKEN_USAGE_METADATA_PREFIX", |
There was a problem hiding this comment.
q: should RESERVED_RESPONSE_METADATA_KEYS also be exported here?
| for reserved_key in RESERVED_RESPONSE_METADATA_KEYS: | ||
| piece.prompt_metadata.pop(reserved_key, None) | ||
|
|
||
| for key, value in values.items(): | ||
| if isinstance(value, str) and value: | ||
| pieces[0].prompt_metadata[key] = value |
There was a problem hiding this comment.
pretty minor but there is an asymmetry in terms of what keys are iterated in the "clearing" loop and the "writing" loop. Although right now, all the callers only pass in keys in RESERVED_RESPONSE_METADATA_KEYS, there could be small chance for drift in the future. We could potentially just take in kwargs here that align with the reserved keys or filter for just the reserved keys and logging when one of the passed in values isn't in that set.
| the SDK raises on a content filter. | ||
| pieces (list[MessagePiece]): The constructed response pieces. | ||
| """ | ||
| capture_token_usage(pieces=pieces, response=response) |
There was a problem hiding this comment.
maybe add a one line comment explaining why we are using a function from chat_completions_response_parser here (i.e., Completions and Chat Completions APIs share the same usage schema)
Justin Song (jsong468)
left a comment
Justin Song (jsong468)
left a comment
There was a problem hiding this comment.
LGTM! I would just check that metadata is populated correctly when actually sending to these different targets if you haven't already!
2 participants
Description
Targets already recorded token usage but discarded why generation stopped, and dropped
usage entirely on content-filtered responses. This captures the provider's stop reason
alongside token usage on
MessagePiece.prompt_metadata:finish_reason— Chat Completions, Completions, LiteLLMstatus+incomplete_reason— Responses APIA base no-op hook
OpenAITarget._capture_response_metadatais called from_handle_content_filter_response, so a filtered response now records the tokens it consumed.These keys are reserved for the provider.
construct_response_from_requestmerges therequest's
prompt_metadatainto every response piece, so a caller-suppliedfinish_reasonwould otherwise be indistinguishable from the real one. All reserved keys are cleared from
every piece before each capture writes back what its own API reported;
token_usage_*iscleared by prefix for the same reason.
Two bugs fixed along the way:
OpenAICompletionTargetcaptured neither usage norfinish_reason.n>1, every piece got choice 0'sfinish_reason; each piece now gets its own.Not breaking:
_METADATA_PREFIX→TOKEN_USAGE_METADATA_PREFIXwas private with noexternal callers.
Tests and Documentation
tests/unit: 14,701 passed / 0 failed. Diff coverage 95% (gate is 90%).pre-commit run --all-filesclean, includingty.per-choice
finish_reasonwithn>1, content-filtered responses, "not reported" cases(missing/empty/non-string omitted rather than stored as zeros), and a SQLite round-trip.
finish_reason=stop; a truncated reasoning request →status=incomplete+incomplete_reason=max_output_tokens; and a forged caller-suppliedfinish_reasoncorrectly cleared.