vision-agent 0.2.117__py3-none-any.whl → 0.2.119__py3-none-any.whl

vision_agent/tools/meta_tools.py

@@ -1,12 +1,17 @@
 import os
+import pickle as pkl
 import subprocess
+import tempfile
 from pathlib import Path
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
+
+from IPython.display import display
 
 import vision_agent as va
 from vision_agent.lmm.types import Message
 from vision_agent.tools.tool_utils import get_tool_documentation
 from vision_agent.tools.tools import TOOL_DESCRIPTIONS
+from vision_agent.utils.execute import Execution, MimeType
 
 # These tools are adapted from SWE-Agent https://github.com/princeton-nlp/SWE-agent
 
@@ -35,97 +40,91 @@ def filter_file(file_name: Union[str, Path]) -> bool:
     )
 
 
-def generate_vision_code(save_file: str, chat: str, media: List[str]) -> str:
-    """Generates python code to solve vision based tasks.
-
-    Parameters:
-        save_file (str): The file path to save the code.
-        chat (str): The chat message from the user.
-        media (List[str]): The media files to use.
-
-    Returns:
-        str: The generated code.
-
-    Examples
-    --------
-        >>> generate_vision_code("code.py", "Can you detect the dogs in this image?", ["image.jpg"])
-        from vision_agent.tools import load_image, owl_v2
-        def detect_dogs(image_path: str):
-            image = load_image(image_path)
-            dogs = owl_v2("dog", image)
-            return dogs
+def redisplay_results(execution: Execution) -> None:
+    """This function is used to add previous execution results to the current output.
+    This is handy if you are inside a notebook environment, call it notebook1, and you
+    have a nested notebook environment, call it notebook2, and you want the execution
+    results from notebook2 to be included in the execution results for notebook1.
+    """
+    for result in execution.results:
+        if result.text is not None:
+            display({MimeType.TEXT_PLAIN: result.text})
+        if result.html is not None:
+            display({MimeType.TEXT_HTML: result.html})
+        if result.markdown is not None:
+            display({MimeType.TEXT_MARKDOWN: result.markdown})
+        if result.svg is not None:
+            display({MimeType.IMAGE_SVG: result.svg})
+        if result.png is not None:
+            display({MimeType.IMAGE_PNG: result.png})
+        if result.jpeg is not None:
+            display({MimeType.IMAGE_JPEG: result.jpeg})
+        if result.mp4 is not None:
+            display({MimeType.VIDEO_MP4_B64: result.mp4})
+        if result.latex is not None:
+            display({MimeType.TEXT_LATEX: result.latex})
+        if result.json is not None:
+            display({MimeType.APPLICATION_JSON: result.json})
+        if result.extra is not None:
+            display(result.extra)
+
+
+class Artifacts:
+    """Artifacts is a class that allows you to sync files between a local and remote
+    environment. In our case, the remote environment could be where the VisionAgent is
+    executing code and as the user adds new images, files or modifies files, those
+    need to be in sync with the remote environment the VisionAgent is running in.
     """
 
-    if ZMQ_PORT is not None:
-        agent = va.agent.VisionAgentCoder(
-            report_progress_callback=lambda inp: report_progress_callback(
-                int(ZMQ_PORT), inp
+    def __init__(self, remote_save_path: Union[str, Path]) -> None:
+        self.remote_save_path = Path(remote_save_path)
+        self.artifacts: Dict[str, Any] = {}
+
+        self.code_sandbox_runtime = None
+
+    def load(self, file_path: Union[str, Path]) -> None:
+        """Loads are artifacts into the remote environment. If an artifact value is None
+        it will skip loading it.
+
+        Parameters:
+            file_path (Union[str, Path]): The file path to load the artifacts from
+        """
+        with open(file_path, "rb") as f:
+            self.artifacts = pkl.load(f)
+        for k, v in self.artifacts.items():
+            if v is not None:
+                mode = "w" if isinstance(v, str) else "wb"
+                with open(self.remote_save_path.parent / k, mode) as f:
+                    f.write(v)
+
+    def show(self) -> str:
+        """Shows the artifacts that have been loaded and their remote save paths."""
+        out_str = "[Artifacts loaded]\n"
+        for k in self.artifacts.keys():
+            out_str += (
+                f"Artifact {k} loaded to {str(self.remote_save_path.parent / k)}\n"
             )
+        out_str += "[End of artifacts]\n"
+        return out_str
+
+    def save(self, local_path: Optional[Union[str, Path]] = None) -> None:
+        save_path = (
+            Path(local_path) if local_path is not None else self.remote_save_path
         )
-    else:
-        agent = va.agent.VisionAgentCoder()
-    try:
-        fixed_chat: List[Message] = [{"role": "user", "content": chat, "media": media}]
-        response = agent.chat_with_workflow(fixed_chat)
-        code = response["code"]
-        with open(save_file, "w") as f:
-            f.write(code)
-        code_lines = code.splitlines(keepends=True)
-        total_lines = len(code_lines)
-        return view_lines(code_lines, 0, total_lines, save_file, total_lines)
-    except Exception as e:
-        return str(e)
-
-
-def edit_vision_code(code_file: str, chat_history: List[str], media: List[str]) -> str:
-    """Edits python code to solve a vision based task.
+        with open(save_path, "wb") as f:
+            pkl.dump(self.artifacts, f)
 
