feat(afent-api): added inherit tool fix

julep-ai · Oct 31, 2024 · b8befd0 · b8befd0
1 parent 2b635cf
commit b8befd0
Showing 1 changed file with 28 additions and 23 deletions.
diff --git a/agents-api/agents_api/activities/task_steps/prompt_step.py b/agents-api/agents_api/activities/task_steps/prompt_step.py
@@ -37,9 +37,9 @@ def format_agent_tool(tool: Tool) -> dict:
         return {
             "type": tool.type,
             "name": tool.name,
-            "display_width_px": tool.computer_20241022.display_width_px,
-            "display_height_px": tool.computer_20241022.display_height_px,
-            "display_number": tool.computer_20241022.display_number,
+            "display_width_px": tool.spec["display_width_px"],
+            "display_height_px": tool.spec["display_height_px"],
+            "display_number": tool.spec["display_number"],
         }
     elif tool.type == "bash_20241022":
         return {
@@ -82,14 +82,31 @@ async def prompt_step(context: StepContext) -> StepOutcome:
         else "gpt-4o"
     )
 
-    agent_tools = list_tools(
-        developer_id=context.execution_input.developer_id,
-        agent_id=context.execution_input.agent.id,
-        limit=128,  # Max number of supported functions in OpenAI. See https://platform.openai.com/docs/api-reference/chat/create
-        offset=0,
-        sort_by="created_at",
-        direction="desc",
-    )
+    # Format agent_tools for litellm
+    # Initialize the formatted_agent_tools with context tools
+    task_tools = context.execution_input.task.tools
+    formatted_agent_tools = [format_agent_tool(tool) for tool in task_tools]
+    # Check if the task has inherit_tools set to True
+    inherited_tools = context.execution_input.task.inherit_tools
+
+    if inherited_tools:
+        agent_tools = list_tools(
+            developer_id=context.execution_input.developer_id,
+            agent_id=context.execution_input.agent.id,
+            limit=128,  # Max number of supported functions in OpenAI. See https://platform.openai.com/docs/api-reference/chat/create
+            offset=0,
+            sort_by="created_at",
+            direction="desc",
+        )
+
+        # Add agent_tools if they are not already in formatted_agent_tools
+        for agent_tool in agent_tools:
+            if (
+                format_agent_tool(agent_tool)
+                and format_agent_tool(agent_tool) not in formatted_agent_tools
+            ):
+                formatted_agent_tools.append(format_agent_tool(agent_tool))
+
     # grab the tools from context.current_step.tools and then append it to the agent_tools
     if context.current_step.settings:
         passed_settings: dict = context.current_step.settings.model_dump(
@@ -102,18 +119,6 @@ async def prompt_step(context: StepContext) -> StepOutcome:
     if isinstance(prompt, str):
         prompt = [{"role": "user", "content": prompt}]
 
-    # Format agent_tools for litellm
-    # Initialize the formatted_agent_tools with context tools
-    task_tools = context.tools
-    formatted_agent_tools = [format_agent_tool(tool) for tool in task_tools]
-    # Add agent_tools if they are not already in formatted_agent_tools
-    for agent_tool in agent_tools:
-        if (
-            format_agent_tool(agent_tool)
-            and format_agent_tool(agent_tool) not in formatted_agent_tools
-        ):
-            formatted_agent_tools.append(format_agent_tool(agent_tool))
-
     # Check if the model is Anthropic
     if agent_model.lower().startswith("claude-3.5") and any(
         tool["type"] in ["computer_20241022", "bash_20241022", "text_editor_20241022"]