vision-agent 0.2.172__tar.gz → 0.2.174__tar.gz

{vision_agent-0.2.172 → vision_agent-0.2.174}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.172
+Version: 0.2.174
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai
@@ -16,6 +16,7 @@ Requires-Dist: e2b-code-interpreter (==0.0.11a37)
 Requires-Dist: flake8 (>=7.0.0,<8.0.0)
 Requires-Dist: ipykernel (>=6.29.4,<7.0.0)
 Requires-Dist: langsmith (>=0.1.58,<0.2.0)
+Requires-Dist: libcst (>=1.5.0,<2.0.0)
 Requires-Dist: nbclient (>=0.10.0,<0.11.0)
 Requires-Dist: nbformat (>=5.10.4,<6.0.0)
 Requires-Dist: numpy (>=1.21.0,<2.0.0)
@@ -27,7 +28,6 @@ Requires-Dist: pillow-heif (>=0.16.0,<0.17.0)
 Requires-Dist: pydantic (==2.7.4)
 Requires-Dist: pydantic-settings (>=2.2.1,<3.0.0)
 Requires-Dist: pytube (==15.0.0)
-Requires-Dist: redbaron (>=0.9.2,<0.10.0)
 Requires-Dist: requests (>=2.0.0,<3.0.0)
 Requires-Dist: rich (>=13.7.1,<14.0.0)
 Requires-Dist: scipy (>=1.13.0,<1.14.0)

{vision_agent-0.2.172 → vision_agent-0.2.174}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "vision-agent"
-version = "0.2.172"
+version = "0.2.174"
 description = "Toolset for Vision Agent"
 authors = ["Landing AI <dev@landing.ai>"]
 readme = "README.md"
@@ -43,7 +43,7 @@ pytube = "15.0.0"
 anthropic = "^0.31.0"
 pydantic = "2.7.4"
 av = "^11.0.0"
-redbaron = "^0.9.2"
+libcst = "^1.5.0"
 
 [tool.poetry.group.dev.dependencies]
 autoflake = "1.*"

{vision_agent-0.2.172 → vision_agent-0.2.174}/vision_agent/agent/vision_agent.py

@@ -492,19 +492,10 @@ class VisionAgent(Agent):
                     }
                 )
 
+                code_action = response.get("execute_python", None)
                 # sometimes it gets stuck in a loop, so we force it to exit
                 if last_response == response:
                     response["let_user_respond"] = True
-
-                finished = response.get("let_user_respond", False)
-
-                code_action = response.get("execute_python", None)
-                if code_action is not None:
-                    code_action = use_extra_vision_agent_args(
-                        code_action, test_multi_plan, custom_tool_names
-                    )
-
-                if last_response == response:
                     self.streaming_message(
                         {
                             "role": "assistant",
@@ -514,7 +505,7 @@ class VisionAgent(Agent):
                                 "value": "Agent is stuck in conversation loop, exited",
                                 "traceback_raw": [],
                             },
-                            "finished": finished and code_action is None,
+                            "finished": code_action is None,
                         }
                     )
                 else:
@@ -524,10 +515,18 @@ class VisionAgent(Agent):
                             "content": new_format_to_old_format(
                                 add_step_descriptions(response)
                             ),
-                            "finished": finished and code_action is None,
+                            "finished": response.get("let_user_respond", False)
+                            and code_action is None,
                         }
                     )
 
