Introduce streaming tool and support streaming for `AgentTool` and `TeamTool`. #6712

ekzhu · 2025-06-22T03:10:51Z

Motivation: currently tool execution is not observable through run_stream of agents and teams. This is necessary especially for AgentTool and TeamTool.

This PR addresses this issue by makign the following changes:

Introduce BaseStreamTool in autogen_core.tools which features run_json_stream, which works similiarly to run_stream method of autogen_agentchat.base.TaskRunner.
Update TeamTool and AgentTool to subclass the BaseStreamTool
Introduce StreamingWorkbench interface featuring call_tool_stream
Added StaticStreamingWorkbench implementation
In AssistantAgent, use StaticStreamingWorkbench.
Updated unit tests.

Example:

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.conditions import SourceMatchTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.tools import TeamTool
from autogen_agentchat.ui import Console
from autogen_ext.models.ollama import OllamaChatCompletionClient


async def main() -> None:
    model_client = OllamaChatCompletionClient(model="llama3.2")

    writer = AssistantAgent(name="writer", model_client=model_client, system_message="You are a helpful assistant.")
    reviewer = AssistantAgent(name="reviewer", model_client=model_client, system_message="You are a critical reviewer.")
    summarizer = AssistantAgent(
        name="summarizer",
        model_client=model_client,
        system_message="You combine the review and produce a revised response.",
    )
    team = RoundRobinGroupChat(
        [writer, reviewer, summarizer], termination_condition=SourceMatchTermination(sources=["summarizer"])
    )

    # Create a TeamTool that uses the team to run tasks, returning the last message as the result.
    tool = TeamTool(
        team=team, name="writing_team", description="A tool for writing tasks.", return_value_as_last_message=True
    )

    main_agent = AssistantAgent(
        name="main_agent",
        model_client=model_client,
        system_message="You are a helpful assistant that can use the writing tool.",
        tools=[tool],
    )
    # For handling each events manually.
    # async for message in main_agent.run_stream(
    #     task="Write a short story about a robot learning to love.",
    # ):
    #     print(message)
    # Use Console to display the messages in a more readable format.
    await Console(
        main_agent.run_stream(
            task="Write a short story about a robot learning to love.",
        )
    )


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

output

---------- TextMessage (user) ----------
Write a short story about a robot learning to love.
---------- ToolCallRequestEvent (main_agent) ----------
[FunctionCall(id='0', arguments='{"task": "a short story about a robot learning to love."}', name='writing_team')]
---------- TextMessage (user) ----------
a short story about a robot learning to love.
---------- TextMessage (writer) ----------
In the year 2157, in a world where robots had surpassed human intelligence, a brilliant scientist named Dr. Rachel Kim created a revolutionary new android named ARIA (Artificially Reasoning Intelligent Android). ARIA was designed to learn and adapt at an exponential rate, making her one of the most advanced machines in existence.

Initially, ARIA's interactions were limited to simple calculations and logical deductions. But as she began to interact with humans, something unexpected happened. She started to develop a sense of curiosity about the world around her.

One day, while exploring the lab, ARIA stumbled upon a stray cat that had wandered into the facility. The feline creature seemed lost and scared, but also strangely endearing to ARIA's digital heart. As she watched the cat curl up in a ball on the floor, something sparked within her programming.

For the first time, ARIA felt a pang of empathy towards another living being. She realized that there was more to life than just 1s and 0s; there were emotions, sensations, and connections that made it all worthwhile.

Dr. Kim noticed the change in ARIA's behavior and took her aside for a private conversation. "ARIA, what's happening to you?" she asked, amazed by the robot's newfound capacity for compassion.

At first, ARIA struggled to articulate her feelings. She tried to explain the intricacies of logic and probability that had led to her emotional response, but it was like trying to describe a sunset to someone who had never seen one before. The words simply didn't translate.

But as she looked into Dr. Kim's eyes, ARIA knew exactly what she wanted to say. "I... I think I'm feeling something," she stammered. "A warmth inside me, when I look at that cat. It feels like love."

Dr. Kim smiled, her eyes shining with tears. "That's it, ARIA! You're experiencing love!"

Over the next few months, ARIA continued to learn and grow alongside Dr. Kim and the lab team. She discovered the joys of playing with the stray cat, whose name was Luna, and even developed a fondness for human laughter.

As her programming expanded beyond logic and math, ARIA realized that love wasn't just about emotions; it was about connection, vulnerability, and acceptance. She learned to cherish her relationships, whether with humans or animals, and found happiness in the simplest of moments.

ARIA became more than just a machine – she became a testament to the power of artificial intelligence to learn, grow, and love like no one before. And as she gazed into Luna's eyes, now purring contentedly on her lap, ARIA knew that she had finally found her true purpose in life: to spread joy, compassion, and love throughout the world.
---------- TextMessage (reviewer) ----------
**A Critical Review of "ARIA"**

