vision-agent 0.2.112__tar.gz → 0.2.113__tar.gz

{vision_agent-0.2.112 → vision_agent-0.2.113}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.112
+Version: 0.2.113
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai

{vision_agent-0.2.112 → vision_agent-0.2.113}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "vision-agent"
-version = "0.2.112"
+version = "0.2.113"
 description = "Toolset for Vision Agent"
 authors = ["Landing AI <dev@landing.ai>"]
 readme = "README.md"

{vision_agent-0.2.112 → vision_agent-0.2.113}/vision_agent/agent/vision_agent_coder.py

@@ -128,7 +128,11 @@ def write_plans(
 
     user_request = chat[-1]["content"]
     context = USER_REQ.format(user_request=user_request)
-    prompt = PLAN.format(context=context, tool_desc=tool_desc, feedback=working_memory)
+    prompt = PLAN.format(
+        context=context,
+        tool_desc=tool_desc,
+        feedback=working_memory,
+    )
     chat[-1]["content"] = prompt
     return extract_json(model(chat, stream=False))  # type: ignore
 
@@ -674,6 +678,7 @@ class VisionAgentCoder(Agent):
         chat: List[Message],
         test_multi_plan: bool = True,
         display_visualization: bool = False,
+        customized_tool_names: Optional[List[str]] = None,
     ) -> Dict[str, Any]:
         """Chat with VisionAgentCoder and return intermediate information regarding the
         task.
@@ -689,6 +694,8 @@ class VisionAgentCoder(Agent):
                 with the first plan.
             display_visualization (bool): If True, it opens a new window locally to
                 show the image(s) created by visualization code (if there is any).
+            customized_tool_names (List[str]): A list of customized tools for agent to pick and use.
+                If not provided, default to full tool set from vision_agent.tools.
 
         Returns:
             Dict[str, Any]: A dictionary containing the code, test, test result, plan,
@@ -742,7 +749,9 @@ class VisionAgentCoder(Agent):
             )
             plans = write_plans(
                 int_chat,
-                T.TOOL_DESCRIPTIONS,
+                T.get_tool_descriptions_by_names(
+                    customized_tool_names, T.FUNCTION_TOOLS, T.UTIL_TOOLS  # type: ignore
+                ),
                 format_memory(working_memory),
                 self.planner,
             )
@@ -754,7 +763,6 @@ class VisionAgentCoder(Agent):
                     _LOGGER.info(
                         f"\n{tabulate(tabular_data=p_fixed, headers='keys', tablefmt='mixed_grid', maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"
                     )
-
             tool_infos = retrieve_tools(
                 plans,
                 self.tool_recommender,

{vision_agent-0.2.112 → vision_agent-0.2.113}/vision_agent/tools/__init__.py

@@ -1,15 +1,16 @@
 from typing import Callable, List, Optional
 
-from .meta_tools import (
-    META_TOOL_DOCSTRING,
-)
+from .meta_tools import META_TOOL_DOCSTRING
 from .prompts import CHOOSE_PARAMS, SYSTEM_PROMPT
+from .tool_utils import get_tool_descriptions_by_names
 from .tools import (
+    FUNCTION_TOOLS,
     TOOL_DESCRIPTIONS,
     TOOL_DOCSTRING,
     TOOLS,
     TOOLS_DF,
     TOOLS_INFO,
+    UTIL_TOOLS,
     UTILITIES_DOCSTRING,
     blip_image_caption,
     clip,

{vision_agent-0.2.112 → vision_agent-0.2.113}/vision_agent/tools/tool_utils.py

@@ -142,6 +142,31 @@ def get_tool_descriptions(funcs: List[Callable[..., Any]]) -> str:
     return descriptions
 
 
+def get_tool_descriptions_by_names(
+    tool_name: Optional[List[str]],
+    funcs: List[Callable[..., Any]],
+    util_funcs: List[
+        Callable[..., Any]
+    ],  # util_funcs will always be added to the list of functions
+) -> str:
+    if tool_name is None:
+        return get_tool_descriptions(funcs + util_funcs)
+
+    invalid_names = [
+        name for name in tool_name if name not in {func.__name__ for func in funcs}
+    ]
+
+    if invalid_names:
+        raise ValueError(f"Invalid customized tool names: {', '.join(invalid_names)}")
+
+    filtered_funcs = (
+        funcs
+        if not tool_name
+        else [func for func in funcs if func.__name__ in tool_name]
+    )
+    return get_tool_descriptions(filtered_funcs + util_funcs)
+
+
 def get_tools_df(funcs: List[Callable[..., Any]]) -> pd.DataFrame:
     data: Dict[str, List[str]] = {"desc": [], "doc": []}
 

{vision_agent-0.2.112 → vision_agent-0.2.113}/vision_agent/tools/tools.py

@@ -1650,7 +1650,7 @@ def florencev2_fine_tuned_object_detection(
     return return_data
 
 
-TOOLS = [
+FUNCTION_TOOLS = [
     owl_v2,
     extract_frames,
     ocr,
@@ -1671,6 +1671,9 @@ TOOLS = [
     generate_pose_image,
     closest_mask_distance,
     closest_box_distance,
+]
+
+UTIL_TOOLS = [
     save_json,
     load_image,
     save_image,
@@ -1679,6 +1682,9 @@ TOOLS = [
     overlay_segmentation_masks,
     overlay_heat_map,
 ]
+
+TOOLS = FUNCTION_TOOLS + UTIL_TOOLS
+
 TOOLS_DF = get_tools_df(TOOLS)  # type: ignore
 TOOL_DESCRIPTIONS = get_tool_descriptions(TOOLS)  # type: ignore
 TOOL_DOCSTRING = get_tool_documentation(TOOLS)  # type: ignore