-    Parameters:
-        code_file (str): The file path to the code.
-        chat_history (List[str]): The chat history to used to generate the code.
+    def __iter__(self) -> Any:
+        return iter(self.artifacts)
 
-    Returns:
-        str: The edited code.
+    def __getitem__(self, name: str) -> Any:
+        return self.artifacts[name]
 
-    Examples
-    --------
-        >>> edit_vision_code(
-        >>>     "code.py",
-        >>>     ["Can you detect the dogs in this image?", "Can you use a higher threshold?"],
-        >>>     ["dog.jpg"],
-        >>> )
-        from vision_agent.tools import load_image, owl_v2
-        def detect_dogs(image_path: str):
-            image = load_image(image_path)
-            dogs = owl_v2("dog", image, threshold=0.8)
-            return dogs
-    """
+    def __setitem__(self, name: str, value: Any) -> None:
+        self.artifacts[name] = value
 
-    agent = va.agent.VisionAgentCoder()
-    with open(code_file, "r") as f:
-        code = f.read()
-
-    # Append latest code to second to last message from assistant
-    fixed_chat_history: List[Message] = []
-    for i, chat in enumerate(chat_history):
-        if i == 0:
-            fixed_chat_history.append({"role": "user", "content": chat, "media": media})
-        elif i > 0 and i < len(chat_history) - 1:
-            fixed_chat_history.append({"role": "user", "content": chat})
-        elif i == len(chat_history) - 1:
-            fixed_chat_history.append({"role": "assistant", "content": code})
-            fixed_chat_history.append({"role": "user", "content": chat})
-
-    try:
-        response = agent.chat_with_workflow(fixed_chat_history, test_multi_plan=False)
-        code = response["code"]
-        with open(code_file, "w") as f:
-            f.write(code)
-        code_lines = code.splitlines(keepends=True)
-        total_lines = len(code_lines)
-        return view_lines(code_lines, 0, total_lines, code_file, total_lines)
-    except Exception as e:
-        return str(e)
+    def __contains__(self, name: str) -> bool:
+        return name in self.artifacts
 
 
 def format_lines(lines: List[str], start_idx: int) -> str:
@@ -136,34 +135,40 @@ def format_lines(lines: List[str], start_idx: int) -> str:
 
 
 def view_lines(
-    lines: List[str], line_num: int, window_size: int, file_path: str, total_lines: int
+    lines: List[str], line_num: int, window_size: int, name: str, total_lines: int
 ) -> str:
     start = max(0, line_num - window_size)
     end = min(len(lines), line_num + window_size)
-    return (
-        f"[File: {file_path} ({total_lines} lines total)]\n"
+    return_str = (
+        f"[Artifact: {name} ({total_lines} lines total)]\n"
         + format_lines(lines[start:end], start)
-        + ("[End of file]" if end == len(lines) else f"[{len(lines) - end} more lines]")
+        + (
+            "[End of artifact]"
+            if end == len(lines)
+            else f"[{len(lines) - end} more lines]"
+        )
     )
+    print(return_str)
+    return return_str
 
 
-def open_file(file_path: str, line_num: int = 0, window_size: int = 100) -> str:
-    """Opens the file at at the given path in the editor. If `line_num` is provided,
-    the window will be moved to include that line. It only shows the first 100 lines by
-    default! Max `window_size` supported is 2000. use `scroll up/down` to view the file
-    if you want to see more.
+def open_code_artifact(
+    artifacts: Artifacts, name: str, line_num: int = 0, window_size: int = 100
+) -> str:
+    """Opens the provided code artifact. If `line_num` is provided, the window will be
+    moved to include that line. It only shows the first 100 lines by default! Max
+    `window_size` supported is 2000.
 
     Parameters:
-        file_path (str): The file path to open, preferred absolute path.
+        artifacts (Artifacts): The artifacts object to open the artifact from.
+        name (str): The name of the artifact to open.
         line_num (int): The line number to move the window to.
         window_size (int): The number of lines to show above and below the line.
     """