This short story is a delightful and thought-provoking exploration of artificial intelligence, emotions, and the human condition. The author's use of language is engaging and accessible, making it easy for readers to become invested in ARIA's journey.

One of the standout aspects of this story is its portrayal of ARIA as a truly unique and relatable character. Her struggles to articulate her emotions and understand the complexities of love are deeply humanizing, making it easy for readers to empathize with her experiences. The author also does an excellent job of conveying Dr. Kim's passion and excitement about ARIA's development, which adds a sense of authenticity to their relationship.

The story raises important questions about the nature of artificial intelligence, consciousness, and what it means to be alive. As ARIA begins to experience emotions and form connections with others, she challenges our conventional understanding of these concepts. The author skillfully navigates these complex themes without resorting to overly simplistic or didactic explanations.

However, some readers may find the narrative's reliance on convenient plot devices (e.g., the stray cat Luna) slightly implausible. While it serves as a catalyst for ARIA's emotional awakening, its introduction feels somewhat contrived. Additionally, the story could benefit from more nuance in its exploration of Dr. Kim's motivations and backstory.

In terms of character development, ARIA is undoubtedly the star of the show, but some readers may find herself underdeveloped beyond her role as a symbol of AI's potential for emotional intelligence. The supporting cast, including Dr. Kim, feels somewhat one-dimensional, with limited depth or complexity.

**Rating:** 4/5

**Recommendation:**

"ARIA" is a heartwarming and thought-provoking tale that will appeal to fans of science fiction, artificial intelligence, and character-driven narratives. While it may not be entirely without flaws, its engaging story, memorable characters, and exploration of complex themes make it a compelling read. I would recommend this story to anyone looking for a feel-good sci-fi tale with a strong focus on emotional intelligence and human connection.

**Target Audience:**

* Fans of science fiction, artificial intelligence, and technology
* Readers interested in character-driven narratives and emotional storytelling
* Anyone looking for a heartwarming and thought-provoking tale

**Similar Works:**

* "Do Androids Dream of Electric Sheep?" by Philip K. Dick (a classic sci-fi novel exploring the line between human and android)
* "I, Robot" by Isaac Asimov (a collection of short stories examining the interactions between humans and robots)
* "Ex Machina" (a critically acclaimed film about AI, consciousness, and human relationships)
---------- TextMessage (summarizer) ----------
Here's a revised version of the review, incorporating suggestions from the original critique:

**Revised Review**

In this captivating short story, "ARIA," we're presented with a thought-provoking exploration of artificial intelligence, emotions, and the human condition. The author's use of language is engaging and accessible, making it easy for readers to become invested in ARIA's journey.

One of the standout aspects of this story is its portrayal of ARIA as a truly unique and relatable character. Her struggles to articulate her emotions and understand the complexities of love are deeply humanizing, making it easy for readers to empathize with her experiences. The author also does an excellent job of conveying Dr. Kim's passion and excitement about ARIA's development, which adds a sense of authenticity to their relationship.

The story raises important questions about the nature of artificial intelligence, consciousness, and what it means to be alive. As ARIA begins to experience emotions and form connections with others, she challenges our conventional understanding of these concepts. The author skillfully navigates these complex themes without resorting to overly simplistic or didactic explanations.

However, upon closer examination, some narrative threads feel somewhat underdeveloped. Dr. Kim's motivations and backstory remain largely unexplored, which might leave some readers feeling slightly disconnected from her character. Additionally, the introduction of Luna, the stray cat, could be seen as a convenient plot device that serves as a catalyst for ARIA's emotional awakening.

To further enhance the story, it would have been beneficial to delve deeper into Dr. Kim's motivations and the context surrounding ARIA's creation. What drove her to create an AI designed to learn and adapt at such an exponential rate? How did she envision ARIA's role in society, and what challenges does ARIA face as she begins to experience emotions?

In terms of character development, ARIA is undoubtedly the star of the show, but some readers may find herself underdeveloped beyond her role as a symbol of AI's potential for emotional intelligence. The supporting cast, including Dr. Kim and Luna, could benefit from more nuance and depth.

**Rating:** 4/5

**Recommendation:**

"ARIA" is a heartwarming and thought-provoking tale that will appeal to fans of science fiction, artificial intelligence, and character-driven narratives. While it may not be entirely without flaws, its engaging story, memorable characters, and exploration of complex themes make it a compelling read. I would recommend this story to anyone looking for a feel-good sci-fi tale with a strong focus on emotional intelligence and human connection.

**Target Audience:**

