feat(agents-api): Add migration for adding forward_tool_calls to sess…

…ions

Signed-off-by: Diwank Singh Tomer <[email protected]>

agents-api/agents_api/autogen/Sessions.py

@@ -43,12 +43,12 @@ class CreateSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: Literal["auto"] | StrictBool = "auto"
+    forward_tool_results: StrictBool | None = None
     """
     Whether to forward the tool results to the model when available.
     "true" => always forward
     "false" => never forward
-    "auto" => forward if applicable
+    null => forward if applicable (default)
 
     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 +80,12 @@ class PatchSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: Literal["auto"] | StrictBool = "auto"
+    forward_tool_results: StrictBool | None = None
     """
     Whether to forward the tool results to the model when available.
     "true" => always forward
     "false" => never forward
-    "auto" => forward if applicable
+    null => forward if applicable (default)
 
     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 +117,12 @@ class Session(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: Literal["auto"] | StrictBool = "auto"
+    forward_tool_results: StrictBool | None = None
     """
     Whether to forward the tool results to the model when available.
     "true" => always forward
     "false" => never forward
-    "auto" => forward if applicable
+    null => forward if applicable (default)
 
     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 +190,12 @@ class UpdateSessionRequest(BaseModel):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: Literal["auto"] | StrictBool = "auto"
+    forward_tool_results: StrictBool | None = None
     """
     Whether to forward the tool results to the model when available.
     "true" => always forward
     "false" => never forward
-    "auto" => forward if applicable
+    null => forward if applicable (default)
 
     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 +234,12 @@ class CreateOrUpdateSessionRequest(CreateSessionRequest):
     """
     Action to start on context window overflow
     """
-    forward_tool_results: Literal["auto"] | StrictBool = "auto"
+    forward_tool_results: StrictBool | None = None
     """
     Whether to forward the tool results to the model when available.
     "true" => always forward
     "false" => never forward
-    "auto" => forward if applicable
+    null => forward if applicable (default)
 
     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.

agents-api/agents_api/autogen/Tasks.py

@@ -688,12 +688,12 @@ class PromptStep(BaseModel):
     """
     Whether to unwrap the output of the prompt step, equivalent to `response.choices[0].message.content`
     """
-    forward_tool_results: Literal["auto"] | StrictBool = "auto"
+    forward_tool_results: StrictBool | None = None
     """
     Whether to forward the tool results to the model when available.
     "true" => always forward
     "false" => never forward
-    "auto" => forward if applicable
+    null => forward if applicable (default)
 
     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.
@@ -728,12 +728,12 @@ class PromptStepUpdateItem(BaseModel):
     """
     Whether to unwrap the output of the prompt step, equivalent to `response.choices[0].message.content`
     """
-    forward_tool_results: Literal["auto"] | StrictBool = "auto"
+    forward_tool_results: StrictBool | None = None
     """
     Whether to forward the tool results to the model when available.
     "true" => always forward
     "false" => never forward
-    "auto" => forward if applicable
+    null => forward if applicable (default)
 
     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.

agents-api/agents_api/workflows/task_execution/__init__.py

@@ -390,45 +390,48 @@ async def run(
 
                 state = PartialTransition(type="resume", output=result)
 
-            case PromptStep(
-                forward_tool_results=forward_tool_results, unwrap=False
-            ), StepOutcome(output=response):
+            case PromptStep(unwrap=True), StepOutcome(output=response):
                 workflow.logger.debug(f"Prompt step: Received response: {response}")
+                state = PartialTransition(output=response)
 
-                if (
-                    response["choices"][0]["finish_reason"] != "tool_calls"
-                    or not forward_tool_results
-                ):
-                    workflow.logger.debug("Prompt step: Received response")
-                    state = PartialTransition(output=response)
-                else:
-                    workflow.logger.debug("Prompt step: Received tool call")
-                    message = response["choices"][0]["message"]
-                    tool_calls_input = message["tool_calls"]
-
-                    # Enter a wait-for-input step to ask the developer to run the tool calls
-                    tool_calls_results = await workflow.execute_activity(
-                        task_steps.raise_complete_async,
-                        args=[context, tool_calls_input],
-                        schedule_to_close_timeout=timedelta(days=31),
-                    )
-
-                    # Feed the tool call results back to the model
-                    context.current_step.prompt.append(message)
-                    context.current_step.prompt.append(tool_calls_results)
-                    new_response = await workflow.execute_activity(
-                        task_steps.prompt_step,
-                        context,
-                        schedule_to_close_timeout=timedelta(
-                            seconds=30 if debug or testing else 600
-                        ),
-                    )
-                    state = PartialTransition(output=new_response.output, type="resume")
+            case PromptStep(forward_tool_results=False, unwrap=False), StepOutcome(
+                output=response
+            ):
+                workflow.logger.debug(f"Prompt step: Received response: {response}")
+                state = PartialTransition(output=response)
 
-            case PromptStep(unwrap=True), StepOutcome(output=response):
+            case PromptStep(unwrap=False), StepOutcome(output=response) if response[
+                "choices"
+            ][0]["finish_reason"] != "tool_calls":
                 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"]
+
+                # Enter a wait-for-input step to ask the developer to run the tool calls
+                tool_calls_results = await workflow.execute_activity(
+                    task_steps.raise_complete_async,
+                    args=[context, tool_calls_input],
+                    schedule_to_close_timeout=timedelta(days=31),
+                )
+
+                # Feed the tool call results back to the model
+                context.current_step.prompt.append(message)
+                context.current_step.prompt.append(tool_calls_results)
+                new_response = await workflow.execute_activity(
+                    task_steps.prompt_step,
+                    context,
+                    schedule_to_close_timeout=timedelta(
+                        seconds=30 if debug or testing else 600
+                    ),
+                )
+                state = PartialTransition(output=new_response.output, type="resume")
+
             case SetStep(), StepOutcome(output=evaluated_output):
                 workflow.logger.info("Set step: Updating user state")
                 self.update_user_state(evaluated_output)
@@ -494,7 +497,7 @@ async def run(
                     ),
                 )
 
-                state = PartialTransition(output=tool_call_response, type="step")
+                state = PartialTransition(output=tool_call_response)
 
             case ToolCallStep(), StepOutcome(output=_):
                 # FIXME: Handle system/api_call tool_calls

agents-api/migrations/migrate_1727235852_add_forward_tool_calls_option.py

@@ -0,0 +1,87 @@
+#/usr/bin/env python3
+
+MIGRATION_ID = "add_forward_tool_calls_option"
+CREATED_AT = 1727235852.744035
+
+
+def run(client, queries):
+    joiner = "}\n\n{"
+
+    query = joiner.join(queries)
+    query = f"{{\n{query}\n}}"
+    client.run(query)
+
+
+add_forward_tool_calls_option_to_session_query = dict(
+    up="""
+    ?[forward_tool_calls, token_budget, context_overflow, developer_id, session_id, updated_at, situation, summary, created_at, metadata, render_templates, token_budget, context_overflow] := *sessions{
+        developer_id,
+        session_id,
+        updated_at,
+        situation,
+        summary,
+        created_at,
+        metadata,
+        render_templates,
+        token_budget,
+        context_overflow,
+    },
+    forward_tool_calls = null
+
+    :replace sessions {
+        developer_id: Uuid,
+        session_id: Uuid,
+        updated_at: Validity default [floor(now()), true],
+        =>
+        situation: String,
+        summary: String? default null,
+        created_at: Float default now(),
+        metadata: Json default {},
+        render_templates: Bool default false,
+        token_budget: Int? default null,
+        context_overflow: String? default null,
+        forward_tool_calls: Bool? default null,
+    }
+     """,
+    down="""
+    ?[token_budget, context_overflow, developer_id, session_id, updated_at, situation, summary, created_at, metadata, render_templates, token_budget, context_overflow] := *sessions{
+        developer_id,
+        session_id,
+        updated_at,
+        situation,
+        summary,
+        created_at,
+        metadata,
+        render_templates,
+        token_budget,
+        context_overflow,
+    }
+
+    :replace sessions {
+        developer_id: Uuid,
+        session_id: Uuid,
+        updated_at: Validity default [floor(now()), true],
+        =>
+        situation: String,
+        summary: String? default null,
+        created_at: Float default now(),
+        metadata: Json default {},
+        render_templates: Bool default false,
+        token_budget: Int? default null,
+        context_overflow: String? default null,
+    }
+    """,
+)
+
+
+queries = [
+    add_forward_tool_calls_option_to_session_query,
+]
+
+
+def up(client):
+    run(client, [q["up"] for q in queries])
+
+
+def down(client):
+    run(client, [q["down"] for q in reversed(queries)])

typespec/sessions/models.tsp

@@ -65,11 +65,11 @@ model Session {
     /** Whether to forward the tool results to the model when available.
      * "true" => always forward
      * "false" => never forward
-     * "auto" => forward if applicable
+     * null => forward if applicable (default)
      * 
      * 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. */
-    forward_tool_results: "auto" | boolean = "auto";
+    forward_tool_results: boolean | null = null;
 
     ...HasId;
     ...HasMetadata;

typespec/tasks/steps.tsp

@@ -110,11 +110,11 @@ model PromptStepDef {
     /** Whether to forward the tool results to the model when available.
      * "true" => always forward
      * "false" => never forward
-     * "auto" => forward if applicable
+     * null => forward if applicable (default)
      * 
      * 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. */
-    forward_tool_results: "auto" | boolean = "auto";
+    forward_tool_results: boolean | null = null;
 }
 
 model EvaluateStep extends BaseWorkflowStep<"evaluate"> {