diff --git a/vision_agent/agent/vision_agent_v2.py b/vision_agent/agent/vision_agent_v2.py
index f73e159e..d2d152d3 100644
--- a/vision_agent/agent/vision_agent_v2.py
+++ b/vision_agent/agent/vision_agent_v2.py
@@ -249,12 +249,13 @@ def __call__(
         self,
         input: Union[List[Dict[str, str]], str],
         image: Optional[Union[str, Path]] = None,
-    ) -> Tuple[str, str]:
+    ) -> str:
         if isinstance(input, str):
             input = [{"role": "user", "content": input}]
-        return self.chat(input, image)
+        code, _ = self.chat_with_tests(input, image)
+        return code
 
-    def chat(
+    def chat_with_tests(
         self,
         chat: List[Dict[str, str]],
         image: Optional[Union[str, Path]] = None,