From eb8a467ffffbb4d31456ec6faa2bd3c754ded8d5 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 1 Oct 2026 17:36:54 +0200 Subject: [PATCH 1/2] Python: preserve Responses citations in AG-UI --- python/packages/ag-ui/AGENTS.md | 4 + python/packages/ag-ui/README.md | 36 +++ .../_event_converters.py | 23 +- .../agent_framework_ag_ui/_run_common.py | 42 ++- .../tests/ag_ui/test_event_converters.py | 56 ++++ python/packages/ag-ui/tests/ag_ui/test_run.py | 152 ++++++++- python/packages/openai/AGENTS.md | 5 + python/packages/openai/README.md | 15 + .../agent_framework_openai/_chat_client.py | 206 ++++++------ .../tests/openai/test_openai_chat_client.py | 302 +++++++++++++++++- 10 files changed, 732 insertions(+), 109 deletions(-) diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 8de91344a1a..b5f44d19413 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -27,6 +27,10 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - Outbound custom events are emitted as AG-UI `CUSTOM`. - Usage metadata from `Content(type="usage")` is surfaced as `CUSTOM` events with `name="usage"`. +- 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`. - 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/README.md b/python/packages/ag-ui/README.md index 2a53d7f3ac2..f392a061ca9 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -119,6 +119,42 @@ and response-cookie handling, and must be closed by the caller. Scope a cookie-b client to a single authenticated principal; do not share it across users. AG-UI thread IDs are correlation identifiers, not authentication boundaries. +## Citations and annotations + +Text content carrying `Content.annotations` emits a message-linked `CUSTOM` event named `annotations`. +This includes annotation-only updates received after the response text, such as SharePoint grounding +citations from the Responses API. The event arrives before `RUN_FINISHED` and does not repeat response text: + +```json +{ + "type": "CUSTOM", + "name": "annotations", + "value": { + "messageId": "assistant-message-id", + "annotations": [ + { + "type": "citation", + "title": "Document", + "url": "https://example.sharepoint.com/document.pdf", + "annotated_regions": [ + {"type": "text_span", "start_index": 0, "end_index": 6} + ] + } + ] + } +} +``` + +Each event contains a batch of newly emitted annotations for the indicated text message. Frontends can +append those annotations to that message and use the citation URL, title, or file ID to render sources. +Framework annotation fields and additional properties are retained, but provider `raw_representation` +objects are omitted. + +`AGUIChatClient` restores these batches as `Content.annotations` in streaming updates and aggregated +responses, while retaining the custom-event payload in `update.additional_properties["ag_ui_custom_event"]`. +Custom events are live run metadata: citation rendering and persistence alongside frontend message history +remain application responsibilities; `MESSAGES_SNAPSHOT` does not restore these custom events. + ## Tool Return Helpers Use `state_update` when a backend tool needs to send different payloads to the model, the UI, and shared state. The `text` value remains the LLM-bound tool result, `tool_result` becomes the AG-UI `ToolCallResultEvent.content` for frontend rendering, and `state` is merged into durable shared state. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py b/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py index 447af1654d9..a161f5b7e18 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py @@ -5,9 +5,10 @@ from __future__ import annotations import logging -from typing import Any +from typing import Any, cast from agent_framework import ( + Annotation, ChatResponseUpdate, Content, ) @@ -261,9 +262,9 @@ def _handle_run_error(self, event: dict[str, Any]) -> ChatResponseUpdate: def _handle_custom_event(self, event: dict[str, Any], raw_event_type: str) -> ChatResponseUpdate: """Handle CUSTOM/CUSTOM_EVENT events. - Custom events are surfaced as metadata so callers can inspect protocol-specific payloads. + Custom events remain inspectable as metadata; annotation batches also restore text annotations. """ - return ChatResponseUpdate( + update = ChatResponseUpdate( role="assistant", contents=[], additional_properties={ @@ -276,3 +277,19 @@ def _handle_custom_event(self, event: dict[str, Any], raw_event_type: str) -> Ch }, }, ) + if event.get("name") == "annotations": + value = event.get("value") + message_id = value.get("messageId") if isinstance(value, dict) else None + annotations = value.get("annotations") if isinstance(value, dict) else None + if ( + not isinstance(message_id, str) + or not message_id + or not isinstance(annotations, list) + or not all(isinstance(annotation, dict) for annotation in annotations) + ): + logger.warning("Invalid annotations custom event: expected messageId and an annotations array") + else: + update.message_id = message_id + if annotations: + update.contents = [Content.from_text(text="", annotations=cast("list[Annotation]", annotations))] + return update 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 5c47a105b8e..b88a1c1d5fe 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 @@ -610,14 +610,15 @@ def _track_reasoning_segment(flow: FlowState, message_id: str) -> None: def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]: - """Emit TextMessage events for TextContent.""" - if not content.text: + """Emit text deltas and message-linked annotation batches.""" + if not content.text and not content.annotations: return [] if skip_text or flow.waiting_for_approval: return [] events: list[BaseEvent] = [] + duplicate_text = False if not flow.message_id: flow.message_id = generate_event_id() flow.accumulated_text = "" @@ -626,7 +627,9 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li elif flow.accumulated_text and content.text == flow.accumulated_text: # Guard against full-message replay chunks that can appear after streaming deltas. logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id) - return [] + if not content.annotations: + return [] + duplicate_text = True # A tool-only response may pre-open a message before its tool-call segment # is tracked. If that segment claims the pre-opened ID, rotate to a fresh @@ -648,9 +651,38 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li if segment is None: segment = _open_text_segment(flow, flow.message_id) - events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text)) - flow.accumulated_text += content.text + if content.text and not duplicate_text: + events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text)) + flow.accumulated_text += content.text segment["text"] = flow.accumulated_text + if content.annotations: + annotations = cast( + "list[dict[str, Any]]", + make_json_safe( + [ + {key: value for key, value in annotation.items() if key != "raw_representation"} + for annotation in content.annotations + ] + ), + ) + new_annotations = annotations + if duplicate_text: + # Reconcile replay by occurrence count, not source URL or dictionary uniqueness. + remaining_annotations = list(segment.get("annotations", [])) + new_annotations = [] + for annotation in annotations: + if annotation in remaining_annotations: + remaining_annotations.remove(annotation) + else: + new_annotations.append(annotation) + if new_annotations: + segment.setdefault("annotations", []).extend(new_annotations) + events.append( + CustomEvent( + name="annotations", + value={"messageId": flow.message_id, "annotations": new_annotations}, + ) + ) return events diff --git a/python/packages/ag-ui/tests/ag_ui/test_event_converters.py b/python/packages/ag-ui/tests/ag_ui/test_event_converters.py index 53f02c434c8..2b9323a4bd7 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_event_converters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_event_converters.py @@ -391,6 +391,62 @@ def test_custom_event_alias_conversion(self) -> None: assert updates[0].additional_properties["ag_ui_custom_event"]["raw_type"] == "CUSTOM_EVENT" assert updates[1].additional_properties["ag_ui_custom_event"]["raw_type"] == "custom_event" + @pytest.mark.parametrize("event_type", ["CUSTOM", "CUSTOM_EVENT", "custom_event"]) + def test_annotations_custom_event_restores_content(self, event_type: str) -> None: + """Annotation custom events preserve their explicit message correlation and metadata.""" + converter = AGUIEventConverter() + converter.current_message_id = "another-message" + annotations = [ + { + "type": "citation", + "title": "Document", + "url": "https://example.sharepoint.com/document.pdf", + "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 6}], + } + ] + event = { + "type": event_type, + "name": "annotations", + "value": {"messageId": "msg_citations", "annotations": annotations}, + } + + update = converter.convert_event(event) + + assert update is not None + assert update.message_id == "msg_citations" + assert len(update.contents) == 1 + assert update.contents[0].type == "text" + assert update.contents[0].text == "" + assert update.contents[0].annotations == annotations + assert update.additional_properties is not None + assert update.additional_properties["ag_ui_custom_event"]["value"] == event["value"] + + @pytest.mark.parametrize( + "value", + [ + None, + {"annotations": []}, + {"messageId": 123, "annotations": []}, + {"messageId": "msg_citations", "annotations": "invalid"}, + {"messageId": "msg_citations", "annotations": ["invalid"]}, + ], + ) + def test_malformed_annotations_custom_event_preserves_metadata( + self, value: Any, caplog: pytest.LogCaptureFixture + ) -> None: + """Malformed known custom events remain observable and log an explicit warning.""" + converter = AGUIEventConverter() + event = {"type": "CUSTOM", "name": "annotations", "value": value} + + with caplog.at_level(logging.WARNING): + update = converter.convert_event(event) + + assert update is not None + assert update.contents == [] + assert update.additional_properties is not None + assert update.additional_properties["ag_ui_custom_event"]["value"] == value + assert "annotations" in caplog.text + def test_full_conversation_flow(self) -> None: """Test complete conversation flow with multiple event types.""" converter = AGUIEventConverter() 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 f4d7b02fdea..f487f4a2ed3 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -21,7 +21,15 @@ ToolCallArgsEvent, ToolCallStartEvent, ) -from agent_framework import AgentResponse, AgentResponseUpdate, Content, Message, ResponseStream +from agent_framework import ( + AgentResponse, + AgentResponseUpdate, + Annotation, + ChatResponse, + Content, + Message, + ResponseStream, +) from agent_framework.exceptions import AgentInvalidResponseException, ResponseInvalidatedException from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] @@ -45,6 +53,7 @@ ResumeDecision, ) from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id +from agent_framework_ag_ui._event_converters import AGUIEventConverter from agent_framework_ag_ui._run_common import ( FlowState, _build_run_finished_event, @@ -631,6 +640,147 @@ def test_emit_text_skips_when_waiting_for_approval(): assert len(events) == 0 +def test_emit_text_annotations_are_json_safe_and_message_linked() -> None: + """Portable citations reach AG-UI without serializing provider raw objects.""" + flow = FlowState() + annotation = Annotation( + type="citation", + title="Document", + url="https://example.sharepoint.com/document.pdf", + annotated_regions=[{"type": "text_span", "start_index": 0, "end_index": 6}], + raw_representation=object(), + ) + content = Content.from_text("Answer", annotations=[annotation]) + + events = _emit_content(content, flow) + + assert len(events) == 3 + start, text, citations = events + assert isinstance(start, TextMessageStartEvent) + assert isinstance(text, TextMessageContentEvent) + assert isinstance(citations, CustomEvent) + assert citations.name == "annotations" + assert citations.value == { + "messageId": start.message_id, + "annotations": [{key: value for key, value in annotation.items() if key != "raw_representation"}], + } + assert text.message_id == start.message_id + assert "raw_representation" not in citations.model_dump_json(by_alias=True) + assert content.annotations is not None + assert annotation["raw_representation"] is content.annotations[0]["raw_representation"] + + +def test_emit_text_annotation_only_update_starts_a_message_without_an_empty_delta() -> None: + """Citations received before text still reference an announced message.""" + flow = FlowState() + annotation = Annotation(type="citation", url="https://example.com/document") + + events = _emit_text(Content.from_text("", annotations=[annotation]), flow) + text_events = _emit_text(Content.from_text("Answer"), flow) + + assert len(events) == 2 + assert isinstance(events[0], TextMessageStartEvent) + assert isinstance(events[1], CustomEvent) + assert events[1].value == {"messageId": events[0].message_id, "annotations": [annotation]} + assert len(text_events) == 1 + assert isinstance(text_events[0], TextMessageContentEvent) + assert text_events[0].message_id == events[0].message_id + assert text_events[0].delta == "Answer" + + +def test_emit_text_replay_emits_only_new_annotations() -> None: + """Suppress repeated text and citations without dropping final-only citations.""" + flow = FlowState() + first = Annotation(type="citation", url="https://example.com/first") + second = Annotation(type="citation", url="https://example.com/second") + _emit_text(Content.from_text("Answer", annotations=[first]), flow) + + events = _emit_text(Content.from_text("Answer", annotations=[first, second]), flow) + replay_events = _emit_text(Content.from_text("Answer", annotations=[first, second]), flow) + + assert len(events) == 1 + assert isinstance(events[0], CustomEvent) + assert events[0].value == {"messageId": flow.message_id, "annotations": [second]} + assert replay_events == [] + assert flow.accumulated_text == "Answer" + + +def test_emit_text_preserves_identical_annotation_occurrences() -> None: + """Separate annotation-only updates can cite the same source at distinct occurrences.""" + flow = FlowState() + annotation = Annotation(type="citation", file_id="file-123") + _emit_text(Content.from_text("Answer"), flow) + + first_events = _emit_text(Content.from_text("", annotations=[annotation]), flow) + second_events = _emit_text(Content.from_text("", annotations=[annotation]), flow) + replay_events = _emit_text(Content.from_text("Answer", annotations=[annotation, annotation]), flow) + + assert len(first_events) == len(second_events) == 1 + assert all(isinstance(event, CustomEvent) for event in [*first_events, *second_events]) + assert replay_events == [] + + +@pytest.mark.parametrize("suppression", ["skip_text", "waiting_for_approval"]) +def test_emit_text_annotations_respect_suppression(suppression: str) -> None: + """Explicitly suppressed text must not leak an annotation message either.""" + flow = FlowState(waiting_for_approval=suppression == "waiting_for_approval") + content = Content.from_text("Answer", annotations=[Annotation(type="citation", url="https://example.com")]) + + events = _emit_text(content, flow, skip_text=suppression == "skip_text") + + assert events == [] + assert flow.message_id is None + + +async def test_run_agent_stream_delivers_annotations_before_run_finished() -> None: + """A grounded response roundtrips through live AG-UI events with its citations intact.""" + annotation = Annotation( + type="citation", + title="Document", + url="https://example.sharepoint.com/document.pdf", + annotated_regions=[{"type": "text_span", "start_index": 0, "end_index": 6}], + raw_representation=object(), + ) + agent = StubAgent( + updates=[ + AgentResponseUpdate(role="assistant", contents=[Content.from_text("Answer")]), + AgentResponseUpdate(role="assistant", contents=[Content.from_text("", annotations=[annotation])]), + ] + ) + + events = [ + event + async for event in run_agent_stream( + { + "run_id": "run-citations", + "thread_id": "thread-citations", + "messages": [{"role": "user", "content": "Find the document."}], + }, + agent, + AgentConfig(), + ) + ] + + citation_events = [event for event in events if isinstance(event, CustomEvent) and event.name == "annotations"] + assert len(citation_events) == 1 + text_event = next(event for event in events if isinstance(event, TextMessageContentEvent)) + assert citation_events[0].value["messageId"] == text_event.message_id + assert events.index(citation_events[0]) < next( + index for index, event in enumerate(events) if event.type == "RUN_FINISHED" + ) + converter = AGUIEventConverter() + updates = [] + for event in events: + if update := converter.convert_event(event.model_dump(mode="json", by_alias=True, exclude_none=True)): + updates.append(update) + response = ChatResponse.from_updates(updates) + assert response.text == "Answer" + text_content = next(content for message in response.messages for content in message.contents if content.text) + assert text_content.annotations == [ + {key: value for key, value in annotation.items() if key != "raw_representation"} + ] + + def _snapshot_kinds(event): """Flatten a MessagesSnapshotEvent into (kind, message) pairs for order assertions.""" kinds = [] diff --git a/python/packages/openai/AGENTS.md b/python/packages/openai/AGENTS.md index 2dc541ca802..f2478efb98d 100644 --- a/python/packages/openai/AGENTS.md +++ b/python/packages/openai/AGENTS.md @@ -66,6 +66,11 @@ Native Responses and Chat Completions refusals remain ordinary text content with the provider's native refusal field without splitting mixed text/refusal turns; non-assistant 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. + ## Dependencies - `agent-framework-core` — core abstractions diff --git a/python/packages/openai/README.md b/python/packages/openai/README.md index 1f759c04f6b..5be069741a6 100644 --- a/python/packages/openai/README.md +++ b/python/packages/openai/README.md @@ -24,6 +24,21 @@ Use `OpenAIChatClient` for new work unless you specifically need the Chat Comple The previous deprecated Responses alias has been removed. Use `OpenAIChatClient` directly. +## Citations and annotations + +`OpenAIChatClient` exposes Responses API URL, file, file-path, and container-file citations through +`Content.annotations`, using the provider-neutral `Annotation(type="citation", ...)` representation. +URL citations retain their title, URL, and text-span indices, including SharePoint grounding citations. + +Inspect `content.annotations` in response messages or streaming `update.contents`. Annotations can arrive +in an update whose `content.text` is empty: completed content parts and output messages supply citations +without repeating text that was already streamed. Annotation-added and completion events are deduplicated +by item, content part, and annotation position, so separate references to the same source remain distinct. +Foundry clients that inherit the Responses parser receive the same behavior. + +The AG-UI integration surfaces these annotations as message-linked `CUSTOM` events named `annotations`; +see [AG-UI citation support](../ag-ui/README.md#citations-and-annotations). No logging handler is needed. + ## Hosted function results `OpenAIChatClient` parses hosted `function_call_output` items in both streaming and non-streaming responses. diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index d7b7e0735cb..917082386d1 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -404,7 +404,7 @@ def _guess_media_type(filename: str | None, fallback: str | None = "application/ return fallback -# Per-request dedup state for `function_call_output` items, carried in the parse `options` mapping +# Per-request dedup state for output items and text annotations, carried in the parse `options` mapping # rather than on `_parse_chunk_from_openai`'s signature. # # Adding a parameter would break released `agent-framework-foundry` wheels: they pin @@ -414,6 +414,7 @@ def _guess_media_type(filename: str | None, fallback: str | None = "application/ # `options` here is the validated *parse* context, not the outbound request body -- the request is # built from a separate `run_options` dict -- so a private key cannot leak to the service. _SEEN_FUNCTION_CALL_OUTPUT_IDS_OPTION = "__agent_framework_seen_function_call_output_ids" +_SEEN_TEXT_ANNOTATION_KEYS_OPTION = "__agent_framework_seen_text_annotation_keys" def _pairable_function_call_output_call_id(item: ResponseFunctionToolCallOutputItem) -> str | None: @@ -828,6 +829,7 @@ def _inner_get_response( function_call_ids: dict[int, tuple[str, str]] = {} seen_reasoning_delta_item_ids: set[str] = set() seen_function_call_output_ids: set[str] = set() + seen_text_annotation_keys: set[tuple[str, int, int]] = set() validated_options: dict[str, Any] | None = None # Captured once request options are validated/prepared so the streaming finalizer can # still parse the aggregated response into structured output after the stream completes. @@ -848,6 +850,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: mark_feature_used(self._FEATURE_USAGE_INDEX) validated_options = await self._validate_options(options) validated_options[_SEEN_FUNCTION_CALL_OUTPUT_IDS_OPTION] = seen_function_call_output_ids + validated_options[_SEEN_TEXT_ANNOTATION_KEYS_OPTION] = seen_text_annotation_keys response_format = validated_options.get("response_format") try: raw_stream_response = await client.responses.with_raw_response.retrieve( @@ -893,6 +896,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: validated_options, ) = await self._prepare_request(messages, options) validated_options[_SEEN_FUNCTION_CALL_OUTPUT_IDS_OPTION] = seen_function_call_output_ids + validated_options[_SEEN_TEXT_ANNOTATION_KEYS_OPTION] = seen_text_annotation_keys if extra_headers is not None: run_options["extra_headers"] = dict(extra_headers) response_format = validated_options.get("response_format") @@ -3515,6 +3519,61 @@ def _parse_response_from_openai( self._enrich_mcp_search_citations(chat_response.messages[0].contents) return chat_response + def _parse_annotation_from_openai(self, annotation: Any, annotation_index: int) -> Annotation | None: + """Parse a text annotation from an SDK object or streaming annotation dictionary.""" + + def get_value(key: str) -> Any: + if isinstance(annotation, dict): + return cast("dict[str, Any]", annotation).get(key) + return getattr(annotation, key, None) + + annotation_type = get_value("type") + properties: dict[str, Any] = {"annotation_index": annotation_index} + if annotation_type in ("file_path", "file_citation", "container_file_citation"): + file_id = get_value("file_id") + if not file_id: + logger.debug("Skipping %s annotation without a file_id", annotation_type) + return None + parsed = Annotation( + type="citation", + file_id=str(file_id), + additional_properties=properties, + raw_representation=annotation, + ) + if annotation_type == "container_file_citation": + properties["container_id"] = get_value("container_id") + else: + properties["index"] = get_value("index") + if annotation_type != "file_path": + parsed["url"] = get_value("filename") + elif annotation_type == "url_citation": + url = get_value("url") + if not url: + logger.debug("Skipping url_citation annotation without a url") + return None + parsed = Annotation( + type="citation", + title=get_value("title") or "", + url=str(url), + additional_properties=properties, + raw_representation=annotation, + ) + get_url = get_value("get_url") + if get_url is not None: + properties["get_url"] = get_url + else: + logger.debug("Unparsed annotation type in streaming: %s", annotation_type) + return None + + if annotation_type in ("url_citation", "container_file_citation"): + start_index = get_value("start_index") + end_index = get_value("end_index") + if start_index is not None and end_index is not None: + parsed["annotated_regions"] = [ + TextSpanRegion(type="text_span", start_index=start_index, end_index=end_index) + ] + return parsed + def _parse_chunk_from_openai( self, event: OpenAIResponseStreamEvent, @@ -3526,6 +3585,9 @@ def _parse_chunk_from_openai( # Read from `options` rather than a parameter; see # `_SEEN_FUNCTION_CALL_OUTPUT_IDS_OPTION` for why the signature must not grow. seen_function_call_output_ids = cast("set[str] | None", options.get(_SEEN_FUNCTION_CALL_OUTPUT_IDS_OPTION)) + seen_text_annotation_keys = cast( + "set[tuple[str, int, int]] | None", options.get(_SEEN_TEXT_ANNOTATION_KEYS_OPTION) + ) metadata: dict[str, Any] = {} contents: list[Content] = [] local_shell_tool_name = self._get_local_shell_tool_name(options.get("tools")) @@ -3545,6 +3607,26 @@ def output_text_properties(output: Any) -> dict[str, Any] | None: return None return {"logprobs": serialized} + def parse_annotations( + annotations: Sequence[Any], + item_id: str, + content_index: int, + *, + annotation_index: int = 0, + ) -> list[Annotation]: + parsed_annotations: list[Annotation] = [] + for index, annotation in enumerate(annotations, start=annotation_index): + key = (item_id, content_index, index) + can_deduplicate = isinstance(item_id, str) and bool(item_id) and isinstance(content_index, int) + if seen_text_annotation_keys is not None and can_deduplicate and key in seen_text_annotation_keys: + continue + parsed = self._parse_annotation_from_openai(annotation, index) + if parsed is not None: + parsed_annotations.append(parsed) + if seen_text_annotation_keys is not None and can_deduplicate: + seen_text_annotation_keys.add(key) + return parsed_annotations + match event.type: # types: # ResponseAudioDeltaEvent, @@ -3604,9 +3686,11 @@ def output_text_properties(output: Any) -> dict[str, Any] | None: event_part = event.part match event_part.type: case "output_text": + annotations = parse_annotations(event_part.annotations, event.item_id, event.content_index) contents.append( Content.from_text( text=event_part.text, + annotations=annotations or None, additional_properties=output_text_properties(cast(Any, event_part)), raw_representation=event, ) @@ -3622,6 +3706,12 @@ def output_text_properties(output: Any) -> dict[str, Any] | None: ) case _: pass + case "response.content_part.done": + event_part = event.part + if event_part.type == "output_text": + annotations = parse_annotations(event_part.annotations, event.item_id, event.content_index) + if annotations: + contents.append(Content.from_text(text="", annotations=annotations, raw_representation=event)) case "response.output_text.delta": contents.append( Content.from_text( @@ -3959,104 +4049,14 @@ def output_text_properties(output: Any) -> dict[str, Any] | None: ) ) case "response.output_text.annotation.added": - # Handle streaming text annotations (file citations, file paths, etc.) - annotation: Any = event.annotation - - def _get_ann_value(key: str) -> Any: - """Extract value from annotation (dict or object).""" - if isinstance(annotation, dict): - return cast("dict[str, Any]", annotation).get(key) - return getattr(annotation, key, None) - - ann_type = _get_ann_value("type") - ann_file_id = _get_ann_value("file_id") - # Hosted-file citations attach as text annotations (matching the non-streaming path) - # so they don't roundtrip as standalone `input_file` items in assistant history. - if ann_type == "file_path": - if ann_file_id: - annotation_obj = Annotation( - type="citation", - file_id=str(ann_file_id), - additional_properties={ - "annotation_index": event.annotation_index, - "index": _get_ann_value("index"), - }, - raw_representation=annotation, - ) - contents.append( - Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) - ) - elif ann_type == "file_citation": - if ann_file_id: - ann_filename = _get_ann_value("filename") - annotation_obj = Annotation( - type="citation", - file_id=str(ann_file_id), - url=ann_filename, - additional_properties={ - "annotation_index": event.annotation_index, - "index": _get_ann_value("index"), - }, - raw_representation=annotation, - ) - contents.append( - Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) - ) - elif ann_type == "container_file_citation": - if ann_file_id: - ann_filename = _get_ann_value("filename") - ann_start = _get_ann_value("start_index") - ann_end = _get_ann_value("end_index") - annotation_obj = Annotation( - type="citation", - file_id=str(ann_file_id), - url=ann_filename, - additional_properties={ - "annotation_index": event.annotation_index, - "container_id": _get_ann_value("container_id"), - }, - raw_representation=annotation, - ) - if ann_start is not None and ann_end is not None: - annotation_obj["annotated_regions"] = [ - TextSpanRegion( - type="text_span", - start_index=ann_start, - end_index=ann_end, - ) - ] - contents.append( - Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) - ) - elif ann_type == "url_citation": - ann_url = _get_ann_value("url") - if ann_url: - ann_start = _get_ann_value("start_index") - ann_end = _get_ann_value("end_index") - annotation_properties: dict[str, Any] = {"annotation_index": event.annotation_index} - ann_get_url = _get_ann_value("get_url") - if ann_get_url is not None: - annotation_properties["get_url"] = ann_get_url - annotation_obj = Annotation( - type="citation", - title=_get_ann_value("title") or "", - url=str(ann_url), - additional_properties=annotation_properties, - raw_representation=annotation, - ) - if ann_start is not None and ann_end is not None: - annotation_obj["annotated_regions"] = [ - TextSpanRegion( - type="text_span", - start_index=ann_start, - end_index=ann_end, - ) - ] - contents.append( - Content.from_text(text="", annotations=[annotation_obj], raw_representation=event) - ) - else: - logger.debug("Unparsed annotation type in streaming: %s", ann_type) + annotations = parse_annotations( + [event.annotation], + event.item_id, + event.content_index, + annotation_index=event.annotation_index, + ) + if annotations: + contents.append(Content.from_text(text="", annotations=annotations, raw_representation=event)) case "response.output_item.done": done_item = event.item if getattr(done_item, "type", None) == "reasoning": @@ -4070,6 +4070,14 @@ def _get_ann_value(key: str) -> Any: raw_representation=done_item, ) ) + elif done_item.type == "message": + for content_index, part in enumerate(done_item.content): + if part.type == "output_text": + annotations = parse_annotations(part.annotations, done_item.id, content_index) + if annotations: + contents.append( + Content.from_text(text="", annotations=annotations, raw_representation=event) + ) elif getattr(done_item, "type", None) == "mcp_call": call_id = getattr(done_item, "id", None) or getattr(done_item, "call_id", None) or "" output_text = getattr(done_item, "output", None) 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 fcb31c961b3..22e4fd9133f 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import base64 import inspect import json @@ -9,7 +10,7 @@ from importlib import import_module from pathlib import Path from types import SimpleNamespace -from typing import Annotated, Any, cast +from typing import Annotated, Any, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -48,8 +49,11 @@ from openai.types.responses import ( ResponseComputerToolCall, ResponseComputerToolCallOutputItem, + ResponseContentPartDoneEvent, ResponseFunctionShellToolCall, ResponseFunctionShellToolCallOutput, + ResponseOutputItemDoneEvent, + ResponseOutputTextAnnotationAddedEvent, ) from openai.types.responses.response_reasoning_item import Summary from openai.types.responses.response_reasoning_summary_text_delta_event import ( @@ -6808,6 +6812,302 @@ def test_streaming_response_in_progress_type() -> None: assert response.conversation_id == "conv_5678" +def _make_completed_text_event( + annotations: Sequence[dict[str, Any]], + *, + event_type: Literal["response.content_part.done", "response.output_item.done"] = "response.content_part.done", + item_id: str = "msg_citations", + content_index: int = 0, + text: str = "The answer is in the document.", +) -> ResponseContentPartDoneEvent | ResponseOutputItemDoneEvent: + part = {"type": "output_text", "text": text, "annotations": list(annotations)} + if event_type == "response.content_part.done": + return ResponseContentPartDoneEvent.model_validate({ + "type": event_type, + "item_id": item_id, + "output_index": 0, + "content_index": content_index, + "sequence_number": 2, + "part": part, + }) + return ResponseOutputItemDoneEvent.model_validate({ + "type": event_type, + "output_index": 0, + "sequence_number": 3, + "item": { + "type": "message", + "id": item_id, + "role": "assistant", + "status": "completed", + "content": [ + *({"type": "output_text", "text": "", "annotations": []} for _ in range(content_index)), + part, + ], + }, + }) + + +@pytest.mark.parametrize("event_type", ["response.content_part.done", "response.output_item.done"]) +@pytest.mark.parametrize( + ("provider_annotation", "expected"), + [ + param( + { + "type": "url_citation", + "url": "https://example.sharepoint.com/sites/docs/document.pdf", + "title": "Document", + "start_index": 21, + "end_index": 29, + }, + { + "url": "https://example.sharepoint.com/sites/docs/document.pdf", + "title": "Document", + "annotated_regions": [{"type": "text_span", "start_index": 21, "end_index": 29}], + }, + id="sharepoint-url", + ), + param( + {"type": "file_citation", "file_id": "file-123", "filename": "document.pdf", "index": 21}, + {"file_id": "file-123", "url": "document.pdf"}, + id="file-citation", + ), + param( + {"type": "file_path", "file_id": "file-123", "index": 21}, + {"file_id": "file-123"}, + id="file-path", + ), + param( + { + "type": "container_file_citation", + "file_id": "file-123", + "container_id": "container-123", + "filename": "document.pdf", + "start_index": 21, + "end_index": 29, + }, + { + "file_id": "file-123", + "url": "document.pdf", + "annotated_regions": [{"type": "text_span", "start_index": 21, "end_index": 29}], + }, + id="container-file", + ), + ], +) +def test_streaming_completed_text_preserves_annotations_without_replaying_text( + event_type: Literal["response.content_part.done", "response.output_item.done"], + provider_annotation: dict[str, Any], + expected: dict[str, Any], +) -> None: + """Completed text carries citations even when annotation-added events are absent.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + event = _make_completed_text_event([provider_annotation], event_type=event_type) + + update = client._parse_chunk_from_openai(event, {}, {}) + + assert len(update.contents) == 1 + content = update.contents[0] + assert content.type == "text" + assert content.text == "" + assert content.annotations is not None + assert len(content.annotations) == 1 + annotation = content.annotations[0] + assert annotation["type"] == "citation" + assert annotation["additional_properties"]["annotation_index"] == 0 + assert all(annotation.get(key) == value for key, value in expected.items()) + assert annotation["raw_representation"].type == provider_annotation["type"] + assert content.raw_representation is event + + +@pytest.mark.parametrize("mode", ["create", "retrieve", "structured"]) +async def test_streaming_completed_annotations_are_deduplicated_per_occurrence(mode: str) -> None: + """Added, completed-part, and completed-item events emit each citation once.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + first = { + "type": "url_citation", + "url": "https://example.sharepoint.com/document.pdf", + "title": "Document", + "start_index": 0, + "end_index": 3, + } + second = {**first, "start_index": 21, "end_index": 29} + text = '{"location": "Seattle", "weather": "Clear"}' if mode == "structured" else "The answer is in the document." + text_event = ResponseTextDeltaEvent.model_validate({ + "type": "response.output_text.delta", + "item_id": "msg_citations", + "output_index": 0, + "content_index": 0, + "sequence_number": 0, + "delta": text, + "logprobs": [], + }) + added_event = ResponseOutputTextAnnotationAddedEvent.model_validate({ + "type": "response.output_text.annotation.added", + "item_id": "msg_citations", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + "annotation_index": 0, + "annotation": first, + }) + done_event = _make_completed_text_event([first, second], text=text) + item_done_event = _make_completed_text_event([first, second], event_type="response.output_item.done", text=text) + options: OpenAIChatOptions[OutputStruct] = {} + if mode == "retrieve": + options["continuation_token"] = {"response_id": "resp_citations"} + elif mode == "structured": + options["response_format"] = OutputStruct + original_options = dict(options) + events = [text_event, added_event, done_event, item_done_event, done_event] + method = "stream" if mode == "structured" else mode + + with patch.object( + client.client.responses, + method, + new_callable=MagicMock if mode == "structured" else AsyncMock, + return_value=_FakeAsyncEventStream(events), + ) as create: + stream = client.get_response( + [Message(role="user", contents=["Find the document."])], stream=True, options=options + ) + updates = [update async for update in stream] + response = await stream.get_final_response() + + assert response.text == text_event.delta + assert len(response.messages) == 1 + assert len(response.messages[0].contents) == 1 + annotations = response.messages[0].contents[0].annotations + assert annotations is not None + assert len(annotations) == 2 + assert [annotation["annotated_regions"][0]["start_index"] for annotation in annotations] == [0, 21] + assert sum(len(content.annotations or []) for update in updates for content in update.contents) == 2 + assert options == original_options + assert not any(key.startswith("__agent_framework") for key in create.call_args.kwargs) + if mode == "structured": + assert isinstance(response.value, OutputStruct) + assert response.value.location == "Seattle" + + +@pytest.mark.parametrize("with_added_event", [False, True]) +async def test_completed_sharepoint_citations_reach_agui_stream(with_added_event: bool) -> None: + """Responses completion citations survive the client, Agent, and AG-UI transport.""" + pytest.importorskip("agent_framework_ag_ui") + from ag_ui.core import CustomEvent, TextMessageContentEvent + from agent_framework_ag_ui import AgentFrameworkAgent + + client = OpenAIChatClient(model="test-model", api_key="test-key") + annotation = { + "type": "url_citation", + "title": "Document", + "url": "https://example.sharepoint.com/document.pdf", + "start_index": 21, + "end_index": 29, + } + text_event = ResponseTextDeltaEvent.model_validate({ + "type": "response.output_text.delta", + "item_id": "msg_citations", + "output_index": 0, + "content_index": 0, + "sequence_number": 0, + "delta": "The answer is in the document.", + "logprobs": [], + }) + sdk_events: list[object] = [text_event] + if with_added_event: + sdk_events.append( + ResponseOutputTextAnnotationAddedEvent.model_validate({ + "type": "response.output_text.annotation.added", + "item_id": "msg_citations", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + "annotation_index": 0, + "annotation": annotation, + }) + ) + sdk_events.extend([ + _make_completed_text_event([annotation]), + _make_completed_text_event([annotation], event_type="response.output_item.done"), + ]) + agent = AgentFrameworkAgent(agent=client.as_agent(name="grounded-assistant")) + + with patch.object( + client.client.responses, "create", new_callable=AsyncMock, return_value=_FakeAsyncEventStream(sdk_events) + ): + events = [ + event + async for event in agent.run({ + "thread_id": "thread-citations", + "run_id": "run-citations", + "messages": [{"role": "user", "content": "Find the document."}], + }) + ] + + citation_events = [event for event in events if isinstance(event, CustomEvent) and event.name == "annotations"] + assert len(citation_events) == 1 + text_events = [event for event in events if isinstance(event, TextMessageContentEvent)] + assert "".join(event.delta for event in text_events) == text_event.delta + payload = json.loads(citation_events[0].model_dump_json(by_alias=True)) + assert payload["value"]["messageId"] == text_events[0].message_id + assert payload["value"]["annotations"] == [ + { + "type": "citation", + "title": annotation["title"], + "url": annotation["url"], + "annotated_regions": [{"type": "text_span", "start_index": 21, "end_index": 29}], + "additional_properties": {"annotation_index": 0}, + } + ] + assert events.index(citation_events[0]) < next( + index for index, event in enumerate(events) if event.type == "RUN_FINISHED" + ) + + +async def test_streaming_completed_annotations_do_not_share_state_between_requests() -> None: + """A shared client preserves identical citations in concurrent independent runs.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + annotation = {"type": "file_citation", "file_id": "file-123", "filename": "document.pdf", "index": 0} + events = [ + _make_completed_text_event([annotation, annotation]), + _make_completed_text_event([annotation], content_index=1), + _make_completed_text_event([annotation], item_id="msg_other"), + ] + + class InterleavedEventStream(_FakeAsyncEventStream): + async def __anext__(self) -> object: + await asyncio.sleep(0) + return await super().__anext__() + + async def get_response() -> ChatResponse: + return await client.get_response( + [Message(role="user", contents=["Find the document."])], stream=True + ).get_final_response() + + with patch.object( + client.client.responses, + "create", + new_callable=AsyncMock, + side_effect=[InterleavedEventStream(events), InterleavedEventStream(events)], + ): + responses = await asyncio.gather(get_response(), get_response()) + + for response in responses: + assert len(response.messages) == 1 + assert len(response.messages[0].contents) == 1 + annotations = response.messages[0].contents[0].annotations + assert annotations is not None + assert len(annotations) == 4 + + +def test_streaming_completed_text_without_annotations_emits_no_content() -> None: + """Completed text must not duplicate text that was already streamed.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + update = client._parse_chunk_from_openai(_make_completed_text_event([]), {}, {}) + + assert update.contents == [] + + def test_streaming_annotation_added_with_file_path() -> None: """Streaming `file_path` should attach as a text annotation, matching non-streaming.""" client = OpenAIChatClient(model="test-model", api_key="test-key") From b9b836c651799dd07aaaf9b8b100e8d4e190939d Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 2 Oct 2026 09:00:01 +0200 Subject: [PATCH 2/2] Python: address AG-UI citation review feedback --- .../ag-ui/agent_framework_ag_ui/_client.py | 12 ++-- .../_event_converters.py | 59 +++++++++++++++++++ .../agent_framework_ag_ui/_run_common.py | 22 +++++++ .../tests/ag_ui/test_event_converters.py | 35 ++++++++++- python/packages/ag-ui/tests/ag_ui/test_run.py | 49 +++++++++++++++ python/packages/openai/README.md | 3 - 6 files changed, 171 insertions(+), 9 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index f15fc367a8d..f47bb020333 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -30,7 +30,7 @@ from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer from agent_framework.observability import ChatTelemetryLayer -from ._event_converters import AGUIEventConverter +from ._event_converters import AGUIEventConverter, _finalize_agui_response from ._feature_usage import FeatureIndex from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume from ._message_adapters import agent_framework_messages_to_agui @@ -430,17 +430,19 @@ def _inner_get_response( options=options, **kwargs, ), - finalizer=ChatResponse.from_updates, + finalizer=_finalize_agui_response, ) async def _get_response() -> ChatResponse: - return await ChatResponse.from_update_generator( - self._streaming_impl( + updates = [ + update + async for update in self._streaming_impl( messages=messages, options=options, **kwargs, ) - ) + ] + return _finalize_agui_response(updates) return _get_response() diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py b/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py index a161f5b7e18..9d3dad2ee17 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py @@ -5,17 +5,76 @@ from __future__ import annotations import logging +from collections.abc import Sequence from typing import Any, cast from agent_framework import ( Annotation, + ChatResponse, ChatResponseUpdate, Content, + Message, ) logger = logging.getLogger(__name__) +def _annotation_batch_from_update(update: ChatResponseUpdate) -> tuple[str, list[Annotation]] | None: + """Extract a message-linked annotation batch produced by this converter.""" + custom_event = (update.additional_properties or {}).get("ag_ui_custom_event") + if not isinstance(custom_event, dict) or custom_event.get("name") != "annotations" or not update.message_id: + return None + annotations = [ + annotation for content in update.contents if content.type == "text" for annotation in content.annotations or [] + ] + return update.message_id, annotations + + +def _finalize_agui_response(updates: Sequence[ChatResponseUpdate]) -> ChatResponse: + """Aggregate AG-UI updates while attaching annotation events by message ID.""" + annotation_batches: list[tuple[str, list[Annotation]]] = [] + aggregatable_updates: list[ChatResponseUpdate] = [] + for update in updates: + annotation_batch = _annotation_batch_from_update(update) + if annotation_batch is None: + aggregatable_updates.append(update) + else: + annotation_batches.append(annotation_batch) + + response = ChatResponse.from_updates(aggregatable_updates) + + # Annotation events are excluded from message aggregation, but their response metadata + # and raw representations retain their original stream ordering. + response.additional_properties.clear() + for update in updates: + if update.additional_properties: + response.additional_properties.update(update.additional_properties) + if updates: + response.raw_representation = [update.raw_representation for update in updates] + response.continuation_token = updates[-1].continuation_token + + for message_id, batch_annotations in annotation_batches: + if not batch_annotations: + continue + message = next((message for message in response.messages if message.message_id == message_id), None) + if message is None: + response.messages.append( + Message( + role="assistant", + contents=[Content.from_text(text="", annotations=batch_annotations)], + message_id=message_id, + ) + ) + continue + text_content = next((content for content in message.contents if content.type == "text"), None) + if text_content is None: + message.contents.append(Content.from_text(text="", annotations=batch_annotations)) + else: + text_content.annotations = [*(text_content.annotations or []), *batch_annotations] + + return response + + class AGUIEventConverter: """Converter for AG-UI events to Agent Framework types. 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 b88a1c1d5fe..0f931ff88e2 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 @@ -609,6 +609,25 @@ def _track_reasoning_segment(flow: FlowState, message_id: str) -> None: flow.snapshot_segments.append({"kind": "reasoning", "id": message_id}) +def _offset_annotation_text_spans(annotations: list[dict[str, Any]], offset: int) -> None: + """Rebase text spans from a delta to the accumulated message.""" + if offset == 0: + return + for annotation in annotations: + regions = annotation.get("annotated_regions") + if not isinstance(regions, list): + continue + for region in regions: + if not isinstance(region, dict) or region.get("type") != "text_span": + continue + start_index = region.get("start_index") + end_index = region.get("end_index") + if isinstance(start_index, int): + region["start_index"] = start_index + offset + if isinstance(end_index, int): + region["end_index"] = end_index + offset + + def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]: """Emit text deltas and message-linked annotation batches.""" if not content.text and not content.annotations: @@ -651,6 +670,7 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li if segment is None: segment = _open_text_segment(flow, flow.message_id) + annotation_offset = len(flow.accumulated_text) if content.text and not duplicate_text: events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text)) flow.accumulated_text += content.text @@ -665,6 +685,8 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li ] ), ) + if content.text and not duplicate_text: + _offset_annotation_text_spans(annotations, annotation_offset) new_annotations = annotations if duplicate_text: # Reconcile replay by occurrence count, not source URL or dictionary uniqueness. diff --git a/python/packages/ag-ui/tests/ag_ui/test_event_converters.py b/python/packages/ag-ui/tests/ag_ui/test_event_converters.py index 2b9323a4bd7..bacee2e7aa4 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_event_converters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_event_converters.py @@ -8,7 +8,7 @@ import pytest from agent_framework import ChatResponse -from agent_framework_ag_ui._event_converters import AGUIEventConverter +from agent_framework_ag_ui._event_converters import AGUIEventConverter, _finalize_agui_response class TestAGUIEventConverter: @@ -421,6 +421,39 @@ def test_annotations_custom_event_restores_content(self, event_type: str) -> Non assert update.additional_properties is not None assert update.additional_properties["ag_ui_custom_event"]["value"] == event["value"] + def test_annotation_finalizer_targets_an_existing_message_after_later_events(self) -> None: + """Late annotations attach by message ID without creating a duplicate message.""" + converter = AGUIEventConverter() + annotations = [ + { + "type": "citation", + "url": "https://example.com/first", + "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 5}], + } + ] + events: list[dict[str, Any]] = [ + {"type": "TEXT_MESSAGE_START", "messageId": "m1"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "First"}, + {"type": "TEXT_MESSAGE_START", "messageId": "m2"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m2", "delta": "Second"}, + {"type": "TOOL_CALL_RESULT", "toolCallId": "call-1", "result": "done"}, + { + "type": "CUSTOM", + "name": "annotations", + "value": {"messageId": "m1", "annotations": annotations}, + }, + ] + updates = [update for event in events if (update := converter.convert_event(event)) is not None] + + response = _finalize_agui_response(updates) + + assert [message.message_id for message in response.messages].count("m1") == 1 + first_message = next(message for message in response.messages if message.message_id == "m1") + first_text = next(content for content in first_message.contents if content.type == "text") + assert first_text.text == "First" + assert first_text.annotations == annotations + assert [message.message_id for message in response.messages].count(None) == 1 + @pytest.mark.parametrize( "value", [ 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 f487f4a2ed3..7e0b483020e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -670,6 +670,55 @@ def test_emit_text_annotations_are_json_safe_and_message_linked() -> None: assert annotation["raw_representation"] is content.annotations[0]["raw_representation"] +def test_emit_text_rebases_delta_annotation_spans_and_deduplicates_full_replay() -> None: + """Chunk-relative spans become message-relative and match completed replay spans.""" + flow = FlowState() + annotation = Annotation( + type="citation", + url="https://example.com/target", + annotated_regions=[{"type": "text_span", "start_index": 0, "end_index": 6}], + ) + _emit_text(Content.from_text("Prefix "), flow) + + events = _emit_text(Content.from_text("target", annotations=[annotation]), flow) + replay_events = _emit_text( + Content.from_text( + "Prefix target", + annotations=[ + Annotation( + type="citation", + url="https://example.com/target", + annotated_regions=[{"type": "text_span", "start_index": 7, "end_index": 13}], + ) + ], + ), + flow, + ) + + citation_event = next(event for event in events if isinstance(event, CustomEvent)) + assert citation_event.value["annotations"][0]["annotated_regions"] == [ + {"type": "text_span", "start_index": 7, "end_index": 13} + ] + assert annotation["annotated_regions"] == [{"type": "text_span", "start_index": 0, "end_index": 6}] + assert replay_events == [] + + +def test_emit_text_preserves_annotation_only_spans() -> None: + """Annotation-only updates already use complete-message span indices.""" + flow = FlowState() + annotation = Annotation( + type="citation", + url="https://example.com/target", + annotated_regions=[{"type": "text_span", "start_index": 7, "end_index": 13}], + ) + _emit_text(Content.from_text("Prefix target"), flow) + + events = _emit_text(Content.from_text("", annotations=[annotation]), flow) + + citation_event = next(event for event in events if isinstance(event, CustomEvent)) + assert citation_event.value["annotations"][0]["annotated_regions"] == annotation["annotated_regions"] + + def test_emit_text_annotation_only_update_starts_a_message_without_an_empty_delta() -> None: """Citations received before text still reference an announced message.""" flow = FlowState() diff --git a/python/packages/openai/README.md b/python/packages/openai/README.md index 5be069741a6..9198b8f1011 100644 --- a/python/packages/openai/README.md +++ b/python/packages/openai/README.md @@ -36,9 +36,6 @@ without repeating text that was already streamed. Annotation-added and completio by item, content part, and annotation position, so separate references to the same source remain distinct. Foundry clients that inherit the Responses parser receive the same behavior. -The AG-UI integration surfaces these annotations as message-linked `CUSTOM` events named `annotations`; -see [AG-UI citation support](../ag-ui/README.md#citations-and-annotations). No logging handler is needed. - ## Hosted function results `OpenAIChatClient` parses hosted `function_call_output` items in both streaming and non-streaming responses.