Skip to content
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

[Visualization Extension] Enhance agent graph visualization to prevent infinite recursion, add mermaid graph #445

Open
wants to merge 9 commits into
base: main
Choose a base branch
from

Conversation

ashish-dahal
Copy link

Improves agent graph visualization logic:

  • Safely handles recursive agent handoffs during visualization
  • Adds a test case for recursive handoff loops
  • Cleans up redundant code and improves readability in the visualization test

Improves on PR #387, addressing pending change requests.

@ashish-dahal
Copy link
Author

Hey @rm-openai I noticed the original PR hasn’t had updates, so I’ve created a follow-up here with additional fixes and improvements. Happy to adjust things if needed.

@ashish-dahal
Copy link
Author

Addressed #387 (comment) and #387 (comment)

@ashish-dahal ashish-dahal changed the title [Visualization Extension] Enhance agent graph visualization to prevent infinite recursion for agents with recursive handoffs [Visualization Extension] Enhance agent graph visualization to prevent infinite recursion, add mermaid graph Apr 6, 2025
Copy link
Collaborator

@rm-openai rm-openai left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, PR is looking great! Some changes requested inline


@dataclass(frozen=True)
class Edge:
source: str
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shouldn't these be Node not str?

def _add_agent_nodes_and_edges(
self,
agent: Agent,
parent: Optional[Agent],
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you use Agent | None instead?

graph: Graph,
) -> None:
# Add agent node
graph.add_node(Node(agent.name, agent.name, NodeType.AGENT))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

using agent.name as the ID isn't safe unfortunately, since agents can have the same name. You'd need to do an instance equality check everywhere instead.

Copy link
Author

@ashish-dahal ashish-dahal Apr 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It works when agent.handoffs have Agent instances. but apparently it's a bit tricky when Handoff objects are passed. It only has the target agent's name (handoff.agent_name), not an Agent instance. The instance itself only gets resolved at runtime with handoff.on_invoke_handoff. so we can't consistently do instance equality checks to link Handoff instances back to the Agent instances at build time if we consider cases with recursive handoffs. The only option I see is to resolve this using agent names but mention the caveat in the docs. any suggestions?



def draw_graph(
agent: Agent, filename: Optional[str] = None, renderer: str = "graphviz"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

renderer should be either a Literal or an enum here

Also please add str | None = None instead of Optional

Copy link
Collaborator

@rm-openai rm-openai left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as before

@ashish-dahal
Copy link
Author

should be looking good now

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants