From 63f41cff1a57a59e4074b9c102afce57fb37b510 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 6 Jan 2026 22:58:44 +0000 Subject: [PATCH] Update tests --- tests/integration/test_agent_streaming_integration.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/integration/test_agent_streaming_integration.py b/tests/integration/test_agent_streaming_integration.py index 0971d30c..d6004c21 100644 --- a/tests/integration/test_agent_streaming_integration.py +++ b/tests/integration/test_agent_streaming_integration.py @@ -47,8 +47,9 @@ Args: { "}" ] - for chunk in chunks: - await chunk_callback(chunk) + for i, chunk in enumerate(chunks): + is_final = (i == len(chunks) - 1) + await chunk_callback(chunk, is_final) return full_text else: @@ -312,8 +313,10 @@ Final Answer: AI is the simulation of human intelligence in machines.""" call_count += 1 if streaming and chunk_callback: - for chunk in response.split(): - await chunk_callback(chunk + " ") + chunks = response.split() + for i, chunk in enumerate(chunks): + is_final = (i == len(chunks) - 1) + await chunk_callback(chunk + " ", is_final) return response return response