diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index bd79ecda92a..813e53ce6e4 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -31,7 +31,9 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - Text annotations are surfaced as `CUSTOM` events with `name="annotations"` and `value={"messageId": ..., "annotations": [...]}`. Annotation-only updates announce a text message if needed without emitting empty text deltas. Provider `raw_representation` is omitted; `AGUIChatClient` restores the - annotations on text content. These live custom events are not replayed by `MESSAGES_SNAPSHOT`. + annotations on text content. Content-coordinate regions on text deltas are rebased to accumulated-message + offsets; source-coordinate regions are preserved unchanged. These live custom events are not replayed by + `MESSAGES_SNAPSHOT`. - Inbound custom event aliases are accepted: `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`. - Multimodal user inputs support both legacy (`text`, `binary`) and draft-style (`image`, `audio`, `video`, `document`) shapes. - Interrupted runs complete with `RUN_FINISHED.outcome.type == "interrupt"` and canonical `outcome.interrupts`; do not document or add new flows that depend on the legacy top-level `RUN_FINISHED.interrupt` field. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index a946b1461da..8afb2e8002e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -618,6 +618,8 @@ def _offset_annotation_text_spans(annotations: list[dict[str, Any]], offset: int if offset == 0: return for annotation in annotations: + if annotation.get("annotated_regions_coordinate_space", "content") != "content": + continue regions = annotation.get("annotated_regions") if not isinstance(regions, list): continue diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 6aa7aab5afc..e462e1dbbb3 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -758,6 +758,23 @@ def test_emit_text_rebases_delta_annotation_spans_and_deduplicates_full_replay() assert replay_events == [] +def test_emit_text_does_not_rebase_source_coordinate_annotation_spans() -> None: + """Source-document locations remain independent of accumulated response text.""" + flow = FlowState() + annotation = Annotation( + type="citation", + url="https://example.com/source", + annotated_regions=[{"type": "text_span", "start_index": 2, "end_index": 5}], + annotated_regions_coordinate_space="source", + ) + _emit_text(Content.from_text("Prefix "), flow) + + events = _emit_text(Content.from_text("target", annotations=[annotation]), flow) + + citation_event = next(event for event in events if isinstance(event, CustomEvent)) + assert citation_event.value["annotations"][0] == annotation + + def test_emit_text_preserves_annotation_only_spans() -> None: """Annotation-only updates already use complete-message span indices.""" flow = FlowState() diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index dd894913609..0f189666142 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -1725,6 +1725,13 @@ 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 ``annotated_regions_coordinate_space`` is ``"source"`` to distinguish these locations from + generated-content character offsets. Protocol adapters do not emit these regions as output-text spans. + """ content_blocks = getattr(content_block, "citations", None) if not content_blocks: return None @@ -1789,6 +1796,8 @@ def _parse_citations_from_anthropic( ) case _: logger.debug(f"Unknown citation type encountered: {citation.type}") + if cit.get("annotated_regions"): + cit["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 7822da91600..c6938ce7641 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -4304,6 +4304,50 @@ def test_parse_citations_search_result_location( 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 + else: + assert annotation["annotated_regions"] == [ + {"type": "text_span", "start_index": expected_span[0], "end_index": expected_span[1]} + ] + assert annotation["annotated_regions_coordinate_space"] == "source" + + @pytest.mark.flaky @pytest.mark.integration @skip_if_anthropic_integration_tests_disabled diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 895764eb93d..411cb728d08 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -82,6 +82,9 @@ agent_framework/ `additional_properties` namespaces and streamed mergeable content uses `merge_mode`; response updates can use the same field to replace all contents of one logical message; see [spec 005](../../../docs/specs/005-python-content-identity-and-protocol-metadata.md). +- **Citation coordinates** - `Annotation.annotated_regions_coordinate_space` distinguishes spans in generated + `Content.text` (`"content"`, also the default when omitted) from locations in cited source material (`"source"`). + Protocol adapters must not interpret source-coordinate regions as generated-text offsets. - **Computer use** - `ComputerSafetyCheck` and the `Content.from_computer_tool_call` / `from_computer_tool_result` constructors are experimental under `COMPUTER_USE`; the rest of `Content` retains its existing stage. Computer results can omit screenshots in core; OpenAI-based connectors require them when converting to Responses items. diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 0d9aed13f58..7eca7a0fe28 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -481,10 +481,56 @@ class Annotation(TypedDict, total=False): tool_name: str snippet: str annotated_regions: Sequence[TextSpanRegion] + annotated_regions_coordinate_space: Literal["content", "source"] additional_properties: dict[str, Any] raw_representation: Any +class _ValidatedTextSpanRegion(TypedDict): + type: Literal["text_span"] + start_index: int + end_index: int + + +def _get_annotation_text_span_regions( # pyright: ignore[reportUnusedFunction] + annotation: Mapping[str, Any], + *, + coordinate_space: Literal["content", "source"] = "content", +) -> list[_ValidatedTextSpanRegion]: + """Return validated text spans in the requested coordinate space. + + An absent coordinate space retains the legacy meaning that regions refer to the + annotated ``Content.text``. Source-coordinate regions refer to the cited source + and must not be interpreted as generated-content offsets. + """ + if annotation.get("annotated_regions_coordinate_space", "content") != coordinate_space: + return [] + + raw_regions = annotation.get("annotated_regions") + if not isinstance(raw_regions, Sequence) or isinstance(raw_regions, (str, bytes, bytearray)): + return [] + + regions: list[_ValidatedTextSpanRegion] = [] + 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": + continue + start_index = region.get("start_index") + end_index = region.get("end_index") + if type(start_index) is not int or type(end_index) is not int or start_index < 0 or end_index < start_index: + continue + regions.append( + _ValidatedTextSpanRegion( + type="text_span", + start_index=start_index, + end_index=end_index, + ) + ) + return regions + + @experimental(feature_id=ExperimentalFeature.COMPUTER_USE) class ComputerSafetyCheck(TypedDict, total=False): """A provider warning that must be reviewed before running computer actions.""" diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 771e94fb799..626353676ed 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -54,6 +54,7 @@ from agent_framework._types import ( _MODEL_OUTPUT_KIND_KEY, _append_instructions, + _get_annotation_text_span_regions, _get_data_bytes, _get_data_bytes_as_str, _get_operation_state, @@ -4426,7 +4427,13 @@ def test_text_content_with_annotations_serialization(): # Create multiple citations citation1 = Annotation(type="citation", title="Citation 1", url="http://example.com/1", annotated_regions=[region1]) - citation2 = Annotation(type="citation", title="Citation 2", url="http://example.com/2", annotated_regions=[region2]) + citation2 = Annotation( + type="citation", + title="Citation 2", + url="http://example.com/2", + annotated_regions=[region2], + annotated_regions_coordinate_space="source", + ) # Create TextContent with multiple annotations content = Content.from_text(text="Hello world", annotations=[citation1, citation2]) @@ -4438,6 +4445,7 @@ def test_text_content_with_annotations_serialization(): assert len(content_dict["annotations"]) == 2 assert content_dict["annotations"][0]["title"] == "Citation 1" assert content_dict["annotations"][1]["title"] == "Citation 2" + assert content_dict["annotations"][1]["annotated_regions_coordinate_space"] == "source" # Deserialize reconstructed = Content.from_dict(content_dict) @@ -4449,9 +4457,33 @@ def test_text_content_with_annotations_serialization(): assert all(isinstance(ann, dict) for ann in reconstructed.annotations) # type: ignore[union-attr] # pyrefly: ignore[not-iterable] assert reconstructed.annotations[0]["title"] == "Citation 1" # type: ignore[index] # pyrefly: ignore[unsupported-operation] assert reconstructed.annotations[1]["title"] == "Citation 2" # type: ignore[index] # pyrefly: ignore[unsupported-operation] + assert reconstructed.annotations[1]["annotated_regions_coordinate_space"] == "source" # type: ignore[index] # pyrefly: ignore[unsupported-operation] assert all(isinstance(ann["annotated_regions"][0], dict) for ann in reconstructed.annotations) # type: ignore[union-attr] # pyrefly: ignore[not-iterable] +def test_get_annotation_text_span_regions_filters_coordinate_space_and_invalid_regions() -> None: + annotation = cast( + Annotation, + { + "type": "citation", + "annotated_regions_coordinate_space": "source", + "annotated_regions": [ + {"type": "text_span", "start_index": 2, "end_index": 7}, + {"type": "text_span", "start_index": True, "end_index": 7}, + {"type": "text_span", "start_index": -1, "end_index": 7}, + {"type": "text_span", "start_index": 8, "end_index": 7}, + {"type": "source_block", "start_index": 0, "end_index": 1}, + None, + ], + }, + ) + + assert _get_annotation_text_span_regions(annotation) == [] + assert _get_annotation_text_span_regions(annotation, coordinate_space="source") == [ + {"type": "text_span", "start_index": 2, "end_index": 7} + ] + + # region FunctionTool.parse_result with Pydantic models diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 3a352eba761..05a3f97379b 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 `annotated_regions_coordinate_space == "source"` retain source-document coordinates +(for example, Anthropic character, page, or block locations); the host does not emit those 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 54e1323f625..b9d7029646d 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -50,8 +50,12 @@ 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._types import _get_annotation_text_span_regions # pyright: ignore[reportPrivateUsage] from agent_framework._workflows._agent import _StreamingRoleGate # pyright: ignore[reportPrivateUsage] from agent_framework.exceptions import AgentFrameworkException from agent_framework.observability import OBSERVABILITY_SETTINGS @@ -64,12 +68,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, @@ -89,6 +98,7 @@ ResponseUsage, ResponseUsageInputTokensDetails, ResponseUsageOutputTokensDetails, + UrlCitationBody, ) from azure.ai.agentserver.responses.streaming._builders import ( OutputItemBuilder, @@ -273,6 +283,82 @@ def _container_file_citations_from_function_result(content: Content) -> list[_Co return citations +def _annotations_to_output_text( + annotations: Sequence[FrameworkAnnotation] | None, +) -> list[ResponsesAnnotation]: + """Convert framework citations into Responses output-text annotations. + + Source-coordinate regions describe the cited source, not generated content, and + are not emitted as output-text spans. + """ + if not annotations or not isinstance(annotations, Sequence) or isinstance(annotations, (str, bytes)): + return [] + + output_annotations: list[ResponsesAnnotation] = [] + for annotation in annotations: + if not isinstance(annotation, Mapping): + continue + annotation_mapping = cast("Mapping[str, Any]", annotation) + if annotation_mapping.get("type") != "citation": + continue + + raw_properties = annotation_mapping.get("additional_properties") + properties: Mapping[str, Any] = ( + cast("Mapping[str, Any]", raw_properties) if isinstance(raw_properties, Mapping) else {} + ) + regions = _get_annotation_text_span_regions(annotation_mapping) + 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") + + if isinstance(container_id, str) and container_id and isinstance(file_id, str) and file_id: + for region in regions: + output_annotations.append( + ContainerFileCitationBody( + type="container_file_citation", + container_id=container_id, + file_id=file_id, + filename=url if isinstance(url, str) else "", + start_index=region["start_index"], + end_index=region["end_index"], + ) + ) + elif isinstance(url, str) and url and not file_id: + for region in regions: + output_annotations.append( + UrlCitationBody( + type="url_citation", + url=url, + title=title if isinstance(title, str) else "", + start_index=region["start_index"], + end_index=region["end_index"], + ) + ) + elif ( + isinstance(file_id, str) and file_id and isinstance(url, str) and url and type(index) is int and index >= 0 + ): + output_annotations.append( + FileCitationBody( + type="file_citation", + file_id=file_id, + filename=url, + index=index, + ) + ) + elif isinstance(file_id, str) and file_id and type(index) is int and index >= 0: + 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" @@ -1764,7 +1850,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 @@ -2425,7 +2513,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 @@ -2556,15 +2645,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( @@ -2573,6 +2670,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. @@ -2584,9 +2682,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) @@ -2594,10 +2694,13 @@ 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) + 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: yield self._text_content.emit_delta(delta) @@ -2910,11 +3013,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() @@ -2942,6 +3047,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]: @@ -3051,7 +3157,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) @@ -3063,6 +3172,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/agent_framework_foundry_hosting/_workflow_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_responses.py index 811ecacbfb7..ea84c06aa1c 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_responses.py @@ -246,6 +246,7 @@ async def _project( approvals: _CheckpointApprovals, ) -> list[ResponseStreamEvent]: updates: list[AgentResponseUpdate] = [] + complete_message = False if event.type == "request_info": wire_id = IdGenerator.new_id("call", str(stream.response["id"])) if isinstance(event.data, Content) and event.data.type == "function_approval_request": @@ -267,6 +268,7 @@ async def _project( elif event.type in ("output", "intermediate", "data"): data = event.data if isinstance(data, (AgentResponse, ChatResponse)): + complete_message = True updates = _agent_response_updates( AgentResponse( messages=data.messages, @@ -285,8 +287,10 @@ async def _project( ) ] elif isinstance(data, Message): + complete_message = True updates = _agent_response_updates(AgentResponse(messages=[data]), str(stream.response["id"])) elif isinstance(data, list) and data and all(isinstance(item, Message) for item in cast(list[Any], data)): + complete_message = True updates = _agent_response_updates( AgentResponse(messages=cast(list[Message], data)), str(stream.response["id"]) ) @@ -302,7 +306,14 @@ async def _project( return events events: list[ResponseStreamEvent] = [] for update in updates: - events.extend([item async for item in tracker.handle_update(update, approval_storage=approvals)]) + events.extend([ + item + async for item in tracker.handle_update( + update, + approval_storage=approvals, + complete_message=complete_message, + ) + ]) return events diff --git a/python/packages/foundry_hosting/tests/foundry_hosting/test_native_responses.py b/python/packages/foundry_hosting/tests/foundry_hosting/test_native_responses.py index e1583ded9a1..bd90420d6a0 100644 --- a/python/packages/foundry_hosting/tests/foundry_hosting/test_native_responses.py +++ b/python/packages/foundry_hosting/tests/foundry_hosting/test_native_responses.py @@ -16,6 +16,8 @@ Agent, AgentExecutor, AgentResponse, + AgentResponseUpdate, + Annotation, BaseChatClient, ChatOptions, ChatResponse, @@ -286,6 +288,25 @@ def _texts(response: dict[str, Any]) -> list[str]: ] +def _url_citation(*, start_index: int, end_index: int) -> Annotation: + return Annotation( + type="citation", + url="http://test", + title="Test URL", + annotated_regions=[{"type": "text_span", "start_index": start_index, "end_index": end_index}], + ) + + +def _expected_url_citation(*, start_index: int, end_index: int) -> dict[str, Any]: + return { + "type": "url_citation", + "url": "http://test", + "title": "Test URL", + "start_index": start_index, + "end_index": end_index, + } + + async def test_typed_native_turns_use_fresh_graphs_and_exact_state() -> None: graphs: list[Workflow] = [] scopes: list[FoundryRequestScope] = [] @@ -523,6 +544,66 @@ async def start(self, message: str, ctx: WorkflowContext[str, AgentResponse]) -> assert "private-response" not in json.dumps(response) +@pytest.mark.parametrize("payload_type", ["message", "message_list", "agent_response", "chat_response"]) +async def test_completed_framework_output_preserves_adjacent_cited_text_parts(payload_type: str) -> None: + class _Output(Executor): + def __init__(self) -> None: + super().__init__("output") + + @handler + async def start(self, message: str, ctx: WorkflowContext[str, Any]) -> None: + output_message = Message( + "assistant", + [ + Content.from_text("See "), + Content.from_text("Python", annotations=[_url_citation(start_index=0, end_index=6)]), + ], + ) + if payload_type == "message": + output: Any = output_message + elif payload_type == "message_list": + output = [output_message] + elif payload_type == "agent_response": + output = AgentResponse(messages=[output_message]) + else: + output = ChatResponse(messages=[output_message]) + await ctx.yield_output(output) + + response = await _post(_server(lambda request: _build(_Output()), parser=_parse_text), "go", store=False) + + assert response["status"] == "completed", response + message = next(item for item in response["output"] if item["type"] == "message") + assert [part["text"] for part in message["content"]] == ["See ", "Python"] + assert message["content"][0]["annotations"] == [] + assert message["content"][1]["annotations"] == [_expected_url_citation(start_index=0, end_index=6)] + + +@pytest.mark.parametrize("payload_type", ["agent_response_update", "chat_response_update"]) +async def test_streaming_framework_output_keeps_occurrence_relative_citation_offsets(payload_type: str) -> None: + class _Output(Executor): + def __init__(self) -> None: + super().__init__("output") + + @handler + async def start(self, message: str, ctx: WorkflowContext[str, Any]) -> None: + contents = [ + Content.from_text("See "), + Content.from_text("Python", annotations=[_url_citation(start_index=4, end_index=10)]), + ] + if payload_type == "agent_response_update": + output: Any = AgentResponseUpdate(contents=contents, role="assistant", message_id="message") + else: + output = ChatResponseUpdate(contents=contents, role="assistant", message_id="message") + await ctx.yield_output(output) + + response = await _post(_server(lambda request: _build(_Output()), parser=_parse_text), "go", store=False) + + assert response["status"] == "completed", response + message = next(item for item in response["output"] if item["type"] == "message") + assert [part["text"] for part in message["content"]] == ["See Python"] + assert message["content"][0]["annotations"] == [_expected_url_citation(start_index=4, end_index=10)] + + async def test_native_agent_executor_applies_options_without_mutating_defaults_or_double_keywords() -> None: calls: list[dict[str, Any]] = [] lifecycle: list[str] = [] diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index aa9a5f0ce90..84e6ed6e691 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,177 @@ 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) == [ + { + "type": "url_citation", + "url": "https://example.com/source", + "title": "Source", + "start_index": 0, + "end_index": 6, + }, + { + "type": "url_citation", + "url": "https://example.com/source", + "title": "Source", + "start_index": 12, + "end_index": 18, + }, + _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": 20, + "end_index": 30, + }, + ] + + +@pytest.mark.parametrize( + ("coordinate_space", "expected"), + [ + pytest.param(None, [_expected_url_citation()], id="unmarked"), + pytest.param("content", [_expected_url_citation()], id="content"), + pytest.param("source", [], id="source"), + ], +) +def test_framework_annotations_respect_coordinate_space_without_raw_representation( + coordinate_space: Literal["content", "source"] | None, expected: list[dict[str, Any]] +) -> None: + annotation = _url_citation() + if coordinate_space is not None: + annotation["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["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, + [ + 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": True, "end_index": 5}, + {"type": "text_span", "start_index": -1, "end_index": 5}, + {"type": "text_span", "start_index": 8, "end_index": 7}, + {"type": "text_span", "start_index": 2, "end_index": 7}, + ], + }, + { + "type": "citation", + "url": "https://example.com/non-string-title", + "title": 42, + "annotated_regions": [{"type": "text_span", "start_index": 8, "end_index": 9}], + }, + { + "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, + ), + _expected_url_citation( + url="https://example.com/non-string-title", + title="", + start_index=8, + end_index=9, + ), + ] + 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"), + 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": "citation", "file_id": "file-1", "additional_properties": {"index": True}}), + cast(Annotation, {"type": "citation", "file_id": "file-1", "additional_properties": {"index": -1}}), + cast(Annotation, {"type": "other", "url": "https://example.com/unsupported"}), + ] + + assert _annotations_to_output_text(annotations) == [] + + # region Helpers @@ -2016,6 +2248,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: @@ -3835,6 +4130,33 @@ 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: + source_citation = _url_citation(start_index=0, end_index=1) + source_citation["annotated_regions_coordinate_space"] = "source" + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate( + contents=[ + Content.from_text( + "Python sources", + annotations=[source_citation, _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=[ @@ -4658,6 +4980,60 @@ 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["annotated_regions_coordinate_space"] = "source" + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate( + contents=[Content.from_text("See ", annotations=[source_citation])], + role="assistant", + message_id="msg_citation", + ), + AgentResponseUpdate( + contents=[ + Content.from_text( + "Python", + annotations=[_url_citation(start_index=4, end_index=10)], + ) + ], + 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() diff --git a/python/packages/openai/AGENTS.md b/python/packages/openai/AGENTS.md index 472c0d780ce..da64504352a 100644 --- a/python/packages/openai/AGENTS.md +++ b/python/packages/openai/AGENTS.md @@ -77,7 +77,9 @@ marked text is sent as ordinary input text. Responses citations remain `Content.annotations`, including annotation-only updates parsed from completed text parts and output messages. Streaming deduplication is request-local and keyed by item ID, content index, and annotation index, so separate references to the same URL remain distinct. `_parse_chunk_from_openai` retains -the signature overridden by released Foundry clients. +the signature overridden by released Foundry clients. When assistant history is serialized back to Responses, +only content-coordinate annotation regions become output-text spans; source-coordinate regions remain framework +metadata and must not be interpreted as generated-text offsets. ## Dependencies diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index fdb8c3d8eae..be4eec443cd 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -72,6 +72,7 @@ Role, TextSpanRegion, UsageDetails, + _get_annotation_text_span_regions, # pyright: ignore[reportPrivateUsage] _set_operation_state, # pyright: ignore[reportPrivateUsage] detect_media_type_from_base64, get_namespaced_value, @@ -530,7 +531,8 @@ def _annotations_to_output_text(annotations: Sequence[Annotation] | None) -> lis Each Responses API annotation dict carries at most one `start_index`/`end_index` pair, so an `Annotation` with multiple `annotated_regions` is fanned out into one entry per region. - Regions missing valid integer span bounds are skipped. + Regions missing valid integer span bounds are skipped. Source-coordinate regions + are retained on the framework annotation but are not emitted as generated-text spans. """ if not annotations: return [] @@ -539,7 +541,7 @@ def _annotations_to_output_text(annotations: Sequence[Annotation] | None) -> lis if annotation.get("type") != "citation": continue props = annotation.get("additional_properties") or {} - regions = annotation.get("annotated_regions") or [] + regions = _get_annotation_text_span_regions(annotation) file_id = annotation.get("file_id") url = annotation.get("url") title = annotation.get("title") diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index ca9870e1ae3..e2991753bce 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -33,6 +33,7 @@ SupportsImageGenerationTool, SupportsMCPTool, SupportsWebSearchTool, + TextSpanRegion, WorkflowBuilder, tool, ) @@ -5995,6 +5996,27 @@ def test_assistant_text_fans_out_multiple_annotated_regions() -> None: assert spans == [(4, 10), (16, 22), (36, 42)] +def test_assistant_text_does_not_emit_source_coordinate_regions_as_output_spans() -> None: + """Source locations must not be serialized as generated-response offsets.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + text_content = Content.from_text( + "Answer", + annotations=[ + Annotation( + type="citation", + title="Source document", + url="https://example.com/source", + annotated_regions=[TextSpanRegion(type="text_span", start_index=40, end_index=46)], + annotated_regions_coordinate_space="source", + ) + ], + ) + + result = client._prepare_content_for_openai("assistant", text_content) + + assert result["annotations"] == [] + + def test_assistant_text_skips_regions_with_invalid_span() -> None: """Regions missing integer start/end bounds are skipped rather than emitted with `None`.""" from agent_framework._types import Annotation, TextSpanRegion