vision-agent 0.2.173__py3-none-any.whl → 0.2.174__py3-none-any.whl

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/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(

{vision_agent-0.2.173.dist-info → vision_agent-0.2.174.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.173
+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.173.dist-info → vision_agent-0.2.174.dist-info}/RECORD

@@ -3,7 +3,7 @@ vision_agent/agent/__init__.py,sha256=RRMPhH8mgm_pCtEKiVFSjJyDi4lCr4F7k05AhK01xl
 vision_agent/agent/agent.py,sha256=2cjIOxEuSJrqbfPXYoV0qER5ihXsPFCoEFJa4jpqan0,597
 vision_agent/agent/agent_utils.py,sha256=eSgg8CwWylX_erLTqTg2pVhEEgVkMLRrQfYRyJzI3so,5443
 vision_agent/agent/vision_agent.py,sha256=lEda43d-Ri68FIjmn-MPIgLs8_jMpyDVXslusQedhWA,26222
-vision_agent/agent/vision_agent_coder.py,sha256=aVkl0b9LKvy-auuHGYSag-ixYnue0iRQqD1PYLPBR-s,29312
+vision_agent/agent/vision_agent_coder.py,sha256=3Q1VWrN-BNUoSD4OAqKazvXkP2c04PXDYu2Z1f5dQb0,31960
 vision_agent/agent/vision_agent_coder_prompts.py,sha256=gPLVXQMNSzYnQYpNm0wlH_5FPkOTaFDV24bqzK3jQ40,12221
 vision_agent/agent/vision_agent_planner.py,sha256=mjmnXG9CvYf_ZA7ZJ3ri4H-2U_Km55gF1sZYRSOlxpY,19027
 vision_agent/agent/vision_agent_planner_prompts.py,sha256=JDARUzko2HZdxkBtcy6wuP9DCCmbqhK_gnVgrjr6l1k,6691
@@ -17,7 +17,7 @@ vision_agent/lmm/__init__.py,sha256=jyY1sJb_tYKg5-Wzs3p1lvwFkc-aUNZfMcLy3TOC4Zg,
 vision_agent/lmm/lmm.py,sha256=B5ClgwvbybVCWkf9opDMLjTtJZemUU4KUkQoRxGh43I,16787
 vision_agent/lmm/types.py,sha256=ZEXR_ptBL0ZwDMTDYkgxUCmSZFmBYPQd2jreNzr_8UY,221
 vision_agent/tools/__init__.py,sha256=u-vS5iORB4ccvxoAjbtpvhTALDhXGilcATIq1_eZhKo,2332
-vision_agent/tools/meta_tools.py,sha256=zcqp3POpGcS4cfceMih0AD1c6MwS_PFLLA6EjTXhonM,28013
+vision_agent/tools/meta_tools.py,sha256=SVevrA5yxtSIXR8352QMMJv-jW8MIuA68Nj93DjW-60,30640
 vision_agent/tools/prompts.py,sha256=V1z4YJLXZuUl_iZ5rY0M5hHc_2tmMEUKr0WocXKGt4E,1430
 vision_agent/tools/tool_utils.py,sha256=VPGqGJ2ZYEJA6AW7K9X7hQv6vRlMtAQcybE4izdToCw,8196
 vision_agent/tools/tools.py,sha256=iKsBZxJ5--xWK-mqgZ1jbX_bfGS5HmAp-VRZ69m9yPg,77921
@@ -29,7 +29,7 @@ vision_agent/utils/image_utils.py,sha256=rm9GfXvD4JrjnqKrP_f2gfq4SzmqYC0IdC1kKwd
 vision_agent/utils/sim.py,sha256=ZuSS07TUXFGjipmiQoY8TKRmSes7XXCdtU9PI8PC1sw,5609
 vision_agent/utils/type_defs.py,sha256=BE12s3JNQy36QvauXHjwyeffVh5enfcvd4vTzSwvEZI,1384
 vision_agent/utils/video.py,sha256=fOPR48-SuwMbE5eB5rc2F7lVo6k1mVHn26eEJ0QCslc,5602
-vision_agent-0.2.173.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vision_agent-0.2.173.dist-info/METADATA,sha256=xSZ3aKwFKG81BoBIHXLKN2CSqGFmmzMnalcbLPYrV7w,18342
-vision_agent-0.2.173.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
-vision_agent-0.2.173.dist-info/RECORD,,
+vision_agent-0.2.174.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vision_agent-0.2.174.dist-info/METADATA,sha256=V3d-gcpI2IZ4QWygiErgGAeNza-ROlsU-s1eH5Fr9UM,18339
+vision_agent-0.2.174.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+vision_agent-0.2.174.dist-info/RECORD,,