-
Notifications
You must be signed in to change notification settings - Fork 705
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: Add Graphviz-based agent visualization functionality #147
base: main
Are you sure you want to change the base?
Conversation
…into feat/draw_graph
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is such a great PR/idea. Couple of small things to get it in:
- Could you please add this to the
agents/src/extensions
directory? So the import would befrom agents.extensions import visualization
- Could you please add an optional dependency group? So you can do
pip install openai-agents[visualization]
to install it, and if you don't, there's an import error.
Thank you!
tests/test_visualizations.py
Outdated
import graphviz | ||
import pytest | ||
|
||
from src.agents.agent import Agent |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
from src.agents.agent import Agent | |
from agents import Agent |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed!
…into feat/draw_graph
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks like tests are failing. I think you can add an optional dependency to pyproject.toml
[project.optional-dependencies]
visualization=["graphviz"]
# Bigger handoffs (rounded box, yellow) | ||
for handoff in agent.handoffs: | ||
parts.append( | ||
f'"{handoff.name}" [label="{handoff.name}", shape=box, style=filled, style=rounded, ' |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think you want handoff.agent_name
not handoff.name
, if it's a Handoff
instance
f'"{handoff.name}" [label="{handoff.name}", shape=box, style=filled, style=rounded, ' | ||
f"fillcolor=lightyellow, width=1.5, height=0.8];" | ||
) | ||
parts.append(get_all_nodes(handoff)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think this is right - the handoffs list is list[Handoff | Agent]
. So the node might be an Agent, but it might also be a Handoff
object, where you don't statically know which agent is being handed off to. You can use handoff.agent_name
for that
This pull request introduces functionality for visualizing agent structures using Graphviz. The changes include adding a new dependency, implementing functions to generate and draw graphs, and adding tests for these functions.
New functionality for visualizing agent structures:
graphviz
as a new dependency inpyproject.toml
.src/agents/visualizations.py
to generate and draw graphs for agents using Graphviz. These functions includeget_main_graph
,get_all_nodes
,get_all_edges
, anddraw_graph
.Testing the new visualization functionality:
tests/test_visualizations.py
to verify the correctness of the graph generation and drawing functions. The tests coverget_main_graph
,get_all_nodes
,get_all_edges
, anddraw_graph
.For example, given the following code:
Generates the following image: