From 0a3dfa071a0853a115bac067cb18bd6a803d087e Mon Sep 17 00:00:00 2001
From: Yuya Haruna <121998910+yuya-haruna@users.noreply.github.com>
Date: Tue, 22 Apr 2025 02:47:51 +0900
Subject: [PATCH] Fix visualize graph filename to without extension. (#554)

Only the file name is needed since graphviz's `render()` automatically
adds the file extension.
Also, unnecessary .gv (.dot) files are output, so the `cleanup=True`
option has been modified to prevent them from being saved.

Here is a similar modification, but in a different content.
- https://github.com/openai/openai-agents-python/pull/451
---
 docs/ja/visualization.md               | 4 ++--
 docs/visualization.md                  | 4 +---
 src/agents/extensions/visualization.py | 2 +-
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/docs/ja/visualization.md b/docs/ja/visualization.md
index 3c5c64e..9093bb6 100644
--- a/docs/ja/visualization.md
+++ b/docs/ja/visualization.md
@@ -81,7 +81,7 @@ draw_graph(triage_agent).view()
 デフォルトでは、`draw_graph` はグラフをインラインで表示します。ファイルとして保存するには、ファイル名を指定します:
 
 ```python
-draw_graph(triage_agent, filename="agent_graph.png")
+draw_graph(triage_agent, filename="agent_graph")
 ```
 
-これにより、作業ディレクトリに `agent_graph.png` が生成されます。
\ No newline at end of file
+これにより、作業ディレクトリに `agent_graph.png` が生成されます。
diff --git a/docs/visualization.md b/docs/visualization.md
index 00f3126..409803f 100644
--- a/docs/visualization.md
+++ b/docs/visualization.md
@@ -78,9 +78,7 @@ draw_graph(triage_agent).view()
 By default, `draw_graph` displays the graph inline. To save it as a file, specify a filename:
 
 ```python
-draw_graph(triage_agent, filename="agent_graph.png")
+draw_graph(triage_agent, filename="agent_graph")
 ```
 
 This will generate `agent_graph.png` in the working directory.
-
-
diff --git a/src/agents/extensions/visualization.py b/src/agents/extensions/visualization.py
index 5fb3506..888e262 100644
--- a/src/agents/extensions/visualization.py
+++ b/src/agents/extensions/visualization.py
@@ -132,6 +132,6 @@ def draw_graph(agent: Agent, filename: Optional[str] = None) -> graphviz.Source:
     graph = graphviz.Source(dot_code)
 
     if filename:
-        graph.render(filename, format="png")
+        graph.render(filename, format="png", cleanup=True)
 
     return graph