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 f47bb020333..37a51c1c374 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -17,6 +17,8 @@ import httpx from agent_framework import ( + AgentResponse, + AgentResponseUpdate, BaseChatClient, ChatResponse, ChatResponseUpdate, @@ -30,7 +32,7 @@ from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer from agent_framework.observability import ChatTelemetryLayer -from ._event_converters import AGUIEventConverter, _finalize_agui_response +from ._event_converters import AGUIEventConverter, _finalize_agui_response, _finalize_agui_updates from ._feature_usage import FeatureIndex from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume from ._message_adapters import agent_framework_messages_to_agui @@ -403,6 +405,28 @@ def _get_thread_id(self, options: Mapping[str, Any]) -> str: return thread_id + @override + def _finalize_function_invocation_updates( + self, + updates: Sequence[ChatResponseUpdate], + *, + response_format: Any | None = None, + ) -> ChatResponse[Any]: + """Keep message-linked annotations attached across the complete tool loop.""" + return _finalize_agui_response(updates, response_format=response_format) + + @override + def _finalize_agent_response_updates( + self, + updates: Sequence[AgentResponseUpdate], + *, + response_format: Any | None = None, + ) -> AgentResponse[Any]: + """Preserve message-linked annotations after chat updates are mapped to an agent.""" + return _finalize_agui_updates( + updates, lambda items: AgentResponse.from_updates(items, output_format_type=response_format) + ) + @override def _inner_get_response( self, 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 9d3dad2ee17..c1f3e6bf377 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,10 +5,13 @@ from __future__ import annotations import logging -from collections.abc import Sequence -from typing import Any, cast +from collections.abc import Callable, Sequence +from copy import copy +from typing import Any, TypeVar, cast from agent_framework import ( + AgentResponse, + AgentResponseUpdate, Annotation, ChatResponse, ChatResponseUpdate, @@ -16,32 +19,60 @@ Message, ) +AGUIUpdateT = TypeVar("AGUIUpdateT", ChatResponseUpdate, AgentResponseUpdate) +AGUIResponseT = TypeVar("AGUIResponseT", bound=ChatResponse[Any] | AgentResponse[Any]) + logger = logging.getLogger(__name__) -def _annotation_batch_from_update(update: ChatResponseUpdate) -> tuple[str, list[Annotation]] | None: +def _annotation_batch_from_update( + update: ChatResponseUpdate | AgentResponseUpdate, +) -> 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 + # Buffered middleware can copy this response-level event metadata onto a rebuilt message. + # Only the original annotation-only update should be excluded from message aggregation. + value = custom_event.get("value") + if not isinstance(value, dict) or value.get("messageId") != update.message_id: + return None + if any(content.type != "text" or content.text for content in update.contents): + 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.""" +def _finalize_agui_response( + updates: Sequence[ChatResponseUpdate], *, response_format: Any | None = None +) -> ChatResponse[Any]: + """Aggregate chat updates while attaching annotation events by message ID.""" + return _finalize_agui_updates( + updates, lambda items: ChatResponse.from_updates(items, output_format_type=response_format) + ) + + +def _finalize_agui_updates( + updates: Sequence[AGUIUpdateT], aggregate: Callable[[Sequence[AGUIUpdateT]], AGUIResponseT] +) -> AGUIResponseT: + """Attach citations while preserving the response type and original update metadata.""" annotation_batches: list[tuple[str, list[Annotation]]] = [] - aggregatable_updates: list[ChatResponseUpdate] = [] + aggregatable_updates: list[AGUIUpdateT] = [] 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) + # Keep response-level fields from the event without changing any message. + metadata_update = copy(update) + metadata_update.contents = [] + metadata_update.message_id = metadata_update.role = metadata_update.author_name = None + aggregatable_updates.append(metadata_update) - response = ChatResponse.from_updates(aggregatable_updates) + response = aggregate(aggregatable_updates) # Annotation events are excluded from message aggregation, but their response metadata # and raw representations retain their original stream ordering. diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index 679cfb11eec..04fe06fd953 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -12,6 +12,7 @@ import pytest from ag_ui.core import Interrupt, ResumeEntry from agent_framework import ( + Agent, ChatOptions, ChatResponse, ChatResponseUpdate, @@ -727,6 +728,288 @@ async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str assert len(response.messages) > 0 assert "Complete response" in response.text + @pytest.mark.parametrize("stream,invocation_enabled", [(True, True), (False, True), (True, False)]) + @pytest.mark.parametrize("late_annotations", [False, True], ids=["before_later_message", "after_later_message"]) + @pytest.mark.parametrize("as_agent", [False, True], ids=["client", "agent"]) + async def test_public_response_attaches_annotations_to_the_referenced_message( + self, stream: bool, invocation_enabled: bool, late_annotations: bool, as_agent: bool + ) -> None: + """Public finalization must attach annotations by message ID after later text arrives.""" + annotations = [ + { + "type": "citation", + "url": "https://example.test/source", + "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 5}], + } + ] + annotation_event = { + "type": "CUSTOM", + "name": "annotations", + "value": {"messageId": "m1", "annotations": annotations}, + } + + async def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(await request.aread()) + events: list[dict[str, Any]] = [ + {"type": "RUN_STARTED", "threadId": payload["thread_id"], "runId": payload["run_id"]}, + {"type": "TEXT_MESSAGE_START", "messageId": "m1", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "First"}, + {"type": "TEXT_MESSAGE_END", "messageId": "m1"}, + ] + if not late_annotations: + events.append(annotation_event) + events.extend( + [ + {"type": "TEXT_MESSAGE_START", "messageId": "m2", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m2", "delta": "Second"}, + {"type": "TEXT_MESSAGE_END", "messageId": "m2"}, + ] + ) + if late_annotations: + events.append(annotation_event) + events.append({"type": "RUN_FINISHED", "threadId": payload["thread_id"], "runId": payload["run_id"]}) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + async with AGUIChatClient( + endpoint="https://agui.example.test/", + http_client=http_client, + function_invocation_configuration={"enabled": invocation_enabled}, + ) as client: + get_response = Agent(client=client).run if as_agent else client.get_response + messages = [Message(role="user", contents=["Question"])] + if stream: + response_stream = get_response(messages, stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + annotation_updates = [ + update + for update in updates + if (update.additional_properties or {}).get("ag_ui_custom_event", {}).get("name") + == "annotations" + ] + assert len(annotation_updates) == 1 + assert "".join(update.text for update in updates) == "FirstSecond" + else: + response = await get_response(messages) + + assert [(message.message_id, message.text) for message in response.messages] == [ + ("m1", "First"), + ("m2", "Second"), + ] + assert response.messages[0].contents[0].annotations == annotations + assert response.messages[1].contents[0].annotations is None + assert response.additional_properties["ag_ui_custom_event"]["value"] == annotation_event["value"] + + @pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) + @pytest.mark.parametrize("as_agent", [False, True], ids=["client", "agent"]) + async def test_public_annotations_preserve_a_completed_client_tool_round( + self, stream: bool, as_agent: bool + ) -> None: + """Late annotations must retain one executed call/result pair and the next request's history.""" + executed: list[str] = [] + requests: list[dict[str, Any]] = [] + annotations = [{"type": "citation", "url": "https://example.test/source"}] + + @tool(approval_mode="never_require") + def echo(value: str) -> str: + """Return the requested value.""" + executed.append(value) + return value + + async def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(await request.aread()) + requests.append(payload) + events: list[dict[str, Any]] = [ + {"type": "RUN_STARTED", "threadId": payload["thread_id"], "runId": payload["run_id"]}, + ] + if len(requests) == 1: + events.extend( + [ + {"type": "TEXT_MESSAGE_START", "messageId": "planning", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "planning", "delta": "Checking"}, + {"type": "TEXT_MESSAGE_END", "messageId": "planning"}, + {"type": "TOOL_CALL_START", "toolCallId": "call-echo", "toolName": "echo"}, + {"type": "TOOL_CALL_ARGS", "toolCallId": "call-echo", "delta": '{"value":"done"}'}, + {"type": "TOOL_CALL_END", "toolCallId": "call-echo"}, + ] + ) + else: + annotated_events: list[dict[str, Any]] = [ + {"type": "TEXT_MESSAGE_START", "messageId": "m1", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "First"}, + {"type": "TEXT_MESSAGE_END", "messageId": "m1"}, + {"type": "TEXT_MESSAGE_START", "messageId": "m2", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m2", "delta": "Second"}, + {"type": "TEXT_MESSAGE_END", "messageId": "m2"}, + { + "type": "CUSTOM", + "name": "annotations", + "value": {"messageId": "m1", "annotations": annotations}, + }, + ] + events.extend(annotated_events) + events.append({"type": "RUN_FINISHED", "threadId": payload["thread_id"], "runId": payload["run_id"]}) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + async with AGUIChatClient(endpoint="https://agui.example.test/", http_client=http_client) as client: + get_response = Agent(client=client).run if as_agent else client.get_response + messages = [Message(role="user", contents=["Question"])] + if stream: + response_stream = get_response( + messages, stream=True, options={"tools": [echo], "metadata": {"thread_id": "thread-1"}} + ) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + assert ( + sum(content.type == "function_result" for update in updates for content in update.contents) == 1 + ) + else: + response = await get_response( + messages, options={"tools": [echo], "metadata": {"thread_id": "thread-1"}} + ) + + assert len(requests) == 2 + assert [request["thread_id"] for request in requests] == ["thread-1", "thread-1"] + assert executed == ["done"] + assert len(messages) == 1 + call_contents = [ + content for message in response.messages for content in message.contents if content.type == "function_call" + ] + result_contents = [ + content + for message in response.messages + for content in message.contents + if content.type == "function_result" + ] + assert len(call_contents) == len(result_contents) == 1 + assert call_contents[0].call_id == result_contents[0].call_id == "call-echo" + text_messages = [(message.message_id, message.text) for message in response.messages if message.text] + assert text_messages == [("planning", "Checking"), ("m1", "First"), ("m2", "Second")] + first = next(message for message in response.messages if message.message_id == "m1") + assert first.contents[0].annotations == annotations + assert sum(message.message_id == "m1" for message in response.messages) == 1 + history = requests[1]["messages"] + results = [message for message in history if message.get("role") == "tool"] + assert len(results) == 1 + assert results[0]["toolCallId"] == "call-echo" + assert results[0]["content"] == "done" + + @pytest.mark.parametrize("as_agent", [False, True], ids=["client", "agent"]) + async def test_public_annotation_stream_preserves_structured_output( + self, monkeypatch: MonkeyPatch, as_agent: bool + ) -> None: + """The AG-UI whole-loop finalizer retains the requested structured response type.""" + from pydantic import BaseModel + + class Answer(BaseModel): + reply: str + + annotations = [{"type": "citation", "url": "https://example.test/source"}] + events: list[dict[str, Any]] = [ + {"type": "TEXT_MESSAGE_START", "messageId": "m1", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": '{"reply":"First"}'}, + {"type": "TEXT_MESSAGE_END", "messageId": "m1"}, + {"type": "TEXT_MESSAGE_START", "messageId": "m2", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m2", "delta": '{"reply":"Second"}'}, + {"type": "TEXT_MESSAGE_END", "messageId": "m2"}, + {"type": "CUSTOM", "name": "annotations", "value": {"messageId": "m1", "annotations": annotations}}, + ] + + async def post_run(**kwargs: Any) -> AsyncGenerator[dict[str, Any], None]: + for event in events: + yield event + + async with StubAGUIChatClient(endpoint="https://agui.example.test/") as client: + get_response = Agent(client=client).run if as_agent else client.get_response + monkeypatch.setattr(client.http_service, "post_run", post_run) + messages: list[Message] = [Message(role="user", contents=["Question"])] + response_stream = get_response(messages, stream=True, options={"response_format": Answer}) + async for _ in response_stream: + pass + response = await response_stream.get_final_response() + + assert [message.message_id for message in response.messages] == ["m1", "m2"] + assert response.messages[0].contents[0].annotations == annotations + assert isinstance(response.value, Answer) + assert response.value.reply == "Second" + + @pytest.mark.parametrize("message_count", [1, 2], ids=["same_message", "later_message"]) + @pytest.mark.parametrize("batch_count", [1, 2], ids=["one_batch", "two_batches"]) + @pytest.mark.parametrize("as_agent", [False, True], ids=["client", "agent"]) + async def test_buffered_middleware_preserves_rebuilt_text_and_annotations( + self, monkeypatch: MonkeyPatch, as_agent: bool, message_count: int, batch_count: int + ) -> None: + """Response-level annotation metadata must not hide text rebuilt by middleware.""" + from collections.abc import Callable + + from agent_framework import ChatContext, ChatMiddleware + + class UppercaseMiddleware(ChatMiddleware): + async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + context.stream_buffer_updates = True + context.stream_result_to_updates = ChatResponse.to_updates + + def uppercase(response: ChatResponse[Any]) -> ChatResponse[Any]: + for message in response.messages: + for content in message.contents: + if content.type == "text" and content.text is not None: + content.text = content.text.upper() + return response + + context.stream_result_transforms.append(uppercase) + + annotations = [ + {"type": "citation", "url": f"https://example.test/source-{index}"} for index in range(batch_count) + ] + events: list[dict[str, Any]] = [ + {"type": "TEXT_MESSAGE_START", "messageId": "m1", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "First"}, + {"type": "TEXT_MESSAGE_END", "messageId": "m1"}, + ] + if message_count == 2: + events.extend( + [ + {"type": "TEXT_MESSAGE_START", "messageId": "m2", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m2", "delta": "Second"}, + {"type": "TEXT_MESSAGE_END", "messageId": "m2"}, + ] + ) + events.extend( + {"type": "CUSTOM", "name": "annotations", "value": {"messageId": "m1", "annotations": [annotation]}} + for annotation in annotations + ) + + async def post_run(**kwargs: Any) -> AsyncGenerator[dict[str, Any], None]: + for event in events: + yield event + + async with StubAGUIChatClient( + endpoint="https://agui.example.test/", middleware=[UppercaseMiddleware()] + ) as client: + get_response = Agent(client=client).run if as_agent else client.get_response + monkeypatch.setattr(client.http_service, "post_run", post_run) + messages: list[Message] = [Message(role="user", contents=["Question"])] + response_stream = get_response(messages, stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + + expected = [("m1", "FIRST")] + ([("m2", "SECOND")] if message_count == 2 else []) + assert [(message.message_id, message.text) for message in response.messages] == expected + assert "".join(update.text for update in updates) == "".join(text for _, text in expected) + assert response.messages[0].contents[0].annotations == annotations + assert response.additional_properties["ag_ui_custom_event"]["value"]["messageId"] == "m1" + async def test_tool_handling(self, monkeypatch: MonkeyPatch) -> None: """Test that client tool metadata is sent to server. 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 bacee2e7aa4..cb4db69de51 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 @@ -6,9 +6,9 @@ from typing import Any, cast import pytest -from agent_framework import ChatResponse +from agent_framework import Annotation, ChatResponse, Content, Message -from agent_framework_ag_ui._event_converters import AGUIEventConverter, _finalize_agui_response +from agent_framework_ag_ui._event_converters import AGUIEventConverter, _finalize_agui_response, _finalize_agui_updates class TestAGUIEventConverter: @@ -480,6 +480,91 @@ def test_malformed_annotations_custom_event_preserves_metadata( assert update.additional_properties["ag_ui_custom_event"]["value"] == value assert "annotations" in caplog.text + def test_finalizer_preserves_rebuilt_empty_message_with_response_event_metadata(self) -> None: + """Metadata copied to a different message is not a message-linked annotation event.""" + annotations = [Annotation(type="citation", url="https://example.test/source")] + original = ChatResponse( + messages=[ + Message( + role="assistant", message_id="m1", contents=[Content.from_text("First", annotations=annotations)] + ), + Message(role="assistant", message_id="m2", contents=[Content.from_text("")]), + ], + additional_properties={ + "ag_ui_custom_event": {"name": "annotations", "value": {"messageId": "m1", "annotations": annotations}} + }, + ) + + response = _finalize_agui_response(original.to_updates()) + + assert [(message.message_id, message.text) for message in response.messages] == [("m1", "First"), ("m2", "")] + assert response.messages[0].contents[0].annotations == annotations + assert response.additional_properties == original.additional_properties + + def test_agent_finalizer_preserves_metadata_from_annotation_only_update(self) -> None: + """Mapped annotation events retain agent-only fields without creating an empty duplicate.""" + from agent_framework import AgentResponse, AgentResponseUpdate, ContinuationToken + from pydantic import BaseModel + + class CursorToken(ContinuationToken): + cursor: str + + class Answer(BaseModel): + reply: str + + token = CursorToken(cursor="next") + annotations = [Annotation(type="citation", url="https://example.test/source")] + updates = [ + AgentResponseUpdate( + role="assistant", + message_id="m1", + contents=[Content.from_text("First")], + raw_representation={"index": 0}, + ), + AgentResponseUpdate( + role="assistant", + message_id="m2", + contents=[Content.from_text('{"reply":"Second"}')], + raw_representation={"index": 1}, + ), + AgentResponseUpdate( + role="assistant", + message_id="m1", + contents=[Content.from_text("", annotations=annotations)], + agent_id="agent-2", + response_id="response-2", + finish_reason="stop", + continuation_token=token, + additional_properties={ + "ag_ui_custom_event": { + "name": "annotations", + "value": {"messageId": "m1", "annotations": annotations}, + } + }, + raw_representation={"index": 2}, + ), + ] + + response = _finalize_agui_updates( + updates, lambda items: AgentResponse.from_updates(items, output_format_type=Answer) + ) + + assert [(message.message_id, message.text) for message in response.messages] == [ + ("m1", "First"), + ("m2", '{"reply":"Second"}'), + ] + assert response.messages[0].contents[0].annotations == annotations + assert response.agent_id == "agent-2" + assert response.response_id == "response-2" + assert response.finish_reason == "stop" + assert response.continuation_token == token + assert response.raw_representation == [{"index": 0}, {"index": 1}, {"index": 2}] + assert response.additional_properties == updates[-1].additional_properties + assert isinstance(response.value, Answer) + assert response.value.reply == "Second" + assert updates[-1].message_id == "m1" + assert updates[-1].contents[0].annotations == annotations + def test_full_conversation_flow(self) -> None: """Test complete conversation flow with multiple event types.""" converter = AGUIEventConverter() diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index c37156a779b..0560e17baff 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1455,6 +1455,10 @@ def _finalize_response_updates( response_format: Any | None = None, ) -> AgentResponse[Any]: """Finalize response updates into a single AgentResponse.""" + if isinstance(self.client, BaseChatClient): + return self.client._finalize_agent_response_updates( # pyright: ignore[reportPrivateUsage] + updates, response_format=response_format + ) return AgentResponse.from_updates( updates, output_format_type=response_format, diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 564d78aa660..8019378cdfd 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -28,6 +28,8 @@ from ._docstrings import apply_layered_docstring from ._serialization import SerializationMixin from ._types import ( + AgentResponse, + AgentResponseUpdate, ChatResponse, ChatResponseUpdate, EmbeddingGenerationOptions, @@ -355,6 +357,15 @@ def _finalize_response_updates( output_format_type=response_format, ) + def _finalize_agent_response_updates( + self, + updates: Sequence[AgentResponseUpdate], + *, + response_format: Any | None = None, + ) -> AgentResponse[Any]: + """Aggregate mapped agent updates without reapplying per-model-turn processing.""" + return AgentResponse.from_updates(updates, output_format_type=response_format) + def _build_response_stream( self, stream: AsyncIterable[ChatResponseUpdate] | Awaitable[AsyncIterable[ChatResponseUpdate]], diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 602da9dc4cb..d0a456fd7df 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -4910,6 +4910,17 @@ def __init__( kwargs["middleware"] = chat_middleware super().__init__(**kwargs) + def _finalize_function_invocation_updates( + self, + updates: Sequence[ChatResponseUpdate], + *, + response_format: Any | None = None, + ) -> ChatResponse[Any]: + """Aggregate all model turns and tool results without reapplying per-turn processing.""" + from ._types import ChatResponse + + return ChatResponse.from_updates(updates, output_format_type=response_format) + def _update_function_invocation_continuation_state( self, kwargs: dict[str, Any], @@ -5582,10 +5593,7 @@ def get_response( client_kwargs: Mapping[str, Any] | None = None, ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: from ._middleware import _as_middleware_list, categorize_middleware # pyright: ignore[reportPrivateUsage] - from ._types import ( - ChatResponse, - ResponseStream, - ) + from ._types import ResponseStream super_get_response = cast( Callable[..., Any], @@ -5708,7 +5716,7 @@ def get_response( def finalize_stream(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: if invalidation_error: raise invalidation_error[0] - return ChatResponse.from_updates(updates, output_format_type=response_format) + return self._finalize_function_invocation_updates(updates, response_format=response_format) return ResponseStream( self._stream_response_with_function_invocation( diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index e0e79153288..6da24b1fb5c 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -4574,3 +4574,79 @@ async def test_psc_flag_on_storing_without_conversation_id_warns_every_call( # Two service calls -> the warning is emitted twice (one per call, not deduped). missing_id_warnings = [r for r in caplog.records if "returned no conversation_id" in r.message] assert len(missing_id_warnings) == 2 + + +@pytest.mark.parametrize("base_client", [False, True], ids=["non_base_client", "base_client"]) +@pytest.mark.parametrize("custom_agent", [False, True], ids=["default_agent", "custom_agent"]) +async def test_stream_finalization_preserves_client_and_agent_extension_points( + client: Any, base_client: bool, custom_agent: bool +) -> None: + """Client delegation preserves the protocol fallback and an agent's existing override.""" + from pydantic import BaseModel + + from agent_framework import ContinuationToken, UsageDetails + + class CursorToken(ContinuationToken): + cursor: str + + class Answer(BaseModel): + reply: str + + calls: list[str] = [] + + class TrackingClient(MockBaseChatClient): + def _finalize_agent_response_updates( + self, updates: Sequence[AgentResponseUpdate], *, response_format: Any | None = None + ) -> AgentResponse[Any]: + calls.append("client") + return super()._finalize_agent_response_updates(updates, response_format=response_format) + + class TrackingAgent(Agent): + def _finalize_response_updates( + self, updates: Sequence[AgentResponseUpdate], *, response_format: Any | None = None + ) -> AgentResponse[Any]: + calls.append("agent") + return AgentResponse.from_updates(updates, output_format_type=response_format) + + chat_client = TrackingClient() if base_client else client + token = CursorToken(cursor="next") + chat_client.streaming_responses = [ + [ + ChatResponseUpdate( + role="assistant", + message_id="m1", + contents=[Content.from_text("First"), Content.from_usage(UsageDetails(input_token_count=3))], + ), + ChatResponseUpdate( + role="assistant", + message_id="m2", + response_id="response-2", + finish_reason="stop", + contents=[ + Content.from_text('{"reply":"Second"}'), + Content.from_usage(UsageDetails(output_token_count=2)), + ], + continuation_token=token, + additional_properties={"custom": "kept"}, + ), + ] + ] + agent = (TrackingAgent if custom_agent else Agent)(client=chat_client, name="example-agent") + response_stream = agent.run("Question", stream=True, options={"response_format": Answer}) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + + assert calls == (["agent"] if custom_agent else ["client"] if base_client else []) + assert len(updates) == 2 + assert [(message.message_id, message.text) for message in response.messages] == [ + ("m1", "First"), + ("m2", '{"reply":"Second"}'), + ] + assert all(message.author_name == "example-agent" for message in response.messages) + assert response.response_id == "response-2" + assert response.finish_reason == "stop" + assert response.continuation_token == token + assert response.usage_details == {"input_token_count": 3, "output_token_count": 2} + assert response.additional_properties["custom"] == "kept" + assert isinstance(response.value, Answer) + assert response.value.reply == "Second" diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 243dd0dcff0..889a07758b5 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -11194,3 +11194,79 @@ async def test_unbound_local_approval_response_is_filtered_before_mixed_batch_va assert executed == [] # The warning proves the response was filtered before the completeness check reported it missing. assert any("tool-approval response" in record.message for record in caplog.records) + + +async def test_whole_loop_finalization_preserves_per_turn_metadata_and_response_format(client: Any) -> None: + """Whole-loop aggregation keeps token metadata local and supports clients without BaseChatClient.""" + from agent_framework import ContinuationToken, UsageDetails + + class CursorToken(ContinuationToken): + cursor: str + + class Answer(BaseModel): + reply: str + + token = CursorToken(cursor="next") + executed: list[str] = [] + + @tool(approval_mode="never_require") + def echo(value: str) -> str: + """Return the requested value.""" + executed.append(value) + return value + + first_logprobs = [{"token": "First", "logprob": -0.1}] + last_logprobs = [{"token": "Second", "logprob": -0.2}] + client.streaming_responses = [ + [ + ChatResponseUpdate( + role="assistant", + message_id="first-message", + contents=[Content.from_text("First", additional_properties={"logprobs": first_logprobs})], + ), + ChatResponseUpdate( + role="assistant", + message_id="first-message", + contents=[Content.from_function_call(call_id="call-1", name="echo", arguments='{"value":"done"}')], + ), + ChatResponseUpdate(contents=[Content.from_usage(UsageDetails(input_token_count=3, output_token_count=1))]), + ], + [ + ChatResponseUpdate( + role="assistant", + message_id="last-message", + response_id="response-2", + model="test-model", + contents=[ + Content.from_text('{"reply":"Second"}', additional_properties={"logprobs": last_logprobs}), + Content.from_usage(UsageDetails(input_token_count=5, output_token_count=2)), + ], + continuation_token=token, + finish_reason="stop", + additional_properties={"custom": "kept"}, + ), + ], + ] + messages = [Message(role="user", contents=["Question"])] + response_stream = client.get_response(messages, stream=True, options={"tools": [echo], "response_format": Answer}) + assert isinstance(response_stream, ResponseStream) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + + assert executed == ["done"] + assert client.call_count == 2 + assert messages[0].text == "Question" + assert len(messages) == 1 + texts = [content for message in response.messages for content in message.contents if content.type == "text"] + assert [content.additional_properties["logprobs"] for content in texts] == [first_logprobs, last_logprobs] + assert sum(content.type == "function_call" for message in response.messages for content in message.contents) == 1 + assert sum(content.type == "function_result" for message in response.messages for content in message.contents) == 1 + assert response.response_id == "response-2" + assert response.model == "test-model" + assert response.finish_reason == "stop" + assert response.continuation_token == token + assert response.additional_properties["custom"] == "kept" + assert response.usage_details == {"input_token_count": 8, "output_token_count": 3} + assert isinstance(response.value, Answer) + assert response.value.reply == "Second" + assert len(updates) > 0