From d5faa402645f6a2ae041325fc870e01c913f36d0 Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Thu, 15 Oct 2026 16:37:57 -0700 Subject: [PATCH 1/4] python: fix citation issue --- .../_responses.py | 97 +++++++- .../foundry_hosting/tests/test_responses.py | 207 ++++++++++++++++++ 2 files changed, 302 insertions(+), 2 deletions(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 09692723734..da60959f543 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -50,6 +50,9 @@ add_usage_details, get_namespaced_value, ) +from agent_framework import ( + Annotation as FrameworkAnnotation, +) from agent_framework._mcp import _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY # pyright: ignore[reportPrivateUsage] from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentFrameworkException @@ -63,12 +66,17 @@ from azure.ai.agentserver.responses._id_generator import IdGenerator from azure.ai.agentserver.responses.aio import ResponseEventStream from azure.ai.agentserver.responses.hosting import ResponsesAgentServerHost +from azure.ai.agentserver.responses.models import ( + Annotation as ResponsesAnnotation, +) from azure.ai.agentserver.responses.models import ( ComputerAction, ComputerCallSafetyCheckParam, ComputerScreenshotImage, ContainerFileCitationBody, CreateResponse, + FileCitationBody, + FilePath, FunctionShellAction, FunctionShellCallOutputContent, FunctionShellCallOutputExitOutcome, @@ -88,6 +96,7 @@ ResponseUsage, ResponseUsageInputTokensDetails, ResponseUsageOutputTokensDetails, + UrlCitationBody, ) from azure.ai.agentserver.responses.streaming._builders import ( OutputItemBuilder, @@ -272,6 +281,80 @@ def _container_file_citations_from_function_result(content: Content) -> list[_Co return citations +def _annotations_to_output_text( + annotations: Sequence[FrameworkAnnotation] | None, + *, + text_offset: int = 0, +) -> list[ResponsesAnnotation]: + """Convert framework citations into Responses output-text annotations.""" + if not annotations: + return [] + + output_annotations: list[ResponsesAnnotation] = [] + for annotation in annotations: + if annotation.get("type") != "citation": + continue + + properties = annotation.get("additional_properties") or {} + regions = annotation.get("annotated_regions") or [] + file_id = annotation.get("file_id") + url = annotation.get("url") + title = annotation.get("title") + container_id = properties.get("container_id") + index = properties.get("index") + + if isinstance(container_id, str) and container_id and isinstance(file_id, str) and file_id: + for region in regions: + start_index = region.get("start_index") + end_index = region.get("end_index") + if not isinstance(start_index, int) or not isinstance(end_index, int): + continue + output_annotations.append( + ContainerFileCitationBody( + type="container_file_citation", + container_id=container_id, + file_id=file_id, + filename=url or "", + start_index=start_index + text_offset, + end_index=end_index + text_offset, + ) + ) + elif isinstance(url, str) and url and not file_id: + for region in regions: + start_index = region.get("start_index") + end_index = region.get("end_index") + if not isinstance(start_index, int) or not isinstance(end_index, int): + continue + output_annotations.append( + UrlCitationBody( + type="url_citation", + url=url, + title=title or "", + start_index=start_index + text_offset, + end_index=end_index + text_offset, + ) + ) + elif isinstance(file_id, str) and file_id and isinstance(url, str) and url and isinstance(index, int): + output_annotations.append( + FileCitationBody( + type="file_citation", + file_id=file_id, + filename=url, + index=index, + ) + ) + elif isinstance(file_id, str) and file_id and isinstance(index, int): + output_annotations.append( + FilePath( + type="file_path", + file_id=file_id, + index=index, + ) + ) + + return output_annotations + + def _is_refusal_text_content(content: Content) -> bool: return ( content.type == "text" @@ -2423,7 +2506,8 @@ def __init__( self._outstanding_computer_calls: set[str] = set() self._oauth_consent_requests: set[tuple[str, str]] = set() self._pending_container_file_citations: dict[str, _ContainerFileCitation] = {} - self._message_text_annotations: dict[int, list[ContainerFileCitationBody]] = {} + self._text_annotations: list[ResponsesAnnotation] = [] + self._message_text_annotations: dict[int, list[ResponsesAnnotation]] = {} # Set when an agent update reports the model stopped early (content filter, token # limit); the response then ends as ``incomplete`` instead of ``completed`` so callers # can tell a cut-short turn from a successful one. Mirrored into the stream's @@ -2595,7 +2679,11 @@ async def handle( for event in self._ensure_message_content("text", message_id): yield event self._active_message_id = message_id + text_offset = len("".join(self._accumulated)) if content.text and content.merge_mode != "replace" else 0 delta = self._consume_text_update(content.text, content.merge_mode) + for annotation in _annotations_to_output_text(content.annotations, text_offset=text_offset): + if annotation not in self._text_annotations: + self._text_annotations.append(annotation) if delta and self._text_content is not None: yield self._text_content.emit_delta(delta) @@ -2937,6 +3025,7 @@ def _open_message_content( yield self._refusal_content.emit_added() else: self._text_content = self._message_item.add_text_content() + self._text_annotations.clear() yield self._text_content.emit_added() def _open_reasoning(self, content: Content) -> Generator[ResponseStreamEvent]: @@ -3046,7 +3135,10 @@ def _close(self) -> Generator[ResponseStreamEvent]: def _close_message_content(self) -> Generator[ResponseStreamEvent]: accumulated = "".join(self._accumulated) if self._active_type == "text" and self._text_content is not None: - annotations = self._container_file_annotations_for_text(accumulated) + annotations = list(self._text_annotations) + for annotation in self._container_file_annotations_for_text(accumulated): + if annotation not in annotations: + annotations.append(annotation) yield self._text_content.emit_text_done(accumulated) for annotation in annotations: yield self._text_content.emit_annotation_added(annotation) @@ -3058,6 +3150,7 @@ def _close_message_content(self) -> Generator[ResponseStreamEvent]: self._message_text_annotations[self._text_content.content_index] = annotations yield content_done self._text_content = None + self._text_annotations.clear() elif self._active_type == "refusal" and self._refusal_content is not None: yield self._refusal_content.emit_refusal_done(accumulated) yield self._refusal_content.emit_done() diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index dc9b643bd33..2db5bfcd082 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -34,6 +34,7 @@ AgentResponse, AgentResponseUpdate, AgentSession, + Annotation, BaseChatClient, ChatMiddlewareLayer, ChatResponse, @@ -94,6 +95,7 @@ CONSENT_ERROR_CODE, ConsentError, _agent_response_updates, # pyright: ignore[reportPrivateUsage] + _annotations_to_output_text, # pyright: ignore[reportPrivateUsage] _await_before_signal, # pyright: ignore[reportPrivateUsage] _container_file_citations_from_function_result, # pyright: ignore[reportPrivateUsage] _is_allowed_oauth_consent_link, # pyright: ignore[reportPrivateUsage] @@ -241,6 +243,65 @@ def _expected_container_file_citation( } +def _url_citation( + *, + url: str = "https://www.python.org/downloads/release/python-31315/", + title: str = "Python 3.13.15", + start_index: int = 0, + end_index: int = 6, +) -> Annotation: + return Annotation( + type="citation", + url=url, + title=title, + annotated_regions=[{"type": "text_span", "start_index": start_index, "end_index": end_index}], + ) + + +def _expected_url_citation( + *, + url: str = "https://www.python.org/downloads/release/python-31315/", + title: str = "Python 3.13.15", + start_index: int = 0, + end_index: int = 6, +) -> dict[str, Any]: + return { + "type": "url_citation", + "url": url, + "title": title, + "start_index": start_index, + "end_index": end_index, + } + + +def _file_citation( + *, + file_id: str = "file_123", + filename: str = "brambleglass-kestrel-st731.txt", + index: int = 0, +) -> Annotation: + return Annotation( + type="citation", + file_id=file_id, + url=filename, + additional_properties={"index": index}, + ) + + +def _expected_file_citation( + *, + file_id: str = "file_123", + filename: str = "brambleglass-kestrel-st731.txt", + index: int = 0, +) -> dict[str, Any]: + return { + "type": "file_citation", + "file_id": file_id, + "filename": filename, + "index": index, + } + + @pytest.mark.parametrize( "raw_citations", [ @@ -274,6 +335,80 @@ def test_container_file_citations_ignore_recursive_json() -> None: assert _container_file_citations_from_function_result(function_result) == [] +def test_framework_annotations_convert_all_responses_citation_variants() -> None: + annotations = [ + Annotation( + type="citation", + title="Source", + url="https://example.com/source", + annotated_regions=[ + {"type": "text_span", "start_index": 0, "end_index": 6}, + {"type": "text_span", "start_index": 12, "end_index": 18}, + ], + ), + _file_citation(), + Annotation( + type="citation", + file_id="file_path_123", + additional_properties={"index": 7}, + ), + Annotation( + type="citation", + file_id="container_file_123", + url="report.csv", + annotated_regions=[{"type": "text_span", "start_index": 20, "end_index": 30}], + additional_properties={"container_id": "container_123"}, + ), + ] + + assert _annotations_to_output_text(annotations, text_offset=4) == [ + { + "type": "url_citation", + "url": "https://example.com/source", + "title": "Source", + "start_index": 4, + "end_index": 10, + }, + { + "type": "url_citation", + "url": "https://example.com/source", + "title": "Source", + "start_index": 16, + "end_index": 22, + }, + _expected_file_citation(), + { + "type": "file_path", + "file_id": "file_path_123", + "index": 7, + }, + { + "type": "container_file_citation", + "container_id": "container_123", + "file_id": "container_file_123", + "filename": "report.csv", + "start_index": 24, + "end_index": 34, + }, + ] + + +def test_framework_annotations_ignore_unsupported_or_incomplete_values() -> None: + annotations = [ + Annotation(type="citation", url="https://example.com/no-span"), + Annotation( + type="citation", + url="https://example.com/invalid-span", + annotated_regions=[{"type": "text_span", "start_index": 0}], + ), + Annotation(type="citation", file_id="file-without-index"), + Annotation(type="citation", file_id="file-without-index", url="document.txt"), + cast(Annotation, {"type": "other", "url": "https://example.com/unsupported"}), + ] + + assert _annotations_to_output_text(annotations) == [] + + # region Helpers @@ -3759,6 +3894,31 @@ async def test_function_result_container_file_citation_is_added_to_output_text(s assert preparing["annotations"] == [] assert final["annotations"] == [_expected_container_file_citation()] + async def test_framework_url_and_file_citations_are_added_to_output_text(self) -> None: + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate( + contents=[ + Content.from_text( + "Python sources", + annotations=[_url_citation(), _file_citation()], + ) + ], + role="assistant", + message_id="msg_citations", + ) + ] + ) + + resp = await _post(_make_server(agent), stream=False) + + assert resp.status_code == 200 + message = next(item for item in resp.json()["output"] if item["type"] == "message") + assert message["content"][0]["annotations"] == [ + _expected_url_citation(), + _expected_file_citation(), + ] + async def test_latest_pending_container_file_citation_replaces_same_filename(self) -> None: agent = _make_agent( stream_updates=[ @@ -4582,6 +4742,53 @@ async def test_container_file_citation_is_consistent_across_stream_events(self) ) assert final_message["content"][0]["annotations"] == [expected_annotation] + async def test_framework_url_citation_is_consistent_across_stream_events(self) -> None: + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate( + contents=[Content.from_text("See ")], + role="assistant", + message_id="msg_citation", + ), + AgentResponseUpdate( + contents=[Content.from_text("Python", annotations=[_url_citation()])], + role="assistant", + message_id="msg_citation", + ), + AgentResponseUpdate( + contents=[Content.from_text("", annotations=[_url_citation(start_index=4, end_index=10)])], + role="assistant", + message_id="msg_citation", + ), + AgentResponseUpdate( + contents=[ + Content.from_text( + "See Python", + annotations=[_url_citation(start_index=4, end_index=10)], + merge_mode="replace", + ) + ], + role="assistant", + message_id="msg_citation", + ), + ] + ) + + resp = await _post(_make_server(agent), stream=True) + + assert resp.status_code == 200 + events = _parse_sse_events(resp.text) + expected_annotation = _expected_url_citation(start_index=4, end_index=10) + annotation_events = [event for event in events if event["event"] == "response.output_text.annotation.added"] + assert [event["data"]["annotation"] for event in annotation_events] == [expected_annotation] + + content_done = next(event for event in events if event["event"] == "response.content_part.done") + assert content_done["data"]["part"]["annotations"] == [expected_annotation] + + completed = next(event["data"]["response"] for event in events if event["event"] == "response.completed") + final_message = next(item for item in completed["output"] if item["type"] == "message") + assert final_message["content"][0]["annotations"] == [expected_annotation] + async def test_pending_container_file_citation_survives_tracker_recovery(self) -> None: stream = ResponseEventStream(response_id="resp_container_file") stream.emit_created() From b63782dbf81cb9c4568cc06fc40c5eb86fcaa03a Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Fri, 16 Oct 2026 15:28:24 -0700 Subject: [PATCH 2/4] update tests and citations for anthropic --- .../agent_framework_anthropic/_chat_client.py | 32 ---------- .../anthropic/tests/test_anthropic_client.py | 47 +++++++++++++-- .../_responses.py | 41 ++++++++----- .../foundry_hosting/tests/test_responses.py | 58 ++++++++++++++++--- 4 files changed, 118 insertions(+), 60 deletions(-) diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index dd894913609..f9b96b9abf5 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -1737,40 +1737,16 @@ def _parse_citations_from_anthropic( cit["snippet"] = citation.cited_text if citation.file_id: cit["file_id"] = citation.file_id - cit.setdefault("annotated_regions", []) - cit["annotated_regions"].append( # type: ignore[attr-defined] - TextSpanRegion( - type="text_span", - start_index=citation.start_char_index, - end_index=citation.end_char_index, - ) - ) case "page_location": cit["title"] = citation.document_title cit["snippet"] = citation.cited_text if citation.file_id: cit["file_id"] = citation.file_id - cit.setdefault("annotated_regions", []) - cit["annotated_regions"].append( # type: ignore[attr-defined] - TextSpanRegion( - type="text_span", - start_index=citation.start_page_number, - end_index=citation.end_page_number, - ) - ) case "content_block_location": cit["title"] = citation.document_title cit["snippet"] = citation.cited_text if citation.file_id: cit["file_id"] = citation.file_id - cit.setdefault("annotated_regions", []) - cit["annotated_regions"].append( # type: ignore[attr-defined] - TextSpanRegion( - type="text_span", - start_index=citation.start_block_index, - end_index=citation.end_block_index, - ) - ) case "web_search_result_location": cit["title"] = citation.title cit["snippet"] = citation.cited_text @@ -1779,14 +1755,6 @@ def _parse_citations_from_anthropic( cit["title"] = citation.title cit["snippet"] = citation.cited_text cit["url"] = citation.source - cit.setdefault("annotated_regions", []) - cit["annotated_regions"].append( # type: ignore[attr-defined] - TextSpanRegion( - type="text_span", - start_index=citation.start_block_index, - end_index=citation.end_block_index, - ) - ) case _: logger.debug(f"Unknown citation type encountered: {citation.type}") annotations.append(cit) diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index 7822da91600..a5badfae3c4 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -4201,7 +4201,14 @@ def test_parse_citations_char_location(mock_anthropic_client: MagicMock) -> None result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert len(result) > 0 + assert result == [ + { + "type": "citation", + "title": "Source Title", + "snippet": "Citation snippet", + "raw_representation": mock_citation, + } + ] def test_parse_citations_page_location(mock_anthropic_client: MagicMock) -> None: @@ -4225,7 +4232,14 @@ def test_parse_citations_page_location(mock_anthropic_client: MagicMock) -> None result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert len(result) > 0 + assert result == [ + { + "type": "citation", + "title": "Document Title", + "snippet": "Cited text from page", + "raw_representation": mock_citation, + } + ] def test_parse_citations_content_block_location( @@ -4251,7 +4265,14 @@ def test_parse_citations_content_block_location( result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert len(result) > 0 + assert result == [ + { + "type": "citation", + "title": "Document Title", + "snippet": "Cited text from content blocks", + "raw_representation": mock_citation, + } + ] def test_parse_citations_web_search_location(mock_anthropic_client: MagicMock) -> None: @@ -4274,7 +4295,15 @@ def test_parse_citations_web_search_location(mock_anthropic_client: MagicMock) - result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert len(result) > 0 + assert result == [ + { + "type": "citation", + "title": "Search Result", + "snippet": "Cited text from search", + "url": "https://example.com", + "raw_representation": mock_citation, + } + ] def test_parse_citations_search_result_location( @@ -4301,7 +4330,15 @@ def test_parse_citations_search_result_location( result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert len(result) > 0 + assert result == [ + { + "type": "citation", + "title": "Search Result", + "snippet": "Cited text", + "url": "https://source.com", + "raw_representation": mock_citation, + } + ] @pytest.mark.flaky diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index da60959f543..b02c322e753 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -283,23 +283,35 @@ def _container_file_citations_from_function_result(content: Content) -> list[_Co def _annotations_to_output_text( annotations: Sequence[FrameworkAnnotation] | None, - *, - text_offset: int = 0, ) -> list[ResponsesAnnotation]: """Convert framework citations into Responses output-text annotations.""" - if not annotations: + if not annotations or isinstance(annotations, (str, bytes)): return [] output_annotations: list[ResponsesAnnotation] = [] for annotation in annotations: - if annotation.get("type") != "citation": + if not isinstance(annotation, Mapping): + continue + annotation_mapping = cast("Mapping[str, Any]", annotation) + if annotation_mapping.get("type") != "citation": continue - properties = annotation.get("additional_properties") or {} - regions = annotation.get("annotated_regions") or [] - file_id = annotation.get("file_id") - url = annotation.get("url") - title = annotation.get("title") + raw_properties = annotation_mapping.get("additional_properties") + properties: Mapping[str, Any] = ( + cast("Mapping[str, Any]", raw_properties) if isinstance(raw_properties, Mapping) else {} + ) + raw_regions = annotation_mapping.get("annotated_regions") + regions: list[Mapping[str, Any]] = [] + if isinstance(raw_regions, Sequence) and not isinstance(raw_regions, (str, bytes)): + for raw_region in cast("Sequence[Any]", raw_regions): + if not isinstance(raw_region, Mapping): + continue + region = cast("Mapping[str, Any]", raw_region) + if region.get("type") == "text_span": + regions.append(region) + file_id = annotation_mapping.get("file_id") + url = annotation_mapping.get("url") + title = annotation_mapping.get("title") container_id = properties.get("container_id") index = properties.get("index") @@ -315,8 +327,8 @@ def _annotations_to_output_text( container_id=container_id, file_id=file_id, filename=url or "", - start_index=start_index + text_offset, - end_index=end_index + text_offset, + start_index=start_index, + end_index=end_index, ) ) elif isinstance(url, str) and url and not file_id: @@ -330,8 +342,8 @@ def _annotations_to_output_text( type="url_citation", url=url, title=title or "", - start_index=start_index + text_offset, - end_index=end_index + text_offset, + start_index=start_index, + end_index=end_index, ) ) elif isinstance(file_id, str) and file_id and isinstance(url, str) and url and isinstance(index, int): @@ -2679,9 +2691,8 @@ async def handle( for event in self._ensure_message_content("text", message_id): yield event self._active_message_id = message_id - text_offset = len("".join(self._accumulated)) if content.text and content.merge_mode != "replace" else 0 delta = self._consume_text_update(content.text, content.merge_mode) - for annotation in _annotations_to_output_text(content.annotations, text_offset=text_offset): + for annotation in _annotations_to_output_text(content.annotations): if annotation not in self._text_annotations: self._text_annotations.append(annotation) if delta and self._text_content is not None: diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 2db5bfcd082..fa143f45d76 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -361,20 +361,20 @@ def test_framework_annotations_convert_all_responses_citation_variants() -> None ), ] - assert _annotations_to_output_text(annotations, text_offset=4) == [ + assert _annotations_to_output_text(annotations) == [ { "type": "url_citation", "url": "https://example.com/source", "title": "Source", - "start_index": 4, - "end_index": 10, + "start_index": 0, + "end_index": 6, }, { "type": "url_citation", "url": "https://example.com/source", "title": "Source", - "start_index": 16, - "end_index": 22, + "start_index": 12, + "end_index": 18, }, _expected_file_citation(), { @@ -387,12 +387,49 @@ def test_framework_annotations_convert_all_responses_citation_variants() -> None "container_id": "container_123", "file_id": "container_file_123", "filename": "report.csv", - "start_index": 24, - "end_index": 34, + "start_index": 20, + "end_index": 30, }, ] +def test_framework_annotations_ignore_malformed_runtime_shapes() -> None: + annotations = cast( + Any, + [ + None, + "not-an-annotation", + { + "type": "citation", + "url": "https://example.com/mixed", + "title": "Mixed", + "additional_properties": "not-a-mapping", + "annotated_regions": [ + None, + "not-a-region", + {"type": "source_block", "start_index": 0, "end_index": 5}, + {"type": "text_span", "start_index": 2, "end_index": 7}, + ], + }, + { + "type": "citation", + "url": "https://example.com/string-regions", + "annotated_regions": "not-a-sequence-of-regions", + }, + ], + ) + + assert _annotations_to_output_text(annotations) == [ + _expected_url_citation( + url="https://example.com/mixed", + title="Mixed", + start_index=2, + end_index=7, + ) + ] + assert _annotations_to_output_text(cast(Any, "not-an-annotation-sequence")) == [] + + def test_framework_annotations_ignore_unsupported_or_incomplete_values() -> None: annotations = [ Annotation(type="citation", url="https://example.com/no-span"), @@ -4751,7 +4788,12 @@ async def test_framework_url_citation_is_consistent_across_stream_events(self) - message_id="msg_citation", ), AgentResponseUpdate( - contents=[Content.from_text("Python", annotations=[_url_citation()])], + contents=[ + Content.from_text( + "Python", + annotations=[_url_citation(start_index=4, end_index=10)], + ) + ], role="assistant", message_id="msg_citation", ), From 7f7f0c33aecf299e39c7ff41363b3648fc7c3f7e Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Wed, 21 Oct 2026 13:18:51 -0700 Subject: [PATCH 3/4] add Sequence check --- .../agent_framework_foundry_hosting/_responses.py | 2 +- .../packages/foundry_hosting/tests/test_responses.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index b02c322e753..c0aa5b46e57 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -285,7 +285,7 @@ def _annotations_to_output_text( annotations: Sequence[FrameworkAnnotation] | None, ) -> list[ResponsesAnnotation]: """Convert framework citations into Responses output-text annotations.""" - if not annotations or isinstance(annotations, (str, bytes)): + if not annotations or not isinstance(annotations, Sequence) or isinstance(annotations, (str, bytes)): return [] output_annotations: list[ResponsesAnnotation] = [] diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index fa143f45d76..112c88eff68 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -430,6 +430,18 @@ def test_framework_annotations_ignore_malformed_runtime_shapes() -> None: assert _annotations_to_output_text(cast(Any, "not-an-annotation-sequence")) == [] +@pytest.mark.parametrize( + "annotations", + [ + pytest.param(42, id="integer"), + pytest.param(True, id="boolean"), + pytest.param(object(), id="object"), + ], +) +def test_framework_annotations_ignore_non_sequence_inputs(annotations: Any) -> None: + assert _annotations_to_output_text(annotations) == [] + + def test_framework_annotations_ignore_unsupported_or_incomplete_values() -> None: annotations = [ Annotation(type="citation", url="https://example.com/no-span"), From 63f108affd939f5f1edaba52758a65c3b0a1d537 Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Fri, 23 Oct 2026 16:30:53 -0700 Subject: [PATCH 4/4] filter for source citations and preserve completed response citations --- .../agent_framework_anthropic/_chat_client.py | 42 ++++++++ .../anthropic/tests/test_anthropic_client.py | 91 ++++++++-------- python/packages/foundry_hosting/README.md | 12 +++ .../_responses.py | 35 ++++-- .../foundry_hosting/tests/test_responses.py | 102 +++++++++++++++++- 5 files changed, 231 insertions(+), 51 deletions(-) diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index f9b96b9abf5..558e4c39cb6 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -1725,6 +1725,14 @@ def _parse_contents_from_anthropic( def _parse_citations_from_anthropic( self, content_block: BetaContentBlock | BetaRawContentBlockDelta | BetaTextBlock ) -> list[Annotation] | None: + """Parse Anthropic citations into annotations. + + Note: + Document and search-result citations retain their source coordinates in ``annotated_regions``. + Their ``additional_properties["annotated_regions_coordinate_space"]`` is ``"source"`` to distinguish + these locations from generated-response character offsets. The Responses host does not emit these + regions as output-text spans. + """ content_blocks = getattr(content_block, "citations", None) if not content_blocks: return None @@ -1737,16 +1745,40 @@ def _parse_citations_from_anthropic( cit["snippet"] = citation.cited_text if citation.file_id: cit["file_id"] = citation.file_id + cit.setdefault("annotated_regions", []) + cit["annotated_regions"].append( # type: ignore[attr-defined] + TextSpanRegion( + type="text_span", + start_index=citation.start_char_index, + end_index=citation.end_char_index, + ) + ) case "page_location": cit["title"] = citation.document_title cit["snippet"] = citation.cited_text if citation.file_id: cit["file_id"] = citation.file_id + cit.setdefault("annotated_regions", []) + cit["annotated_regions"].append( # type: ignore[attr-defined] + TextSpanRegion( + type="text_span", + start_index=citation.start_page_number, + end_index=citation.end_page_number, + ) + ) case "content_block_location": cit["title"] = citation.document_title cit["snippet"] = citation.cited_text if citation.file_id: cit["file_id"] = citation.file_id + cit.setdefault("annotated_regions", []) + cit["annotated_regions"].append( # type: ignore[attr-defined] + TextSpanRegion( + type="text_span", + start_index=citation.start_block_index, + end_index=citation.end_block_index, + ) + ) case "web_search_result_location": cit["title"] = citation.title cit["snippet"] = citation.cited_text @@ -1755,8 +1787,18 @@ def _parse_citations_from_anthropic( cit["title"] = citation.title cit["snippet"] = citation.cited_text cit["url"] = citation.source + cit.setdefault("annotated_regions", []) + cit["annotated_regions"].append( # type: ignore[attr-defined] + TextSpanRegion( + type="text_span", + start_index=citation.start_block_index, + end_index=citation.end_block_index, + ) + ) case _: logger.debug(f"Unknown citation type encountered: {citation.type}") + if cit.get("annotated_regions"): + cit["additional_properties"] = {"annotated_regions_coordinate_space": "source"} annotations.append(cit) return annotations or None diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index a5badfae3c4..d14027f8b2d 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -4201,14 +4201,7 @@ def test_parse_citations_char_location(mock_anthropic_client: MagicMock) -> None result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert result == [ - { - "type": "citation", - "title": "Source Title", - "snippet": "Citation snippet", - "raw_representation": mock_citation, - } - ] + assert len(result) > 0 def test_parse_citations_page_location(mock_anthropic_client: MagicMock) -> None: @@ -4232,14 +4225,7 @@ def test_parse_citations_page_location(mock_anthropic_client: MagicMock) -> None result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert result == [ - { - "type": "citation", - "title": "Document Title", - "snippet": "Cited text from page", - "raw_representation": mock_citation, - } - ] + assert len(result) > 0 def test_parse_citations_content_block_location( @@ -4265,14 +4251,7 @@ def test_parse_citations_content_block_location( result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert result == [ - { - "type": "citation", - "title": "Document Title", - "snippet": "Cited text from content blocks", - "raw_representation": mock_citation, - } - ] + assert len(result) > 0 def test_parse_citations_web_search_location(mock_anthropic_client: MagicMock) -> None: @@ -4295,15 +4274,7 @@ def test_parse_citations_web_search_location(mock_anthropic_client: MagicMock) - result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert result == [ - { - "type": "citation", - "title": "Search Result", - "snippet": "Cited text from search", - "url": "https://example.com", - "raw_representation": mock_citation, - } - ] + assert len(result) > 0 def test_parse_citations_search_result_location( @@ -4330,15 +4301,51 @@ def test_parse_citations_search_result_location( result = client._parse_citations_from_anthropic(mock_block) assert result is not None - assert result == [ - { - "type": "citation", - "title": "Search Result", - "snippet": "Cited text", - "url": "https://source.com", - "raw_representation": mock_citation, - } - ] + assert len(result) > 0 + + +@pytest.mark.parametrize( + ("citation_type", "location_fields", "expected_span"), + [ + ("char_location", {"start_char_index": 2, "end_char_index": 12}, (2, 12)), + ("page_location", {"start_page_number": 1, "end_page_number": 3}, (1, 3)), + ("content_block_location", {"start_block_index": 0, "end_block_index": 2}, (0, 2)), + ("search_result_location", {"start_block_index": 3, "end_block_index": 5}, (3, 5)), + ("web_search_result_location", {}, None), + ], +) +def test_parse_citations_marks_source_coordinate_space( + mock_anthropic_client: MagicMock, + citation_type: str, + location_fields: dict[str, int], + expected_span: tuple[int, int] | None, +) -> None: + client = create_test_anthropic_client(mock_anthropic_client) + citation = MagicMock( + type=citation_type, + title="Source", + document_title="Source", + cited_text="Cited text", + file_id=None, + url="https://example.com", + source="https://example.com", + ) + citation.configure_mock(**location_fields) + block = MagicMock(type="text", text="Answer", citations=[citation]) + + result = client._parse_citations_from_anthropic(block) + + assert result is not None + assert len(result) == 1 + annotation = result[0] + if expected_span is None: + assert "annotated_regions" not in annotation + assert "annotated_regions_coordinate_space" not in annotation.get("additional_properties", {}) + else: + assert annotation["annotated_regions"] == [ + {"type": "text_span", "start_index": expected_span[0], "end_index": expected_span[1]} + ] + assert annotation["additional_properties"]["annotated_regions_coordinate_space"] == "source" @pytest.mark.flaky diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 3a352eba761..279fc2c3ee6 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -276,6 +276,18 @@ output item ID without changing the provider's `call_id`. Although screenshots are optional in shared `Content`, this Responses host requires them on computer results. The computer-use `Content` constructors and `ComputerSafetyCheck` type are experimental Agent Framework APIs. +## Responses citation annotations + +`ResponsesHostServer` serializes supported framework URL and file citations into `output_text.annotations`. +URL and container-file citation spans refer to character positions in the generated response text. +Annotations with `additional_properties["annotated_regions_coordinate_space"] == "source"` retain +source-document coordinates (for example, Anthropic character, page, or block locations); the host leaves +those annotations unchanged but does not emit their regions as output-text spans. Other supported citation +metadata, such as file citations with an explicit output index, is still converted. +Completed provider responses retain separate text content parts, keeping citation offsets relative to +their original part. Ordinary streaming deltas still merge within their existing text occurrence, +without shifting citation offsets. + ## State store ### Local persistence diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index c0aa5b46e57..a6f67174461 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -284,7 +284,11 @@ def _container_file_citations_from_function_result(content: Content) -> list[_Co def _annotations_to_output_text( annotations: Sequence[FrameworkAnnotation] | None, ) -> list[ResponsesAnnotation]: - """Convert framework citations into Responses output-text annotations.""" + """Convert framework citations into Responses output-text annotations. + + Regions with ``additional_properties["annotated_regions_coordinate_space"] == "source"`` + describe the source document, not output text, and are not emitted as output spans. + """ if not annotations or not isinstance(annotations, Sequence) or isinstance(annotations, (str, bytes)): return [] @@ -301,6 +305,8 @@ def _annotations_to_output_text( cast("Mapping[str, Any]", raw_properties) if isinstance(raw_properties, Mapping) else {} ) raw_regions = annotation_mapping.get("annotated_regions") + if properties.get("annotated_regions_coordinate_space") == "source": + raw_regions = None regions: list[Mapping[str, Any]] = [] if isinstance(raw_regions, Sequence) and not isinstance(raw_regions, (str, bytes)): for raw_region in cast("Sequence[Any]", raw_regions): @@ -1858,7 +1864,9 @@ async def _handle_inner_agent( for update in _agent_response_updates(response, context.response_id): if context.shutdown.is_set() or cancellation_signal.is_set(): break - async for event in tracker.handle_update(update, approval_storage=approval_storage): + async for event in tracker.handle_update( + update, approval_storage=approval_storage, complete_message=True + ): yield event if context.shutdown.is_set() or cancellation_signal.is_set(): continue @@ -2650,15 +2658,23 @@ async def handle_update( update: AgentResponseUpdate, *, approval_storage: FunctionApprovalStore | None = None, + complete_message: bool = False, ) -> AsyncGenerator[ResponseStreamEvent]: """Process one agent update: note its finish reason, then handle each of its contents. This is the single entry point for both the plain-agent and the workflow loops, so the - finish reason cannot be forgotten on one of them. + finish reason cannot be forgotten on one of them. Set ``complete_message`` when the + update contains a completed message's independent content parts rather than streaming + deltas; text and refusal parts then retain their original boundaries and citation offsets. """ self.record_finish_reason(update.finish_reason) for content in update.contents: - async for event in self.handle(content, message_id=update.message_id, approval_storage=approval_storage): + async for event in self.handle( + content, + message_id=update.message_id, + approval_storage=approval_storage, + start_new_content_part=complete_message, + ): yield event async def handle( @@ -2667,6 +2683,7 @@ async def handle( message_id: str | None = None, *, approval_storage: FunctionApprovalStore | None = None, + start_new_content_part: bool = False, ) -> AsyncGenerator[ResponseStreamEvent]: """Process a content item, yielding its events. @@ -2678,9 +2695,11 @@ async def handle( merged into it. approval_storage: Used for content types that fall back to one-shot emission (anything not recognized as a streaming delta type) to save/load approval requests. + start_new_content_part: Whether text and refusal content start a new independent + content part instead of continuing the current streaming occurrence. """ if _is_refusal_text_content(content) and content.text is not None: - for event in self._ensure_message_content("refusal", message_id): + for event in self._ensure_message_content("refusal", message_id, start_new_part=start_new_content_part): yield event self._active_message_id = message_id delta = self._consume_text_update(content.text, content.merge_mode) @@ -2688,7 +2707,7 @@ async def handle( yield self._refusal_content.emit_delta(delta) elif content.type == "text" and content.text is not None: - for event in self._ensure_message_content("text", message_id): + for event in self._ensure_message_content("text", message_id, start_new_part=start_new_content_part): yield event self._active_message_id = message_id delta = self._consume_text_update(content.text, content.merge_mode) @@ -3004,11 +3023,13 @@ def _ensure_message_content( self, content_type: Literal["text", "refusal"], message_id: str | None, + *, + start_new_part: bool = False, ) -> Generator[ResponseStreamEvent]: message_changed = ( message_id is not None and self._active_message_id is not None and message_id != self._active_message_id ) - if self._active_type == content_type and not message_changed: + if self._active_type == content_type and not message_changed and not start_new_part: return if self._message_item is not None and self._active_type in {"text", "refusal"} and not message_changed: yield from self._close_message_content() diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 112c88eff68..6f164edf958 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -393,6 +393,37 @@ def test_framework_annotations_convert_all_responses_citation_variants() -> None ] +@pytest.mark.parametrize( + ("coordinate_space", "expected"), + [ + pytest.param(None, [_expected_url_citation()], id="unmarked"), + pytest.param("response", [_expected_url_citation()], id="response"), + pytest.param("source", [], id="source"), + ], +) +def test_framework_annotations_respect_coordinate_space_without_raw_representation( + coordinate_space: str | None, expected: list[dict[str, Any]] +) -> None: + annotation = _url_citation() + if coordinate_space is not None: + annotation["additional_properties"] = {"annotated_regions_coordinate_space": coordinate_space} + annotations = json.loads(json.dumps([annotation])) + original_annotations = copy.deepcopy(annotations) + + assert _annotations_to_output_text(annotations) == expected + assert annotations == original_annotations + + +def test_framework_annotations_preserve_file_citations_with_source_regions() -> None: + annotation = _file_citation() + annotation["annotated_regions"] = [{"type": "text_span", "start_index": 1, "end_index": 3}] + annotation["additional_properties"]["annotated_regions_coordinate_space"] = "source" + original_annotation = copy.deepcopy(annotation) + + assert _annotations_to_output_text([annotation]) == [_expected_file_citation()] + assert annotation == original_annotation + + def test_framework_annotations_ignore_malformed_runtime_shapes() -> None: annotations = cast( Any, @@ -2124,6 +2155,69 @@ def test_provider_background_forwards_finish_reason_without_output(self) -> None assert updates[0].contents == [] assert updates[0].finish_reason == "content_filter" + @pytest.mark.parametrize("annotate_prefix", [False, True]) + async def test_provider_background_preserves_completed_text_part_citation_offsets( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, annotate_prefix: bool + ) -> None: + response = AgentResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_text( + "See ", annotations=[_url_citation(end_index=3)] if annotate_prefix else None + ), + Content.from_text("Python", annotations=[_url_citation()]), + ], + ) + ], + finish_reason="stop", + ) + original_response = copy.deepcopy(response.to_dict()) + agent = Agent(client=_ServiceStorageRecordingClient()) + + async def run(messages: Any = None, *, session: AgentSession, **kwargs: Any) -> AgentResponse: + del messages, kwargs + session.service_session_id = "private-service-conversation" + return response + + monkeypatch.setattr(agent, "run", MagicMock(side_effect=run)) + server = _make_server( + agent, + session_store=SessionStore(), + history_source="service", + background_source="provider", + options=ResponsesServerOptions(resilient_background=True), + response_store=FileResponseStore(storage_dir=tmp_path), + ) + request = CreateResponse(input="hello", store=True, background=True) + context = ResponseContext(response_id="outer-response", mode_flags=MagicMock()) + + with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])): + events = [ + event + async for event in server._handle_response(request, context, asyncio.Event()) + if isinstance(event, Mapping) + ] + + completed = next(event["response"] for event in events if event["type"] == "response.completed") + message = next(item for item in completed["output"] if item["type"] == "message") + parts = [part for part in message["content"] if part["type"] == "output_text"] + assert len(parts) == len(message["content"]) + assert [part["text"] for part in parts] == ["See ", "Python"] + assert parts[0]["annotations"] == ([_expected_url_citation(end_index=3)] if annotate_prefix else []) + assert parts[1]["annotations"] == [_expected_url_citation()] + annotation_events = [event for event in events if event["type"] == "response.output_text.annotation.added"] + assert [event["content_index"] for event in annotation_events] == ([0, 1] if annotate_prefix else [1]) + assert annotation_events[-1]["annotation"] == _expected_url_citation() + content_done = next( + event for event in events if event["type"] == "response.content_part.done" and event["content_index"] == 1 + ) + assert content_done["part"] == parts[1] + output_done = next(event for event in events if event["type"] == "response.output_item.done") + assert output_done["item"] == message + assert response.to_dict() == original_response + async def test_provider_background_keeps_private_token_under_outer_response( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -3944,13 +4038,15 @@ async def test_function_result_container_file_citation_is_added_to_output_text(s assert final["annotations"] == [_expected_container_file_citation()] async def test_framework_url_and_file_citations_are_added_to_output_text(self) -> None: + source_citation = _url_citation(start_index=0, end_index=1) + source_citation["additional_properties"] = {"annotated_regions_coordinate_space": "source"} agent = _make_agent( stream_updates=[ AgentResponseUpdate( contents=[ Content.from_text( "Python sources", - annotations=[_url_citation(), _file_citation()], + annotations=[source_citation, _url_citation(), _file_citation()], ) ], role="assistant", @@ -4792,10 +4888,12 @@ async def test_container_file_citation_is_consistent_across_stream_events(self) assert final_message["content"][0]["annotations"] == [expected_annotation] async def test_framework_url_citation_is_consistent_across_stream_events(self) -> None: + source_citation = _url_citation(start_index=0, end_index=1) + source_citation["additional_properties"] = {"annotated_regions_coordinate_space": "source"} agent = _make_agent( stream_updates=[ AgentResponseUpdate( - contents=[Content.from_text("See ")], + contents=[Content.from_text("See ", annotations=[source_citation])], role="assistant", message_id="msg_citation", ),