diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index 9180b38b..545911d3 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -655,7 +655,15 @@ def chat_with_workflow(
             int_chat = cast(
                 List[Message],
                 [
-                    {"role": c["role"], "content": c["content"], "media": c["media"]}
+                    (
+                        {
+                            "role": c["role"],
+                            "content": c["content"],
+                            "media": c["media"],
+                        }
+                        if "media" in c
+                        else {"role": c["role"], "content": c["content"]}
+                    )
                     for c in chat
                 ],
             )