+    if name not in artifacts:
+        return f"[Artifact {name} does not exist]"
 
-    file_path_p = Path(file_path)
-    if not file_path_p.exists():
-        return f"[File {file_path} does not exist]"
-
-    total_lines = sum(1 for _ in open(file_path_p))
+    total_lines = len(artifacts[name].splitlines())
     window_size = min(window_size, 2000)
     window_size = window_size // 2
     if line_num - window_size < 0:
@@ -171,211 +176,218 @@ def open_file(file_path: str, line_num: int = 0, window_size: int = 100) -> str:
     elif line_num >= total_lines:
         line_num = total_lines - 1 - window_size
 
-    global CURRENT_LINE, CURRENT_FILE
-    CURRENT_LINE = line_num
-    CURRENT_FILE = file_path
+    lines = artifacts[name].splitlines(keepends=True)
 
-    with open(file_path, "r") as f:
-        lines = f.readlines()
+    return view_lines(lines, line_num, window_size, name, total_lines)
 
-    return view_lines(lines, line_num, window_size, file_path, total_lines)
 
-
-def create_file(file_path: str) -> str:
-    """Creates and opens a new file with the given name.
+def create_code_artifact(artifacts: Artifacts, name: str) -> str:
+    """Creates a new code artifiact with the given name.
 
     Parameters:
-        file_path (str): The file path to create, preferred absolute path.
+        artifacts (Artifacts): The artifacts object to add the new artifact to.
+        name (str): The name of the new artifact.
     """
+    if name in artifacts:
+        return_str = f"[Artifact {name} already exists]"
+    else:
+        artifacts[name] = ""
+        return_str = f"[Artifact {name} created]"
+    print(return_str)
 
-    file_path_p = Path(file_path)
-    if file_path_p.exists():
-        return f"[File {file_path} already exists]"
-    file_path_p.touch()
-    global CURRENT_FILE
-    CURRENT_FILE = file_path
-    return f"[File created {file_path}]"
-
+    display({MimeType.APPLICATION_JSON: {"last_artifact": name}})
+    return return_str
 
-def scroll_up() -> str:
-    """Moves the window up by 100 lines."""
-    if CURRENT_FILE is None:
-        return "[No file is open]"
 
-    return open_file(CURRENT_FILE, CURRENT_LINE + DEFAULT_WINDOW_SIZE)
+def edit_code_artifact(
+    artifacts: Artifacts, name: str, start: int, end: int, content: str
+) -> str:
+    """Edits the given code artifact with the provided content. The content will be
+    inserted between the `start` and `end` line numbers. If the `start` and `end` are
+    the same, the content will be inserted at the `start` line number. If the `end` is
+    greater than the total number of lines in the file, the content will be inserted at
+    the end of the file. If the `start` or `end` are negative, the function will return
+    an error message.
 