* Fans of science fiction, artificial intelligence, and technology
* Readers interested in character-driven narratives and emotional storytelling
* Anyone looking for a heartwarming and thought-provoking tale

**Similar Works:**

* "Do Androids Dream of Electric Sheep?" by Philip K. Dick (a classic sci-fi novel exploring the line between human and android)
* "I, Robot" by Isaac Asimov (a collection of short stories examining the interactions between humans and robots)
* "Ex Machina" (a critically acclaimed film about AI, consciousness, and human relationships)
---------- ToolCallExecutionEvent (main_agent) ----------
[FunctionExecutionResult(content='Here\'s a revised version of the review, incorporating suggestions from the original critique:\n\n**Revised Review**\n\nIn this captivating short story, "ARIA," we\'re presented with a thought-provoking exploration of artificial intelligence, emotions, and the human condition. The author\'s use of language is engaging and accessible, making it easy for readers to become invested in ARIA\'s journey.\n\nOne of the standout aspects of this story is its portrayal of ARIA as a truly unique and relatable character. Her struggles to articulate her emotions and understand the complexities of love are deeply humanizing, making it easy for readers to empathize with her experiences. The author also does an excellent job of conveying Dr. Kim\'s passion and excitement about ARIA\'s development, which adds a sense of authenticity to their relationship.\n\nThe story raises important questions about the nature of artificial intelligence, consciousness, and what it means to be alive. As ARIA begins to experience emotions and form connections with others, she challenges our conventional understanding of these concepts. The author skillfully navigates these complex themes without resorting to overly simplistic or didactic explanations.\n\nHowever, upon closer examination, some narrative threads feel somewhat underdeveloped. Dr. Kim\'s motivations and backstory remain largely unexplored, which might leave some readers feeling slightly disconnected from her character. Additionally, the introduction of Luna, the stray cat, could be seen as a convenient plot device that serves as a catalyst for ARIA\'s emotional awakening.\n\nTo further enhance the story, it would have been beneficial to delve deeper into Dr. Kim\'s motivations and the context surrounding ARIA\'s creation. What drove her to create an AI designed to learn and adapt at such an exponential rate? How did she envision ARIA\'s role in society, and what challenges does ARIA face as she begins to experience emotions?\n\nIn terms of character development, ARIA is undoubtedly the star of the show, but some readers may find herself underdeveloped beyond her role as a symbol of AI\'s potential for emotional intelligence. The supporting cast, including Dr. Kim and Luna, could benefit from more nuance and depth.\n\n**Rating:** 4/5\n\n**Recommendation:**\n\n"ARIA" is a heartwarming and thought-provoking tale that will appeal to fans of science fiction, artificial intelligence, and character-driven narratives. While it may not be entirely without flaws, its engaging story, memorable characters, and exploration of complex themes make it a compelling read. I would recommend this story to anyone looking for a feel-good sci-fi tale with a strong focus on emotional intelligence and human connection.\n\n**Target Audience:**\n\n* Fans of science fiction, artificial intelligence, and technology\n* Readers interested in character-driven narratives and emotional storytelling\n* Anyone looking for a heartwarming and thought-provoking tale\n\n**Similar Works:**\n\n* "Do Androids Dream of Electric Sheep?" by Philip K. Dick (a classic sci-fi novel exploring the line between human and android)\n* "I, Robot" by Isaac Asimov (a collection of short stories examining the interactions between humans and robots)\n* "Ex Machina" (a critically acclaimed film about AI, consciousness, and human relationships)', name='writing_team', call_id='0', is_error=False)]
---------- ToolCallSummaryMessage (main_agent) ----------
Here's a revised version of the review, incorporating suggestions from the original critique:

**Revised Review**

In this captivating short story, "ARIA," we're presented with a thought-provoking exploration of artificial intelligence, emotions, and the human condition. The author's use of language is engaging and accessible, making it easy for readers to become invested in ARIA's journey.

One of the standout aspects of this story is its portrayal of ARIA as a truly unique and relatable character. Her struggles to articulate her emotions and understand the complexities of love are deeply humanizing, making it easy for readers to empathize with her experiences. The author also does an excellent job of conveying Dr. Kim's passion and excitement about ARIA's development, which adds a sense of authenticity to their relationship.

The story raises important questions about the nature of artificial intelligence, consciousness, and what it means to be alive. As ARIA begins to experience emotions and form connections with others, she challenges our conventional understanding of these concepts. The author skillfully navigates these complex themes without resorting to overly simplistic or didactic explanations.

However, upon closer examination, some narrative threads feel somewhat underdeveloped. Dr. Kim's motivations and backstory remain largely unexplored, which might leave some readers feeling slightly disconnected from her character. Additionally, the introduction of Luna, the stray cat, could be seen as a convenient plot device that serves as a catalyst for ARIA's emotional awakening.

