wip(agents-api): Auto-run tools in prompt steps

Signed-off-by: Diwank Singh Tomer <[email protected]>
julep-ai · Oct 31, 2024 · 387e0aa · 387e0aa
1 parent b324093
commit 387e0aa
Show file tree

Hide file tree

Showing 12 changed files with 190 additions and 149 deletions.
diff --git a/agents-api/agents_api/autogen/Sessions.py b/agents-api/agents_api/autogen/Sessions.py
@@ -43,12 +43,10 @@ class CreateSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -80,12 +78,10 @@ class PatchSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -117,12 +113,10 @@ class Session(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -190,12 +184,10 @@ class UpdateSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -234,12 +226,10 @@ class CreateOrUpdateSessionRequest(CreateSessionRequest):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.

diff --git a/agents-api/agents_api/autogen/Tasks.py b/agents-api/agents_api/autogen/Tasks.py
@@ -702,12 +702,10 @@ class PromptStep(BaseModel):
     """
     Whether to unwrap the output of the prompt step, equivalent to `response.choices[0].message.content`
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = True
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: true for prompt steps, false for sessions)
 
     If a tool call is made, the tool's output will be used as the model's input.
     If a tool call is not made, the model's output will be used as the next step's input.
@@ -746,12 +744,10 @@ class PromptStepUpdateItem(BaseModel):
     """
     Whether to unwrap the output of the prompt step, equivalent to `response.choices[0].message.content`
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = True
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: true for prompt steps, false for sessions)
 
     If a tool call is made, the tool's output will be used as the model's input.
     If a tool call is not made, the model's output will be used as the next step's input.

diff --git a/agents-api/agents_api/models/session/create_or_update_session.py b/agents-api/agents_api/models/session/create_or_update_session.py
@@ -53,7 +53,7 @@ def create_or_update_session(
     data: CreateOrUpdateSessionRequest,
 ) -> tuple[list[str], dict]:
     data.metadata = data.metadata or {}
-    session_data = data.model_dump()
+    session_data = data.model_dump(exclude={"auto_run_tools"})
 
     user = session_data.pop("user")
     agent = session_data.pop("agent")

diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py
@@ -60,7 +60,7 @@ def create_session(
     session_id = session_id or uuid4()
 
     data.metadata = data.metadata or {}
-    session_data = data.model_dump()
+    session_data = data.model_dump(exclude={"auto_run_tools"})
 
     user = session_data.pop("user")
     agent = session_data.pop("agent")

diff --git a/agents-api/agents_api/workflows/task_execution/__init__.py b/agents-api/agents_api/workflows/task_execution/__init__.py
@@ -372,10 +372,23 @@ async def run(
                 state = PartialTransition(type="resume", output=result)
 
             case PromptStep(unwrap=True), StepOutcome(output=response):
+                finish_reason = response["choices"][0]["finish_reason"]
+                if finish_reason == "tool_calls":
+                    workflow.logger.error(
+                        "Prompt step: Tool calls not supported in unwrap mode"
+                    )
+
+                    state = PartialTransition(
+                        type="error", output="Tool calls not supported in unwrap mode"
+                    )
+                    await transition(context, state)
+
+                    raise ApplicationError("Tool calls not supported in unwrap mode")
+
                 workflow.logger.debug(f"Prompt step: Received response: {response}")
                 state = PartialTransition(output=response)
 
-            case PromptStep(forward_tool_results=False, unwrap=False), StepOutcome(
+            case PromptStep(auto_run_tools=False, unwrap=False), StepOutcome(
                 output=response
             ):
                 workflow.logger.debug(f"Prompt step: Received response: {response}")
@@ -387,12 +400,22 @@ async def run(
                 workflow.logger.debug(f"Prompt step: Received response: {response}")
                 state = PartialTransition(output=response)
 
-            case PromptStep(unwrap=False), StepOutcome(output=response) if response[
-                "choices"
-            ][0]["finish_reason"] == "tool_calls":
-                workflow.logger.debug("Prompt step: Received tool call")
-                message = response["choices"][0]["message"]
-                tool_calls_input = message["tool_calls"]
+            ## TODO: Handle multiple tool calls and multiple choices
+            # case PromptStep(unwrap=False), StepOutcome(output=response) if response[
+            #     "choices"
+            # ][0]["finish_reason"] == "tool_calls":
+            #     workflow.logger.debug("Prompt step: Received tool call")
+            #     message = response["choices"][0]["message"]
+            #     tool_calls_input = message["tool_calls"]
+
+            case PromptStep(auto_run_tools=True, unwrap=False), StepOutcome(
+                output=response
+            ) if (message := response["choices"][0])[
+                "finish_reason"
+            ] == "tool_calls" and (tool_calls_input := message["tool_calls"])[0][
+                "type"
+            ] == "function":
+                workflow.logger.debug("Prompt step: Received FUNCTION tool call")
 
                 # Enter a wait-for-input step to ask the developer to run the tool calls
                 tool_calls_results = await workflow.execute_activity(
@@ -415,6 +438,67 @@ async def run(
                 )
                 state = PartialTransition(output=new_response.output, type="resume")
 
+            case PromptStep(auto_run_tools=True, unwrap=False), StepOutcome(
+                output=response
+            ) if (message := response["choices"][0])[
+                "finish_reason"
+            ] == "tool_calls" and (tool_calls_input := message["tool_calls"])[0][
+                "type"
+            ] == "integration":
+                workflow.logger.debug("Prompt step: Received INTEGRATION tool call")
+
+                # FIXME: Implement integration tool calls
+                # See: MANUAL TOOL CALL INTEGRATION (below)
+                raise NotImplementedError("Integration tool calls not yet supported")
+
+                # TODO: Feed the tool call results back to the model (see above)
+
+            case PromptStep(auto_run_tools=True, unwrap=False), StepOutcome(
+                output=response
+            ) if (message := response["choices"][0])[
+                "finish_reason"
+            ] == "tool_calls" and (tool_calls_input := message["tool_calls"])[0][
+                "type"
+            ] == "api_call":
+                workflow.logger.debug("Prompt step: Received API_CALL tool call")
+
+                # FIXME: Implement API_CALL tool calls
+                # See: MANUAL TOOL CALL API_CALL (below)
+                raise NotImplementedError("API_CALL tool calls not yet supported")
+
+                # TODO: Feed the tool call results back to the model (see above)
+
+            case PromptStep(auto_run_tools=True, unwrap=False), StepOutcome(
+                output=response
+            ) if (message := response["choices"][0])[
+                "finish_reason"
+            ] == "tool_calls" and (tool_calls_input := message["tool_calls"])[0][
+                "type"
+            ] == "system":
+                workflow.logger.debug("Prompt step: Received SYSTEM tool call")
+
+                # FIXME: Implement SYSTEM tool calls
+                # See: MANUAL TOOL CALL SYSTEM (below)
+                raise NotImplementedError("SYSTEM tool calls not yet supported")
+
+                # TODO: Feed the tool call results back to the model (see above)
+
+            case PromptStep(unwrap=False), StepOutcome(output=response) if (
+                message := response["choices"][0]
+            )["finish_reason"] == "tool_calls" and (
+                tool_calls_input := message["tool_calls"]
+            )[0]["type"] not in ["function", "integration", "api_call", "system"]:
+                workflow.logger.debug(
+                    f"Prompt step: Received unknown tool call: {tool_calls_input[0]['type']}"
+                )
+                state = PartialTransition(output=response)
+
+            case SetStep(), StepOutcome(output=evaluated_output):
+                workflow.logger.info("Set step: Updating user state")
+
+            case SetStep(), StepOutcome(output=evaluated_output):
+                workflow.logger.info("Set step: Updating user state")
+
             case SetStep(), StepOutcome(output=evaluated_output):
                 workflow.logger.info("Set step: Updating user state")
 
@@ -452,6 +536,8 @@ async def run(
             case ToolCallStep(), StepOutcome(output=tool_call) if tool_call[
                 "type"
             ] == "integration":
+                # MANUAL TOOL CALL INTEGRATION
+                workflow.logger.debug("ToolCallStep: Received INTEGRATION tool call")
                 call = tool_call["integration"]
                 tool_name = call["name"]
                 arguments = call["arguments"]
@@ -490,6 +576,8 @@ async def run(
             case ToolCallStep(), StepOutcome(output=tool_call) if tool_call[
                 "type"
             ] == "api_call":
+                # MANUAL TOOL CALL API_CALL
+                workflow.logger.debug("ToolCallStep: Received API_CALL tool call")
                 call = tool_call["api_call"]
                 tool_name = call["name"]
                 arguments = call["arguments"]
@@ -528,6 +616,8 @@ async def run(
             case ToolCallStep(), StepOutcome(output=tool_call) if tool_call[
                 "type"
             ] == "system":
+                # MANUAL TOOL CALL SYSTEM
+                workflow.logger.debug("ToolCallStep: Received SYSTEM tool call")
                 call = tool_call.get("system")
 
                 system_call = SystemDef(**call)
@@ -545,12 +635,16 @@ async def run(
                 workflow.logger.error(
                     f"Unhandled step type: {type(context.current_step).__name__}"
                 )
+                state = PartialTransition(type="error", output="Not implemented")
+                await transition(context, state)
+
                 raise ApplicationError("Not implemented")
 
         # 4. Transition to the next step
         workflow.logger.info(f"Transitioning after step {context.cursor.step}")
 
         # The returned value is the transition finally created
+        state = state or PartialTransition(type="error", output="Not implemented")
         final_state = await transition(context, state)
 
         # ---

diff --git a/integrations-service/integrations/autogen/Sessions.py b/integrations-service/integrations/autogen/Sessions.py
@@ -43,12 +43,10 @@ class CreateSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -80,12 +78,10 @@ class PatchSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -117,12 +113,10 @@ class Session(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -190,12 +184,10 @@ class UpdateSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.
@@ -234,12 +226,10 @@ class CreateOrUpdateSessionRequest(CreateSessionRequest):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: StrictBool | None = None
+    auto_run_tools: StrictBool = False
     """
-    Whether to forward the tool results to the model when available.
-    "true" => always forward
-    "false" => never forward
-    null => forward if applicable (default)
+    Whether to auto-run the tool and send the tool results to the model when available.
+    (default: false for sessions, true for tasks)
 
     If a tool call is made, the tool's output will be sent back to the model as the model's input.
     If a tool call is not made, the model's output will be returned as is.