+    Parameters:
+        artifacts (Artifacts): The artifacts object to edit the artifact from.
+        name (str): The name of the artifact to edit.
+        start (int): The line number to start the edit.
+        end (int): The line number to end the edit.
+        content (str): The content to insert.
+    """
+    # just make the artifact if it doesn't exist instead of forcing agent to call
+    # create_artifact
+    if name not in artifacts:
+        artifacts[name] = ""
 
-def scroll_down() -> str:
-    """Moves the window down by 100 lines."""
-    if CURRENT_FILE is None:
-        return "[No file is open]"
+    total_lines = len(artifacts[name].splitlines())
+    if start < 0 or end < 0 or start > end or end > total_lines:
+        return "[Invalid line range]"
+    if start == end:
+        end += 1
 
-    return open_file(CURRENT_FILE, CURRENT_LINE - DEFAULT_WINDOW_SIZE)
+    new_content_lines = content.splitlines(keepends=True)
+    new_content_lines = [
+        line if line.endswith("\n") else line + "\n" for line in new_content_lines
+    ]
+    lines = artifacts[name].splitlines()
+    edited_lines = lines[:start] + new_content_lines + lines[end:]
 
+    cur_line = start + len(content.split("\n")) // 2
+    with tempfile.NamedTemporaryFile(delete=True) as f:
+        with open(f.name, "w") as f:  # type: ignore
+            f.writelines(edited_lines)
+
+        process = subprocess.Popen(
+            [
+                "flake8",
+                "--isolated",
+                "--select=F821,F822,F831,E111,E112,E113,E999,E902",
+                f.name,
+            ],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        stdout, _ = process.communicate()
+
+        if stdout != "":
+            stdout = stdout.replace(f.name, name)
+            error_msg = "[Edit failed with the following status]\n" + stdout
+            original_view = view_lines(
+                lines,
+                start + ((end - start) // 2),
+                DEFAULT_WINDOW_SIZE,
+                name,
+                total_lines,
+            )
+            total_lines_edit = sum(1 for _ in edited_lines)
+            edited_view = view_lines(
+                edited_lines, cur_line, DEFAULT_WINDOW_SIZE, name, total_lines_edit
+            )
 
-def search_dir(search_term: str, dir_path: str) -> str:
-    """Searches for search_term in all files in a directory.
+            error_msg += f"\n[This is how your edit would have looked like if applied]\n{edited_view}\n\n[This is the original code before your edit]\n{original_view}"
+            return error_msg
 
-    Parameters:
-        search_term (str): The search term to look for.
-        dir_path (str): The directory path to search in, preferred absolute path.
-    """
+    artifacts[name] = "".join(edited_lines)
 
-    dir_path_p = Path(dir_path)
-    if not dir_path_p.exists():
-        return f"[Directory {dir_path} does not exist]"
-
-    matches = []
-    for file in dir_path_p.glob("**/*"):
-        if filter_file(file):
-            with open(file, "r") as f:
-                lines = f.readlines()
-                for i, line in enumerate(lines):
-                    if search_term in line:
-                        matches.append(f"{file}:{i}|{line.strip()}\n")
-    if not matches:
-        return f"[No matches found for {search_term} in {dir_path}]"
-    if len(matches) > 100:
-        return f"[More than {len(matches)} matches found for {search_term} in {dir_path}. Please narrow your search]"
-
-    return_str = f"[Found {len(matches)} matches for {search_term} in {dir_path}]\n"
-    for match in matches:
-        return_str += match
-
-    return_str += f"[End of matches for {search_term} in {dir_path}]"
-    return return_str
+    display({MimeType.APPLICATION_JSON: {"last_artifact": name}})
+    return open_code_artifact(artifacts, name, cur_line)
 
 
-def search_file(search_term: str, file_path: str) -> str:
-    """Searches the file for the given search term.
+def generate_vision_code(
+    artifacts: Artifacts, name: str, chat: str, media: List[str]
+) -> str:
+    """Generates python code to solve vision based tasks.
 
     Parameters:
-        search_term (str): The search term to look for.
-        file_path (str): The file path to search in, preferred absolute path.
-    """
-
-    file_path_p = Path(file_path)
-    if not file_path_p.exists():
-        return f"[File {file_path} does not exist]"
+        artifacts (Artifacts): The artifacts object to save the code to.
+        name (str): The name of the artifact to save the code to.
+        chat (str): The chat message from the user.
+        media (List[str]): The media files to use.
 
