fix: promote choice-level annotations to delta in OpenRouter streaming

OpenRouterChatCompletionStreamingHandler.chunk_parser() only mapped
reasoning_content from streaming chunks. Annotations (url_citation)
at the choice level were not promoted to delta, causing them to be
invisible to the streaming handler downstream.

Changes:
- Check for annotations in both delta and choice level of streaming
  chunks, promote choice-level annotations to delta using pop() to
  avoid duplicate emission on the wire
This commit is contained in:
bbarwik 2026-03-14 20:33:45 +00:00
parent 8abf2d8e34
commit eaada9ef21
2 changed files with 99 additions and 1 deletions

View file

@ -278,7 +278,16 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator):
new_choices = []
for choice in chunk["choices"]:
choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning")
delta = choice["delta"]
delta["reasoning_content"] = delta.get("reasoning")
# Promote annotations (url_citation) from choice level to
# delta so they flow through the streaming handler.
# Some providers put them in delta, others at choice level.
annotations = delta.get("annotations") or choice.pop("annotations", None)
if annotations is not None:
delta["annotations"] = annotations
new_choices.append(choice)
return ModelResponseStream(
id=chunk["id"],

View file

@ -84,6 +84,95 @@ class TestOpenRouterChatCompletionStreamingHandler:
assert "KeyError" in str(exc_info.value)
assert exc_info.value.status_code == 400
def test_chunk_parser_passes_through_annotations_in_delta(self):
"""Annotations in delta (url_citation) must not be silently dropped."""
handler = OpenRouterChatCompletionStreamingHandler(
streaming_response=None, sync_stream=True
)
annotations = [
{
"type": "url_citation",
"url_citation": {
"url": "https://example.com",
"title": "Example",
"start_index": 0,
"end_index": 10,
},
}
]
chunk = {
"id": "test_id",
"created": 1234567890,
"model": "perplexity/sonar-pro",
"choices": [
{
"delta": {
"content": "Hello",
"reasoning": None,
"annotations": annotations,
},
"index": 0,
}
],
}
result = handler.chunk_parser(chunk)
assert result.choices[0]["delta"]["annotations"] == annotations
def test_chunk_parser_passes_through_annotations_on_choice(self):
"""Some providers put annotations at the choice level, not inside delta."""
handler = OpenRouterChatCompletionStreamingHandler(
streaming_response=None, sync_stream=True
)
annotations = [
{
"type": "url_citation",
"url_citation": {
"url": "https://example.com/alt",
"title": "Alt Source",
"start_index": 5,
"end_index": 15,
},
}
]
chunk = {
"id": "test_id",
"created": 1234567890,
"model": "openai/gpt-5.1",
"choices": [
{
"delta": {"content": "World", "reasoning": None},
"annotations": annotations,
"index": 0,
}
],
}
result = handler.chunk_parser(chunk)
assert result.choices[0]["delta"]["annotations"] == annotations
# Choice-level annotations must be removed after promotion to delta
assert result.choices[0].get("annotations") is None
def test_chunk_parser_no_annotations_no_regression(self):
"""When no annotations are present, delta must not gain an annotations key."""
handler = OpenRouterChatCompletionStreamingHandler(
streaming_response=None, sync_stream=True
)
chunk = {
"id": "test_id",
"created": 1234567890,
"model": "test_model",
"choices": [
{"delta": {"content": "hi", "reasoning": None}, "index": 0}
],
}
result = handler.chunk_parser(chunk)
assert result.choices[0]["delta"].get("annotations") is None
def test_openrouter_extra_body_transformation():
transformed_request = OpenrouterConfig().transform_request(