Skip to content

Commit 435c66e

Browse files
giles17CopilotCopilotmoonbox3
authored
Python: Handle url_citation annotations in FoundryChatClient streaming responses (#5071)
* Fix url_citation annotations dropped in streaming (#5029) Add url_citation branch to the streaming annotation handler in _parse_chunk_from_openai, mirroring the existing non-streaming path. The handler creates an Annotation with type='citation', title, url, and annotated_regions (TextSpanRegion), wrapped in Content.from_text. Update test_streaming_annotation_added_with_unknown_type to use a truly unknown type, and add new tests for url_citation (with and without url). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address review feedback for #5029: Python: [Bug]: url_citation annotations silently dropped in Foundry streaming (SharePoint grounding citations lost) --------- Co-authored-by: Copilot <copilot@github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com>
1 parent 52d50be commit 435c66e

2 files changed

Lines changed: 108 additions & 3 deletions

File tree

python/packages/openai/agent_framework_openai/_chat_client.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2474,6 +2474,29 @@ def _get_ann_value(key: str) -> Any:
24742474
raw_representation=event,
24752475
)
24762476
)
2477+
elif ann_type == "url_citation":
2478+
ann_url = _get_ann_value("url")
2479+
if ann_url:
2480+
ann_start = _get_ann_value("start_index")
2481+
ann_end = _get_ann_value("end_index")
2482+
annotation_obj = Annotation(
2483+
type="citation",
2484+
title=_get_ann_value("title") or "",
2485+
url=str(ann_url),
2486+
additional_properties={"annotation_index": event.annotation_index},
2487+
raw_representation=annotation,
2488+
)
2489+
if ann_start is not None and ann_end is not None:
2490+
annotation_obj["annotated_regions"] = [
2491+
TextSpanRegion(
2492+
type="text_span",
2493+
start_index=ann_start,
2494+
end_index=ann_end,
2495+
)
2496+
]
2497+
contents.append(
2498+
Content.from_text(text="", annotations=[annotation_obj], raw_representation=event)
2499+
)
24772500
else:
24782501
logger.debug("Unparsed annotation type in streaming: %s", ann_type)
24792502
case "response.output_item.done":

python/packages/openai/tests/openai/test_openai_chat_client.py

Lines changed: 85 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2570,8 +2570,65 @@ def test_streaming_annotation_added_with_container_file_citation() -> None:
25702570
assert content.additional_properties.get("end_index") == 50
25712571

25722572

2573-
def test_streaming_annotation_added_with_unknown_type() -> None:
2574-
"""Test streaming annotation added event with unknown type is ignored."""
2573+
def test_streaming_annotation_added_with_url_citation() -> None:
2574+
"""Test streaming annotation added event with url_citation type produces citation annotation."""
2575+
client = OpenAIChatClient(model="test-model", api_key="test-key")
2576+
chat_options = ChatOptions()
2577+
function_call_ids: dict[int, tuple[str, str]] = {}
2578+
2579+
mock_event = MagicMock()
2580+
mock_event.type = "response.output_text.annotation.added"
2581+
mock_event.annotation_index = 0
2582+
mock_event.annotation = {
2583+
"type": "url_citation",
2584+
"url": "https://example.sharepoint.com/sites/my-site/doc.pdf",
2585+
"title": "doc.pdf",
2586+
"start_index": 100,
2587+
"end_index": 112,
2588+
}
2589+
2590+
response = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
2591+
2592+
assert len(response.contents) == 1
2593+
content = response.contents[0]
2594+
assert content.type == "text"
2595+
assert content.annotations is not None
2596+
assert len(content.annotations) == 1
2597+
annotation = content.annotations[0]
2598+
assert annotation["type"] == "citation"
2599+
assert annotation["title"] == "doc.pdf"
2600+
assert annotation["url"] == "https://example.sharepoint.com/sites/my-site/doc.pdf"
2601+
assert annotation["additional_properties"]["annotation_index"] == 0
2602+
assert annotation["raw_representation"] == mock_event.annotation
2603+
assert annotation["annotated_regions"] is not None
2604+
assert len(annotation["annotated_regions"]) == 1
2605+
region = annotation["annotated_regions"][0]
2606+
assert region["type"] == "text_span"
2607+
assert region["start_index"] == 100
2608+
assert region["end_index"] == 112
2609+
2610+
2611+
def test_streaming_annotation_added_with_url_citation_no_url() -> None:
2612+
"""Test streaming annotation added event with url_citation but missing url is ignored."""
2613+
client = OpenAIChatClient(model="test-model", api_key="test-key")
2614+
chat_options = ChatOptions()
2615+
function_call_ids: dict[int, tuple[str, str]] = {}
2616+
2617+
mock_event = MagicMock()
2618+
mock_event.type = "response.output_text.annotation.added"
2619+
mock_event.annotation_index = 0
2620+
mock_event.annotation = {
2621+
"type": "url_citation",
2622+
"title": "doc.pdf",
2623+
}
2624+
2625+
response = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
2626+
2627+
assert len(response.contents) == 0
2628+
2629+
2630+
def test_streaming_annotation_added_with_url_citation_no_indices() -> None:
2631+
"""Test streaming annotation with url_citation that has url but no start_index/end_index."""
25752632
client = OpenAIChatClient(model="test-model", api_key="test-key")
25762633
chat_options = ChatOptions()
25772634
function_call_ids: dict[int, tuple[str, str]] = {}
@@ -2582,11 +2639,36 @@ def test_streaming_annotation_added_with_unknown_type() -> None:
25822639
mock_event.annotation = {
25832640
"type": "url_citation",
25842641
"url": "https://example.com",
2642+
"title": "Example",
2643+
}
2644+
2645+
response = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
2646+
2647+
assert len(response.contents) == 1
2648+
annotation = response.contents[0].annotations[0]
2649+
assert annotation["type"] == "citation"
2650+
assert annotation["title"] == "Example"
2651+
assert annotation["url"] == "https://example.com"
2652+
assert annotation["additional_properties"]["annotation_index"] == 0
2653+
assert "annotated_regions" not in annotation
2654+
2655+
2656+
def test_streaming_annotation_added_with_unknown_type() -> None:
2657+
"""Test streaming annotation added event with unknown type is ignored."""
2658+
client = OpenAIChatClient(model="test-model", api_key="test-key")
2659+
chat_options = ChatOptions()
2660+
function_call_ids: dict[int, tuple[str, str]] = {}
2661+
2662+
mock_event = MagicMock()
2663+
mock_event.type = "response.output_text.annotation.added"
2664+
mock_event.annotation_index = 0
2665+
mock_event.annotation = {
2666+
"type": "some_future_annotation_type",
2667+
"data": "test",
25852668
}
25862669

25872670
response = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
25882671

2589-
# url_citation should not produce HostedFileContent
25902672
assert len(response.contents) == 0
25912673

25922674

0 commit comments

Comments
 (0)