-    with open(file_path_p, "r") as f:
-        lines = f.readlines()
+    Returns:
+        str: The generated code.
 
-    search_results = []
-    for i, line in enumerate(lines):
-        if search_term in line:
-            search_results.append(f"{i}|{line.strip()}\n")
+    Examples
+    --------
+        >>> generate_vision_code(artifacts, "code.py", "Can you detect the dogs in this image?", ["image.jpg"])
+        from vision_agent.tools import load_image, owl_v2
+        def detect_dogs(image_path: str):
+            image = load_image(image_path)
+            dogs = owl_v2("dog", image)
+            return dogs
+    """
 
-    if not search_results:
-        return f"[No matches found for {search_term} in {file_path}]"
+    if ZMQ_PORT is not None:
+        agent = va.agent.VisionAgentCoder(
+            report_progress_callback=lambda inp: report_progress_callback(
+                int(ZMQ_PORT), inp
+            )
+        )
+    else:
+        agent = va.agent.VisionAgentCoder()
 
-    return_str = (
-        f"[Found {len(search_results)} matches for {search_term} in {file_path}]\n"
-    )
-    for result in search_results:
-        return_str += result
+    fixed_chat: List[Message] = [{"role": "user", "content": chat, "media": media}]
+    response = agent.chat_with_workflow(fixed_chat, test_multi_plan=True)
+    redisplay_results(response["test_result"])
+    code = response["code"]
+    artifacts[name] = code
+    code_lines = code.splitlines(keepends=True)
+    total_lines = len(code_lines)
 
-    return_str += f"[End of matches for {search_term} in {file_path}]"
-    return return_str
+    display({MimeType.APPLICATION_JSON: {"last_artifact": name}})
+    return view_lines(code_lines, 0, total_lines, name, total_lines)
 
 
-def find_file(file_name: str, dir_path: str = "./") -> str:
-    """Finds all files with the given name in the specified directory.
+def edit_vision_code(
+    artifacts: Artifacts, name: str, chat_history: List[str], media: List[str]
+) -> str:
+    """Edits python code to solve a vision based task.
 
     Parameters:
-        file_name (str): The file name to look for.
-        dir_path (str): The directory path to search in, preferred absolute path.
-    """
-
-    dir_path_p = Path(dir_path)
-    if not dir_path_p.exists():
-        return f"[Directory {dir_path} does not exist]"
-
-    files = list(dir_path_p.glob(f"**/*{file_name}*"))
-    files = [f for f in files if filter_file(f)]
-    if not files:
-        return f"[No files found in {dir_path} with name {file_name}]"
-
-    return_str = f"[Found {len(files)} matches for {file_name} in {dir_path}]\n"
-    for match in files:
-        return_str += str(match) + "\n"
+        artifacts (Artifacts): The artifacts object to save the code to.
+        name (str): The file path to the code.
+        chat_history (List[str]): The chat history to used to generate the code.
 
-    return_str += f"[End of matches for {file_name} in {dir_path}]"
-    return return_str
+    Returns:
+        str: The edited code.
 
+    Examples
+    --------
+        >>> edit_vision_code(
+        >>>     artifacts,
+        >>>     "code.py",
+        >>>     ["Can you detect the dogs in this image?", "Can you use a higher threshold?"],
+        >>>     ["dog.jpg"],
+        >>> )
+        from vision_agent.tools import load_image, owl_v2
+        def detect_dogs(image_path: str):
+            image = load_image(image_path)
+            dogs = owl_v2("dog", image, threshold=0.8)
+            return dogs
+    """
 
-def edit_file(file_path: str, start: int, end: int, content: str) -> str:
-    """Edits the file at the given path with the provided content. The content will be
-    inserted between the `start` and `end` line numbers. If the `start` and `end` are
-    the same, the content will be inserted at the `start` line number. If the `end` is
-    greater than the total number of lines in the file, the content will be inserted at
-    the end of the file. If the `start` or `end` are negative, the function will return
-    an error message.
+    agent = va.agent.VisionAgentCoder()
+    if name not in artifacts:
+        return f"[Artifact {name} does not exist]"
 