+                finished = response.get("let_user_respond", False)
+
+                if code_action is not None:
+                    code_action = use_extra_vision_agent_args(
+                        code_action, test_multi_plan, custom_tool_names
+                    )
+
                 if code_action is not None:
                     result, obs = execute_code_action(
                         artifacts,

{vision_agent-0.2.172 → vision_agent-0.2.174}/vision_agent/agent/vision_agent_coder.py

@@ -5,7 +5,7 @@ import sys
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Sequence, Union, cast
 
-from redbaron import RedBaron  # type: ignore
+import libcst as cst
 from tabulate import tabulate
 
 import vision_agent.tools as T
@@ -49,42 +49,112 @@ WORKSPACE = Path(os.getenv("WORKSPACE", ""))
 _LOGGER = logging.getLogger(__name__)
 
 
-def strip_function_calls(code: str, exclusions: Optional[List[str]] = None) -> str:
+def strip_function_calls(  # noqa: C901
+    code: str, exclusions: Optional[List[str]] = None
+) -> str:
     """This will strip out all code that calls functions except for functions included
     in exclusions.
     """
     if exclusions is None:
         exclusions = []
 
-    red = RedBaron(code)
-    nodes_to_remove = []
-    for node in red:
-        if node.type == "def":
-            continue
-        elif node.type == "import" or node.type == "from_import":
-            continue
-        elif node.type == "call":
-            if node.value and node.value[0].value in exclusions:
-                continue
-            nodes_to_remove.append(node)
-        elif node.type == "atomtrailers":
-            if node[0].value in exclusions:
-                continue
-            nodes_to_remove.append(node)
-        elif node.type == "assignment":
-            if node.value.type == "call" or node.value.type == "atomtrailers":
-                func_name = node.value[0].value
-                if func_name in exclusions:
-                    continue
-                nodes_to_remove.append(node)
-        elif node.type == "endl":
-            continue
-        else:
-            nodes_to_remove.append(node)
-    for node in nodes_to_remove:
-        node.parent.remove(node)
-    cleaned_code = red.dumps().strip()
-    return cleaned_code if isinstance(cleaned_code, str) else code
+    def check_and_remove_node(node: cst.CSTNode, exclusions: List[str]) -> cst.CSTNode:
+        if hasattr(node, "value") and isinstance(node.value, cst.Call):
+            if (
+                isinstance(node.value.func, cst.Name)
+                and node.value.func.value in exclusions
+            ):
+                return node
+            return cst.RemoveFromParent()  # type: ignore
+        return node
+
+    class StripFunctionCallsTransformer(cst.CSTTransformer):
+        def __init__(self, exclusions: List[str]):
+            # Store exclusions to skip removing certain function calls
+            self.exclusions = exclusions
+            self.in_function_or_class = False
+
+        def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:
+            self.in_function_or_class = True
+            return True
+
+        def leave_FunctionDef(
+            self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
+        ) -> cst.BaseStatement:
+            self.in_function_or_class = False
+            return updated_node
+
+        def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:
+            self.in_function_or_class = True
+            return True
+
+        def leave_ClassDef(
+            self, node: cst.ClassDef, updated_node: cst.ClassDef
+        ) -> cst.BaseStatement:
+            self.in_function_or_class = False
+            return updated_node
+
+        def leave_Expr(
+            self, original_node: cst.Expr, updated_node: cst.Expr
+        ) -> cst.Expr:
+            if not self.in_function_or_class:
+                return cast(
+                    cst.Expr, check_and_remove_node(updated_node, self.exclusions)
+                )
+            return updated_node
+
+        def leave_Assign(
+            self, original_node: cst.Assign, updated_node: cst.Assign
+        ) -> cst.Assign:
+            if not self.in_function_or_class:
+                return cast(
+                    cst.Assign, check_and_remove_node(updated_node, self.exclusions)
+                )
+            return updated_node
+
+        def leave_If(self, original_node: cst.If, updated_node: cst.If) -> cst.If:
+            if not self.in_function_or_class:
+                return cast(
+                    cst.If, check_and_remove_node(updated_node, self.exclusions)
+                )
+            return updated_node
+
+        def leave_For(self, original_node: cst.For, updated_node: cst.For) -> cst.For:
+            if not self.in_function_or_class:
+                return cast(
+                    cst.For, check_and_remove_node(updated_node, self.exclusions)
+                )
+            return updated_node
+
+        def leave_While(
+            self, original_node: cst.While, updated_node: cst.While
+        ) -> cst.While:
+            if not self.in_function_or_class:
+                return cast(
+                    cst.While, check_and_remove_node(updated_node, self.exclusions)
+                )
+            return updated_node
+
+        def leave_With(
+            self, original_node: cst.With, updated_node: cst.With
+        ) -> cst.With:
+            if not self.in_function_or_class:
+                return cast(
+                    cst.With, check_and_remove_node(updated_node, self.exclusions)
+                )
+            return updated_node
+
+        def leave_Try(self, original_node: cst.Try, updated_node: cst.Try) -> cst.Try:
+            if not self.in_function_or_class:
+                return cast(
+                    cst.Try, check_and_remove_node(updated_node, self.exclusions)
+                )
+            return updated_node
+
+    tree = cst.parse_module(code)
+    transformer = StripFunctionCallsTransformer(exclusions)
+    modified_tree = tree.visit(transformer)
+    return modified_tree.code
 
 
 def write_code(

{vision_agent-0.2.172 → vision_agent-0.2.174}/vision_agent/tools/meta_tools.py

@@ -8,8 +8,8 @@ import tempfile
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
+import libcst as cst
 from IPython.display import display
-from redbaron import RedBaron  # type: ignore
 
 import vision_agent as va
 from vision_agent.agent.agent_utils import extract_json
@@ -663,21 +663,79 @@ def use_extra_vision_agent_args(
     Returns:
         str: The edited code.
     """
-    red = RedBaron(code)
-    for node in red:
-        # seems to always be atomtrailers not call type
-        if node.type == "atomtrailers":
-            if node.name.value == "generate_vision_code":
-                node.value[1].value.append(f"test_multi_plan={test_multi_plan}")
-
-            if (
-                node.name.value == "generate_vision_code"
-                or node.name.value == "edit_vision_code"
-            ):
-                if custom_tool_names is not None:
-                    node.value[1].value.append(f"custom_tool_names={custom_tool_names}")
-    cleaned_code = red.dumps().strip()
-    return cleaned_code if isinstance(cleaned_code, str) else code
+
+    class VisionAgentTransformer(cst.CSTTransformer):
+        def __init__(
+            self, test_multi_plan: bool, custom_tool_names: Optional[List[str]]
+        ):
+            self.test_multi_plan = test_multi_plan
+            self.custom_tool_names = custom_tool_names
+
+        def leave_Call(
+            self, original_node: cst.Call, updated_node: cst.Call
+        ) -> cst.Call:
+            # Check if the function being called is generate_vision_code or edit_vision_code
+            if isinstance(updated_node.func, cst.Name) and updated_node.func.value in [
+                "generate_vision_code",
+                "edit_vision_code",
+            ]:
+                # Add test_multi_plan argument to generate_vision_code calls
+                if updated_node.func.value == "generate_vision_code":
+                    new_arg = cst.Arg(
+                        keyword=cst.Name("test_multi_plan"),
+                        value=cst.Name(str(self.test_multi_plan)),
+                        equal=cst.AssignEqual(
+                            whitespace_before=cst.SimpleWhitespace(""),
+                            whitespace_after=cst.SimpleWhitespace(""),
+                        ),
+                    )
+                    updated_node = updated_node.with_changes(
+                        args=[*updated_node.args, new_arg]
+                    )
+
+                # Add custom_tool_names if provided
+                if self.custom_tool_names is not None:
+                    list_arg = []
+                    for i, tool_name in enumerate(self.custom_tool_names):
+                        if i < len(self.custom_tool_names) - 1:
+                            list_arg.append(
+                                cst._nodes.expression.Element(
+                                    value=cst.SimpleString(value=f'"{tool_name}"'),
+                                    comma=cst.Comma(
+                                        whitespace_before=cst.SimpleWhitespace(""),
+                                        whitespace_after=cst.SimpleWhitespace(" "),
+                                    ),
+                                )
+                            )
+                        else:
+                            list_arg.append(
+                                cst._nodes.expression.Element(
+                                    value=cst.SimpleString(value=f'"{tool_name}"'),
+                                )
+                            )
+                    new_arg = cst.Arg(
+                        keyword=cst.Name("custom_tool_names"),
+                        value=cst.List(list_arg),
+                        equal=cst.AssignEqual(
+                            whitespace_before=cst.SimpleWhitespace(""),
+                            whitespace_after=cst.SimpleWhitespace(""),
+                        ),
+                    )
+                    updated_node = updated_node.with_changes(
+                        args=[*updated_node.args, new_arg]
+                    )
+
+            return updated_node
+
+    # Parse the input code into a CST node
+    tree = cst.parse_module(code)
+
+    # Apply the transformer to modify the CST
+    transformer = VisionAgentTransformer(test_multi_plan, custom_tool_names)
+    modified_tree = tree.visit(transformer)
+
+    # Return the modified code as a string
+    return modified_tree.code
 
 
 def use_object_detection_fine_tuning(