diff --git a/py/src/braintrust/integrations/agno/_test_agno_helpers.py b/py/src/braintrust/integrations/agno/_test_agno_helpers.py index 2c7d4b65b..f27584a35 100644 --- a/py/src/braintrust/integrations/agno/_test_agno_helpers.py +++ b/py/src/braintrust/integrations/agno/_test_agno_helpers.py @@ -62,6 +62,7 @@ async def _aexecute(self, session_id, user_id, execution_input, workflow_run_res def _execute_stream(self, session, execution_input, workflow_run_response, run_context=None): yield FakeEvent("WorkflowStarted", content=None) yield FakeEvent("StepStarted", content=None) + yield FakeEvent("StepProgress", content="fetched 3/10", data={"completed": 3, "total": 10}) yield FakeEvent("StepCompleted", content="hello ") yield FakeEvent("WorkflowCompleted", content="world", metrics=FakeMetrics(), status="COMPLETED") diff --git a/py/src/braintrust/integrations/agno/test_workflow.py b/py/src/braintrust/integrations/agno/test_workflow.py index cd98866c4..e95376289 100644 --- a/py/src/braintrust/integrations/agno/test_workflow.py +++ b/py/src/braintrust/integrations/agno/test_workflow.py @@ -146,7 +146,7 @@ def test_agno_workflow_stream_aggregates_workflow_events(memory_logger): run_response = FakeWorkflowRunResponse(input="hello world") chunks = list(workflow._execute_stream("session-1", execution_input, run_response)) - assert len(chunks) == 4 + assert len(chunks) == 5 spans = memory_logger.pop() assert len(spans) == 1 diff --git a/py/src/braintrust/integrations/agno/tracing.py b/py/src/braintrust/integrations/agno/tracing.py index 29751b4fe..27092c55a 100644 --- a/py/src/braintrust/integrations/agno/tracing.py +++ b/py/src/braintrust/integrations/agno/tracing.py @@ -609,7 +609,7 @@ def _aggregate_workflow_chunks(chunks: list[Any], workflow_run_response: Any | N if hasattr(chunk, "content") and chunk.content: if event == "WorkflowCompleted": final_workflow_content = str(chunk.content) - elif final_workflow_content is None: + elif event != "StepProgress" and final_workflow_content is None: aggregated["content"] += str(chunk.content) if hasattr(chunk, "status") and chunk.status: