feat(agents-api): bug fixes for the prompt step and types of tool cal…

…l fixed
julep-ai · Oct 31, 2024 · 2d83db3 · 2d83db3
1 parent 74982fa
commit 2d83db3
Show file tree

Hide file tree

Showing 7 changed files with 161 additions and 24 deletions.
diff --git a/agents-api/agents_api/activities/execute_system.py b/agents-api/agents_api/activities/execute_system.py
@@ -225,8 +225,9 @@ async def execute_system(
                 return create_session_query(**arguments)
             elif system.operation == "update":
                 return update_session_query(**arguments)
-            elif system.operation == "delete":
-                return update_session_query(**arguments)
+            # FIXME: @hamada needs to be fixed
+            # elif system.operation == "delete":
+            #     return update_session_query(**arguments)
             elif system.operation == "delete":
                 return delete_session_query(**arguments)
         # TASKS

diff --git a/agents-api/agents_api/activities/task_steps/prompt_step.py b/agents-api/agents_api/activities/task_steps/prompt_step.py
@@ -1,6 +1,6 @@
 import os
 
-from anthropic import Anthropic  # Import Anthropic client
+from anthropic import AsyncAnthropic, Anthropic  # Import AsyncAnthropic client
 from beartype import beartype
 from temporalio import activity
 from temporalio.exceptions import ApplicationError
@@ -9,6 +9,7 @@
 from ...clients import (
     litellm,  # We dont directly import `acompletion` so we can mock it
 )
+from litellm.utils import CustomStreamWrapper, ModelResponse
 from ...common.protocol.tasks import StepContext, StepOutcome
 from ...common.storage_handler import auto_blob_store
 from ...common.utils.template import render_template
@@ -20,7 +21,7 @@
 
 # FIXME: This shouldn't be here.
 def format_agent_tool(tool: Tool) -> dict:
-    if tool.function:
+    if tool.type == "function":
         return {
             "type": "function",
             "function": {
@@ -29,20 +30,20 @@ def format_agent_tool(tool: Tool) -> dict:
                 "parameters": tool.function.parameters,
             },
         }
-    elif tool.computer_20241022:
+    elif tool.type == "computer_20241022":
         return {
             "type": tool.type,
             "name": tool.name,
-            "display_width_px": tool.display_width_px,
-            "display_height_px": tool.display_height_px,
-            "display_number": tool.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.bash_20241022:
+    elif tool.type == "bash_20241022":
         return {
             "type": tool.type,
             "name": tool.name,
         }
-    elif tool.text_editor_20241022:
+    elif tool.type == "text_editor_20241022":
         return {
             "type": tool.type,
             "name": tool.name,
@@ -51,7 +52,6 @@ def format_agent_tool(tool: Tool) -> dict:
     else:
         return {}
 
-
 @activity.defn
 @auto_blob_store
 @beartype
@@ -86,7 +86,7 @@ async def prompt_step(context: StepContext) -> StepOutcome:
         sort_by="created_at",
         direction="desc",
     )
-
+    # 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(
             exclude_unset=True
@@ -99,25 +99,38 @@ async def prompt_step(context: StepContext) -> StepOutcome:
         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 agent_tools if format_agent_tool(tool)
+        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 "claude-3.5-sonnet-20241022" == agent_model.lower():
+    if (agent_model.lower().startswith("claude-3.5")
+        and any(
+            tool["type"] in ["computer_20241022", "bash_20241022", "text_editor_20241022"]
+            for tool in formatted_agent_tools)
+        ):
         # Retrieve the API key from the environment variable
         betas = [COMPUTER_USE_BETA_FLAG]
         # Use Anthropic API directly
-        client = Anthropic(api_key=anthropic_api_key)
-
+        client = AsyncAnthropic(api_key=anthropic_api_key)
+        new_prompt = [{"role": "user", "content": prompt[0]["content"]}]
         # Claude Response
         response = await client.beta.messages.create(
-            model=agent_model,
-            messages=prompt,
+            model="claude-3-5-sonnet-20241022",
+            messages=new_prompt,
             tools=formatted_agent_tools,
             max_tokens=1024,
             betas=betas,
         )
+        print("-" * 100)
+        # print("Model Response: ", model_response)
+        print("Response: ", response)
 
         return StepOutcome(
             output=response.model_dump()

diff --git a/agents-api/agents_api/autogen/Tools.py b/agents-api/agents_api/autogen/Tools.py
@@ -687,6 +687,18 @@ class CreateToolRequest(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: Literal[
+        "function",
+        "integration",
+        "system",
+        "api_call",
+        "computer_20241022",
+        "text_editor_20241022",
+        "bash_20241022",
+    ]
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool
@@ -953,6 +965,21 @@ class PatchToolRequest(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: (
+        Literal[
+            "function",
+            "integration",
+            "system",
+            "api_call",
+            "computer_20241022",
+            "text_editor_20241022",
+            "bash_20241022",
+        ]
+        | None
+    ) = None
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool
@@ -1380,6 +1407,18 @@ class Tool(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: Literal[
+        "function",
+        "integration",
+        "system",
+        "api_call",
+        "computer_20241022",
+        "text_editor_20241022",
+        "bash_20241022",
+    ]
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool
@@ -1457,6 +1496,18 @@ class UpdateToolRequest(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: Literal[
+        "function",
+        "integration",
+        "system",
+        "api_call",
+        "computer_20241022",
+        "text_editor_20241022",
+        "bash_20241022",
+    ]
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool

diff --git a/agents-api/agents_api/common/protocol/tasks.py b/agents-api/agents_api/common/protocol/tasks.py
@@ -235,13 +235,12 @@ def task_to_spec(
     for tool in task.tools:
         tool_spec = getattr(tool, tool.type)
 
-        tools.append(
-            TaskToolDef(
-                type=tool.type,
-                spec=tool_spec.model_dump(),
-                **tool.model_dump(exclude={"type"}),
-            )
+        tool_obj = dict(
+            type=tool.type,
+            spec=tool_spec.model_dump(),
+            **tool.model_dump(exclude={"type"}),
         )
+        tools.append(TaskToolDef(**tool_obj))
 
     workflows = [Workflow(name="main", steps=task_data.pop("main"))]
 

diff --git a/integrations-service/integrations/autogen/Tools.py b/integrations-service/integrations/autogen/Tools.py
@@ -687,6 +687,18 @@ class CreateToolRequest(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: Literal[
+        "function",
+        "integration",
+        "system",
+        "api_call",
+        "computer_20241022",
+        "text_editor_20241022",
+        "bash_20241022",
+    ]
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool
@@ -953,6 +965,21 @@ class PatchToolRequest(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: (
+        Literal[
+            "function",
+            "integration",
+            "system",
+            "api_call",
+            "computer_20241022",
+            "text_editor_20241022",
+            "bash_20241022",
+        ]
+        | None
+    ) = None
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool
@@ -1380,6 +1407,18 @@ class Tool(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: Literal[
+        "function",
+        "integration",
+        "system",
+        "api_call",
+        "computer_20241022",
+        "text_editor_20241022",
+        "bash_20241022",
+    ]
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool
@@ -1457,6 +1496,18 @@ class UpdateToolRequest(BaseModel):
     """
     Name of the tool (must be unique for this agent and a valid python identifier string )
     """
+    type: Literal[
+        "function",
+        "integration",
+        "system",
+        "api_call",
+        "computer_20241022",
+        "text_editor_20241022",
+        "bash_20241022",
+    ]
+    """
+    Type of the tool
+    """
     description: str | None = None
     """
     Description of the tool

diff --git a/typespec/tools/models.tsp b/typespec/tools/models.tsp
@@ -210,6 +210,9 @@ model Tool {
     /** Name of the tool (must be unique for this agent and a valid python identifier string )*/
     name: validPythonIdentifier;
 
+    /** Type of the tool */
+    type: ToolType;
+
     /** Description of the tool */
     description?: string;
 

diff --git a/typespec/tsp-output/@typespec/openapi3/openapi-1.0.0.yaml b/typespec/tsp-output/@typespec/openapi3/openapi-1.0.0.yaml
@@ -6656,11 +6656,16 @@ components:
       type: object
       required:
         - name
+        - type
       properties:
         name:
           allOf:
             - $ref: '#/components/schemas/Common.validPythonIdentifier'
           description: Name of the tool (must be unique for this agent and a valid python identifier string )
+        type:
+          allOf:
+            - $ref: '#/components/schemas/Tools.ToolType'
+          description: Type of the tool
         description:
           type: string
           description: Description of the tool
@@ -6885,6 +6890,10 @@ components:
           allOf:
             - $ref: '#/components/schemas/Common.validPythonIdentifier'
           description: Name of the tool (must be unique for this agent and a valid python identifier string )
+        type:
+          allOf:
+            - $ref: '#/components/schemas/Tools.ToolType'
+          description: Type of the tool
         description:
           type: string
           description: Description of the tool
@@ -7291,6 +7300,7 @@ components:
       type: object
       required:
         - name
+        - type
         - created_at
         - updated_at
         - id
@@ -7299,6 +7309,10 @@ components:
           allOf:
             - $ref: '#/components/schemas/Common.validPythonIdentifier'
           description: Name of the tool (must be unique for this agent and a valid python identifier string )
+        type:
+          allOf:
+            - $ref: '#/components/schemas/Tools.ToolType'
+          description: Type of the tool
         description:
           type: string
           description: Description of the tool
@@ -7380,11 +7394,16 @@ components:
       type: object
       required:
         - name
+        - type
       properties:
         name:
           allOf:
             - $ref: '#/components/schemas/Common.validPythonIdentifier'
           description: Name of the tool (must be unique for this agent and a valid python identifier string )
+        type:
+          allOf:
+            - $ref: '#/components/schemas/Tools.ToolType'
+          description: Type of the tool
         description:
           type: string
           description: Description of the tool