diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 4119cdafd0..e71a62a6eb 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -775,12 +775,28 @@ def _convert_workflow_event_to_agent_response_updates( raw_representation=msg, ) ) + if not updates and data.usage_details is not None: + # No message is shown to the caller, but the tokens were still spent. Without a + # message_id or author_name this update joins the message in progress instead of + # starting a new one. + updates.append( + AgentResponseUpdate( + contents=[], + role="assistant", + response_id=data.response_id or response_id, + created_at=data.created_at + or datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + ) + ) if updates: updates[-1].agent_id = data.agent_id updates[-1].finish_reason = data.finish_reason updates[-1].continuation_token = data.continuation_token _copy_operation_state(cast(AgentResponse[Any], data), updates[-1]) updates[-1].additional_properties = dict(data.additional_properties) + # Updates carry usage as content; without it the streamed response loses the token counts. + if data.usage_details is not None: + updates[-1].contents.append(Content.from_usage(data.usage_details)) return updates if isinstance(data, Message): if data.role != "assistant" or _contains_internal_call_content(data.contents): diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 020d8a7326..93a1db2a80 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -943,6 +943,153 @@ async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, assert final_response.continuation_token == {"token": "response-resume-token"} assert final_response.additional_properties == {"provider_marker": "preserve-response"} + async def test_workflow_as_agent_stream_preserves_agent_response_usage(self) -> None: + """Streaming must carry the usage of an AgentResponse output, like the non-streaming path does.""" + usage = UsageDetails(input_token_count=10, output_token_count=5, total_token_count=15) + + @executor + async def usage_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output( + AgentResponse( + messages=[ + Message(role="assistant", contents=["first"]), + Message(role="assistant", contents=["second"]), + ], + response_id="usage-response", + usage_details=usage, + ) + ) + + agent = WorkflowBuilder(start_executor=usage_executor).build().as_agent("usage-test-agent") + + non_streaming = await agent.run("hello") + assert non_streaming.usage_details == usage + + stream = agent.run("hello", stream=True) + updates = [update async for update in stream] + final_response = await stream.get_final_response() + + assert [update.text for update in updates] == ["first", "second"] + # Usage describes the whole response, so it rides on the last update only. + assert not any(content.type == "usage" for content in updates[0].contents) + usage_contents = [content for content in updates[-1].contents if content.type == "usage"] + assert len(usage_contents) == 1 + assert usage_contents[0].usage_details == usage + assert final_response.usage_details == usage + assert final_response.text == non_streaming.text + + async def test_workflow_as_agent_stream_sums_usage_of_multiple_agent_responses(self) -> None: + """Usage from several AgentResponse outputs adds up the same way in both run modes.""" + + @executor + async def two_responses_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output( + AgentResponse( + messages=[Message(role="assistant", contents=["one"])], + response_id="usage-response-1", + usage_details=UsageDetails(input_token_count=3, output_token_count=2, total_token_count=5), + ) + ) + await ctx.yield_output( + AgentResponse( + messages=[Message(role="assistant", contents=["two"])], + response_id="usage-response-2", + usage_details=UsageDetails(input_token_count=7, output_token_count=1, total_token_count=8), + ) + ) + + agent = WorkflowBuilder(start_executor=two_responses_executor).build().as_agent("usage-sum-agent") + expected = UsageDetails(input_token_count=10, output_token_count=3, total_token_count=13) + + non_streaming = await agent.run("hello") + streaming = await agent.run("hello", stream=True).get_final_response() + + assert non_streaming.usage_details == expected + assert streaming.usage_details == expected + + @pytest.mark.parametrize( + "messages", + [ + [Message(role="assistant", contents=[Content.from_function_call(call_id="c1", name="t", arguments="{}")])], + [], + ], + ids=["only_internal_call", "no_messages"], + ) + async def test_workflow_as_agent_stream_keeps_usage_without_caller_facing_messages( + self, messages: list[Message] + ) -> None: + """Usage still reaches the streamed response when no message of the AgentResponse is shown to the caller.""" + usage = UsageDetails(input_token_count=10, output_token_count=5, total_token_count=15) + + @executor + async def usage_only_executor(messages_in: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output(AgentResponse(messages=list(messages), usage_details=usage)) + + agent = WorkflowBuilder(start_executor=usage_only_executor).build().as_agent("usage-only-agent") + + non_streaming = await agent.run("hello") + stream = agent.run("hello", stream=True) + updates = [update async for update in stream] + streaming = await stream.get_final_response() + + assert non_streaming.usage_details == usage + assert streaming.usage_details == usage + assert streaming.text == non_streaming.text + assert len(streaming.messages) == len(non_streaming.messages) + # The filtered messages must not reach the caller through a raw representation either. + assert all(update.raw_representation is None for update in updates) + raw = streaming.raw_representation + assert not any(isinstance(item, AgentResponse) for item in (raw if isinstance(raw, list) else [raw])) + + async def test_workflow_as_agent_stream_usage_only_update_does_not_split_a_message(self) -> None: + """A usage-only AgentResponse between chunks of one streamed message leaves that message whole.""" + usage = UsageDetails(input_token_count=2, output_token_count=1, total_token_count=3) + + @executor + async def interleaved_executor( + messages: list[Message], + ctx: WorkflowContext[Never, AgentResponseUpdate | AgentResponse], # type: ignore[valid-type] + ) -> None: + await ctx.yield_output( + AgentResponseUpdate(contents=[Content.from_text("x")], message_id="m1", author_name="writer") + ) + await ctx.yield_output(AgentResponse(messages=[], usage_details=usage)) + await ctx.yield_output( + AgentResponseUpdate(contents=[Content.from_text("y")], message_id="m1", author_name="writer") + ) + + agent = WorkflowBuilder(start_executor=interleaved_executor).build().as_agent("interleaved-agent") + + streaming = await agent.run("hello", stream=True).get_final_response() + + assert [message.text for message in streaming.messages] == ["xy"] + assert streaming.usage_details == usage + + async def test_workflow_as_agent_stream_usage_only_update_keeps_continuation_token(self) -> None: + """The usage-only update carries the response-level fields, like the last update of a normal response.""" + usage = UsageDetails(input_token_count=4, output_token_count=0, total_token_count=4) + + @executor + async def background_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output( + AgentResponse( + messages=[], + usage_details=usage, + continuation_token=cast(Any, {"token": "still-running"}), + ) + ) + + agent = WorkflowBuilder(start_executor=background_executor).build().as_agent("background-agent") + + stream = agent.run("hello", stream=True) + updates = [update async for update in stream] + final_response = await stream.get_final_response() + + assert len(updates) == 1 + assert [content.type for content in updates[0].contents] == ["usage"] + assert updates[0].continuation_token == {"token": "still-running"} + assert final_response.usage_details == usage + async def test_workflow_as_agent_stream_preserves_custom_finish_reason(self) -> None: """Test that a non-literal finish_reason is forwarded unchanged.