diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index b60e8391..736c9754 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -123,7 +123,7 @@ def __init__(
         verbosity: int = 0,
         local_artifacts_path: Optional[Union[str, Path]] = None,
         code_sandbox_runtime: Optional[str] = None,
-        callback_message: Optional[Callable[[Message], None]] = None,
+        callback_message: Optional[Callable[[Dict[str, Any]], None]] = None,
     ) -> None:
         """Initialize the VisionAgent.
 
@@ -324,7 +324,7 @@ def chat_with_code(
             artifacts.save()
         return orig_chat, artifacts
 
-    def streaming_message(self, message: Message) -> None:
+    def streaming_message(self, message: Dict[str, Any]) -> None:
         if self.callback_message:
             self.callback_message(message)