Update multi_agent.py

multi_agent.py

@@ -21,7 +21,7 @@ def format_as_markdown(code: str) -> str:
     markdown_code += '\n```'
     return markdown_code
     
-def run_multi_agent(llm, message):
+def run_multi_agent(llm, task):
     llm_config = {"model": llm}
     
     executor = LocalCommandLineCodeExecutor(
@@ -46,7 +46,7 @@ def run_multi_agent(llm, message):
     
     chat_result = code_executor_agent.initiate_chat(
         code_writer_agent,
-        message=message,
+        message=task,
         max_turns=10
     )
 
@@ -57,8 +57,16 @@ def run_multi_agent(llm, message):
         if not first_message:
             chat += f"**{message['role'].replace('assistant', 'Code Executor').replace('user', 'Code Writer')}**\n{message['content']}\n\n"
         first_message = False
+
+    file_name_png = ""
+    
+    for file in os.listdir("coding"):
+        if file:
+            _, file_extension = os.path.splitext(file)
+            if file_extension == ".png":
+                file_name_png = file
     
-    image_data = read_image_file("/home/user/app/coding/ytd_stock_gains.png")
+    image_data = read_image_file(f"/home/user/app/coding/{file_name_png}")
     markdown_code_png = generate_markdown_image(image_data)
 
     result = f"{markdown_code_png}\n\n{chat}"