diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index 355fc214..d5e65169 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -329,7 +329,7 @@ def __call__(
             input = [{"role": "user", "content": input}]
         results = self.chat_with_workflow(input, media)
         results.pop("working_memory")
-        return results["code"]  # type: ignore
+        return results  # type: ignore
 
     def chat_with_workflow(
         self,