[Visualization Extension] Enhance agent graph visualization to prevent infinite recursion, add mermaid graph #445
Conversation
…with recursive handoffs
|
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. |
|
Addressed #387 (comment) and #387 (comment) |
ashish-dahal
changed the title
|
|
||
| @dataclass(frozen=True) | ||
| class Edge: | ||
| source: str |
There was a problem hiding this comment.
shouldn't these be Node not str?
| def _add_agent_nodes_and_edges( | ||
| self, | ||
| agent: Agent, | ||
| parent: Optional[Agent], |
There was a problem hiding this comment.
can you use Agent | None instead?
| graph: Graph, | ||
| ) -> None: | ||
| # Add agent node | ||
| graph.add_node(Node(agent.name, agent.name, NodeType.AGENT)) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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" |
There was a problem hiding this comment.
renderer should be either a Literal or an enum here
Also please add str | None = None instead of Optional
|
should be looking good now |
| response = requests.get(f"https://mermaid.ink/img/{base64_string}") | ||
| response.raise_for_status() |
There was a problem hiding this comment.
I don't love this - relying on a 3rd party service that could return anything or leak information.
I'd actually prefer if we got rid of the Renderer base class, and instead had:
class GraphVizRenderer:
@classmethod
def render_to_file(...)
@classmethod
def render_to_object(...)
class MermaidRenderer:
@classmethod
def render_mermaid_ink_link(...)
@classmethod
def render_mermaid_text(...)
Improves agent graph visualization logic:
Improves on PR #387, addressing pending change requests.