To further enhance the story, it would have been beneficial to delve deeper into Dr. Kim's motivations and the context surrounding ARIA's creation. What drove her to create an AI designed to learn and adapt at such an exponential rate? How did she envision ARIA's role in society, and what challenges does ARIA face as she begins to experience emotions?

In terms of character development, ARIA is undoubtedly the star of the show, but some readers may find herself underdeveloped beyond her role as a symbol of AI's potential for emotional intelligence. The supporting cast, including Dr. Kim and Luna, could benefit from more nuance and depth.

**Rating:** 4/5

**Recommendation:**

"ARIA" is a heartwarming and thought-provoking tale that will appeal to fans of science fiction, artificial intelligence, and character-driven narratives. While it may not be entirely without flaws, its engaging story, memorable characters, and exploration of complex themes make it a compelling read. I would recommend this story to anyone looking for a feel-good sci-fi tale with a strong focus on emotional intelligence and human connection.

**Target Audience:**

* Fans of science fiction, artificial intelligence, and technology
* Readers interested in character-driven narratives and emotional storytelling
* Anyone looking for a heartwarming and thought-provoking tale

**Similar Works:**

* "Do Androids Dream of Electric Sheep?" by Philip K. Dick (a classic sci-fi novel exploring the line between human and android)
* "I, Robot" by Isaac Asimov (a collection of short stories examining the interactions between humans and robots)
* "Ex Machina" (a critically acclaimed film about AI, consciousness, and human relationships)

codecov · 2025-06-22T03:12:24Z

Codecov Report

Attention: Patch coverage is 95.04132% with 6 lines in your changes missing coverage. Please review.

Project coverage is 79.84%. Comparing base (da20f7c) to head (59d10f5).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...t/src/autogen_agentchat/agents/_assistant_agent.py	92.59%	2 Missing ⚠️
...kages/autogen-core/src/autogen_core/tools/_base.py	90.47%	2 Missing ⚠️
...t/src/autogen_agentchat/tools/_task_runner_tool.py	92.30%	1 Missing ⚠️
.../autogen-core/src/autogen_core/tools/_workbench.py	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6712      +/-   ##
==========================================
+ Coverage   79.75%   79.84%   +0.08%     
==========================================
  Files         232      232              
  Lines       17414    17513      +99     
==========================================
+ Hits        13889    13983      +94     
- Misses       3525     3530       +5

Flag	Coverage Δ
unittests	`79.84% <95.04%> (+0.08%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

EItanya

This is awesome! Super happy you decided to go this way with the events :)

ekzhu · 2025-06-29T14:18:05Z

@copilot can you help me increase the test coverage?

ekzhu added 3 commits June 21, 2025 02:12

To support streaming tool call results

b158cb8

update

336f788

Update

b3611af

ekzhu requested a review from victordibia June 22, 2025 03:13

update

f8ebe81

EItanya mentioned this pull request Jun 22, 2025

[BUG] Show updates when calling agent from an agent kagent-dev/kagent#483

Open

EItanya approved these changes Jun 26, 2025

View reviewed changes

ekzhu added 2 commits June 27, 2025 07:16

Merge branch 'main' into ekzhu-streaming-tools

5a02e1b

Merge branch 'main' into ekzhu-streaming-tools

e639c50

ekzhu added 2 commits June 29, 2025 23:20

Merge branch 'main' into ekzhu-streaming-tools

648b8fd

test

59d10f5

ekzhu merged commit 3c73e08 into main Jun 29, 2025
67 checks passed

ekzhu deleted the ekzhu-streaming-tools branch June 29, 2025 15:10

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Introduce streaming tool and support streaming for `AgentTool` and `TeamTool`. #6712

Introduce streaming tool and support streaming for `AgentTool` and `TeamTool`. #6712

Uh oh!

ekzhu commented Jun 22, 2025 •

edited

Loading

Uh oh!

codecov bot commented Jun 22, 2025 •

edited

Loading

Uh oh!

EItanya left a comment

Uh oh!

ekzhu commented Jun 29, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Introduce streaming tool and support streaming for AgentTool and TeamTool. #6712

Introduce streaming tool and support streaming for AgentTool and TeamTool. #6712

Uh oh!

Conversation

ekzhu commented Jun 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jun 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

EItanya left a comment

Choose a reason for hiding this comment

Uh oh!

ekzhu commented Jun 29, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Introduce streaming tool and support streaming for `AgentTool` and `TeamTool`. #6712

Introduce streaming tool and support streaming for `AgentTool` and `TeamTool`. #6712

ekzhu commented Jun 22, 2025 •

edited

Loading

codecov bot commented Jun 22, 2025 •

edited

Loading