fix mypy

landing-ai · May 7, 2024 · d69ce5a · d69ce5a
1 parent b7b1667
commit d69ce5a
Showing 1 changed file with 4 additions and 3 deletions.
diff --git 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,