From a49aa123e96125fc45e33839bd7e90a708088b70 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Tue, 24 Sep 2024 23:57:56 -0400 Subject: [PATCH] feat(agents-api): Add migration for adding forward_tool_calls to sessions Signed-off-by: Diwank Singh Tomer --- agents-api/agents_api/autogen/Sessions.py | 20 ++--- agents-api/agents_api/autogen/Tasks.py | 8 +- .../workflows/task_execution/__init__.py | 71 +++++++-------- ...727235852_add_forward_tool_calls_option.py | 87 +++++++++++++++++++ typespec/sessions/models.tsp | 4 +- typespec/tasks/steps.tsp | 4 +- 6 files changed, 142 insertions(+), 52 deletions(-) create mode 100644 agents-api/migrations/migrate_1727235852_add_forward_tool_calls_option.py diff --git a/agents-api/agents_api/autogen/Sessions.py b/agents-api/agents_api/autogen/Sessions.py index 5eb866134..75e5252cd 100644 --- a/agents-api/agents_api/autogen/Sessions.py +++ b/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. diff --git a/agents-api/agents_api/autogen/Tasks.py b/agents-api/agents_api/autogen/Tasks.py index 2b5318c0c..48dba4ad7 100644 --- a/agents-api/agents_api/autogen/Tasks.py +++ b/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. diff --git a/agents-api/agents_api/workflows/task_execution/__init__.py b/agents-api/agents_api/workflows/task_execution/__init__.py index cfc03e341..3c8197e29 100644 --- a/agents-api/agents_api/workflows/task_execution/__init__.py +++ b/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 diff --git a/agents-api/migrations/migrate_1727235852_add_forward_tool_calls_option.py b/agents-api/migrations/migrate_1727235852_add_forward_tool_calls_option.py new file mode 100644 index 000000000..ad1aab998 --- /dev/null +++ b/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)]) diff --git a/typespec/sessions/models.tsp b/typespec/sessions/models.tsp index 89678e926..dfbb6ea41 100644 --- a/typespec/sessions/models.tsp +++ b/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; diff --git a/typespec/tasks/steps.tsp b/typespec/tasks/steps.tsp index abfbbc5ed..3495def1b 100644 --- a/typespec/tasks/steps.tsp +++ b/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"> {