fixed generate

landing-ai · Jun 10, 2024 · 1ed977c · 1ed977c
1 parent 848c8e0
commit 1ed977c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/vision_agent/agent/agent.py b/vision_agent/agent/agent.py
@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Union
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
 
 from vision_agent.lmm import MediaChatItem
 
@@ -9,6 +10,7 @@ class Agent(ABC):
  def __call__(
  self,
  input: Union[str, List[MediaChatItem]],
+ media: Optional[Union[str, Path]] = None,
  ) -> str:
  pass
 

diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
@@ -460,6 +460,7 @@ def __init__(
  def __call__(
  self,
  input: Union[str, List[MediaChatItem]],
+ media: Optional[Union[str, Path]] = None,
  ) -> str:
  """Chat with Vision Agent and return intermediate information regarding the task.
 
@@ -475,6 +476,8 @@ def __call__(
 
  if isinstance(input, str):
  input = [{"role": "user", "content": input}]
+ if media is not None:
+ input[0]["media"] = [media]
  results = self.chat_with_workflow(input)
  results.pop("working_memory")
  return results # type: ignore