diff --git a/flaml/autogen/agent/user_proxy_agent.py b/flaml/autogen/agent/user_proxy_agent.py
index 210bb698e7..dd68c3fec9 100644
--- a/flaml/autogen/agent/user_proxy_agent.py
+++ b/flaml/autogen/agent/user_proxy_agent.py
@@ -91,7 +91,7 @@ def _execute_code(self, code_blocks):
                 logs = logs.decode("utf-8")
             elif lang in ["python", "Python"]:
                 if code.startswith("# filename: "):
-                    filename = code[11 : code.find("\n")].strip()
+                    filename = code[11 : code.find(".py")].strip() + '.py'
                 else:
                     filename = None
                 exitcode, logs, image = execute_code(