-    Parameters:
-        file_path (str): The file path to edit, preferred absolute path.
-        start (int): The line number to start the edit.
-        end (int): The line number to end the edit.
-        content (str): The content to insert.
-    """
-    file_path_p = Path(file_path)
-    if not file_path_p.exists():
-        return f"[File {file_path} does not exist]"
+    code = artifacts[name]
 
-    total_lines = sum(1 for _ in open(file_path_p))
-    if start < 0 or end < 0 or start > end or end > total_lines:
-        return "[Invalid line range]"
-    if start == end:
-        end += 1
+    # Append latest code to second to last message from assistant
+    fixed_chat_history: List[Message] = []
+    for i, chat in enumerate(chat_history):
+        if i == 0:
+            fixed_chat_history.append({"role": "user", "content": chat, "media": media})
+        elif i > 0 and i < len(chat_history) - 1:
+            fixed_chat_history.append({"role": "user", "content": chat})
+        elif i == len(chat_history) - 1:
+            fixed_chat_history.append({"role": "assistant", "content": code})
+            fixed_chat_history.append({"role": "user", "content": chat})
 
-    new_content_lines = content.splitlines(keepends=True)
-    new_content_lines = [
-        line if line.endswith("\n") else line + "\n" for line in new_content_lines
-    ]
-    with open(file_path_p, "r") as f:
-        lines = f.readlines()
-        edited_lines = lines[:start] + new_content_lines + lines[end:]
+    response = agent.chat_with_workflow(fixed_chat_history, test_multi_plan=False)
+    redisplay_results(response["test_result"])
+    code = response["code"]
+    artifacts[name] = code
+    code_lines = code.splitlines(keepends=True)
+    total_lines = len(code_lines)
 
-    cur_line = start + len(content.split("\n")) // 2
-    tmp_file = file_path_p.with_suffix(".tmp")
-    with open(tmp_file, "w") as f:
-        f.writelines(edited_lines)
-
-    process = subprocess.Popen(
-        [
-            "flake8",
-            "--isolated",
-            "--select=F821,F822,F831,E111,E112,E113,E999,E902",
-            tmp_file,
-        ],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        text=True,
-    )
-    stdout, _ = process.communicate()
-    tmp_file.unlink()
-    if stdout != "":
-        stdout = stdout.replace(tmp_file.name, file_path)
-        error_msg = "[Edit failed with the following status]\n" + stdout
-        original_view = view_lines(
-            lines,
-            start + ((end - start) // 2),
-            DEFAULT_WINDOW_SIZE,
-            file_path,
-            total_lines,
-        )
-        total_lines_edit = sum(1 for _ in edited_lines)
-        edited_view = view_lines(
-            edited_lines, cur_line, DEFAULT_WINDOW_SIZE, file_path, total_lines_edit
-        )
+    display({MimeType.APPLICATION_JSON: {"last_artifact": name}})
+    return view_lines(code_lines, 0, total_lines, name, total_lines)
 
-        error_msg += f"\n[This is how your edit would have looked like if applied]\n{edited_view}\n\n[This is the original code before your edit]\n{original_view}"
-        return error_msg
 
-    with open(file_path_p, "w") as f:
-        f.writelines(edited_lines)
+def write_media_artifact(artifacts: Artifacts, local_path: str) -> str:
+    """Writes a media file to the artifacts object.
 
-    return open_file(file_path, cur_line)
+    Parameters:
+        artifacts (Artifacts): The artifacts object to save the media to.
+        local_path (str): The local path to the media file.
+    """
+    with open(local_path, "rb") as f:
+        media = f.read()
+    artifacts[Path(local_path).name] = media
+    return f"[Media {Path(local_path).name} saved]"
 
 
 def get_tool_descriptions() -> str:
@@ -388,15 +400,11 @@ def get_tool_descriptions() -> str:
 META_TOOL_DOCSTRING = get_tool_documentation(
     [
         get_tool_descriptions,
+        open_code_artifact,
+        create_code_artifact,
+        edit_code_artifact,
         generate_vision_code,
         edit_vision_code,
-        open_file,
-        create_file,
-        scroll_up,
-        scroll_down,
-        edit_file,
-        search_dir,
-        search_file,
-        find_file,
+        write_media_artifact,
     ]
 )