diff --git a/portia/__init__.py b/portia/__init__.py index 4a2ce014a..b91e985ad 100644 --- a/portia/__init__.py +++ b/portia/__init__.py @@ -91,7 +91,7 @@ # Plan and execution related classes from portia.plan import Plan, PlanBuilder, PlanContext, PlanInput, PlanUUID, Step, Variable -from portia.plan_run import PlanRun, PlanRunState +from portia.plan_run import PlanRun, PlanRunState, PlanRunV2, StepOutputValue # Core classes from portia.portia import ExecutionHooks, Portia @@ -163,6 +163,7 @@ "PlanRun", "PlanRunNotFoundError", "PlanRunState", + "PlanRunV2", "PlanUUID", "PlanV2", "PlanningAgentType", @@ -176,6 +177,7 @@ "StdioMcpClientConfig", "Step", "StepOutput", + "StepOutputValue", "StepV2", "StorageClass", "StorageError", diff --git a/portia/execution_hooks.py b/portia/execution_hooks.py index 5b8d2fa13..734ddf1d7 100644 --- a/portia/execution_hooks.py +++ b/portia/execution_hooks.py @@ -8,6 +8,8 @@ from pydantic import BaseModel, ConfigDict +from portia.builder.plan_v2 import PlanV2 +from portia.builder.step_v2 import StepV2 from portia.clarification import ( Clarification, ClarificationCategory, @@ -18,8 +20,7 @@ from portia.errors import ToolHardError from portia.execution_agents.output import Output from portia.logger import logger -from portia.plan import Plan, Step -from portia.plan_run import PlanRun +from portia.plan_run import PlanRunV2 from portia.tool import Tool @@ -53,53 +54,53 @@ class ExecutionHooks(BaseModel): clarification_handler: ClarificationHandler | None = None """Handler for clarifications raised during execution.""" - before_step_execution: Callable[[Plan, PlanRun, Step], BeforeStepExecutionOutcome] | None = None + before_step_execution: Callable[[PlanV2, PlanRunV2, StepV2], BeforeStepExecutionOutcome] | None = None """Called before executing each step. Args: - plan: The plan being executed - plan_run: The current plan run - step: The step about to be executed + plan: The PlanV2 being executed + plan_run: The current PlanRunV2 + step: The StepV2 about to be executed Returns: BeforeStepExecutionOutcome | None: Whether to continue with the step execution or skip it. If None is returned, the default behaviour is to continue with the step execution. """ - after_step_execution: Callable[[Plan, PlanRun, Step, Output], None] | None = None + after_step_execution: Callable[[PlanV2, PlanRunV2, StepV2, Output], None] | None = None """Called after executing each step. When there's an error, this is called with the error as the output value. Args: - plan: The plan being executed - plan_run: The current plan run - step: The step that was executed + plan: The PlanV2 being executed + plan_run: The current PlanRunV2 + step: The StepV2 that was executed output: The output from the step execution """ - before_plan_run: Callable[[Plan, PlanRun], None] | None = None + before_plan_run: Callable[[PlanV2, PlanRunV2], None] | None = None """Called before executing the first step of the plan run. Args: - plan: The plan being executed - plan_run: The current plan run + plan: The PlanV2 being executed + plan_run: The current PlanRunV2 """ - after_plan_run: Callable[[Plan, PlanRun, Output], None] | None = None + after_plan_run: Callable[[PlanV2, PlanRunV2, Output], None] | None = None """Called after executing the plan run. This is not called if a clarification is raised, as it is expected that the plan will be resumed after the clarification is handled. Args: - plan: The plan that was executed - plan_run: The completed plan run + plan: The PlanV2 that was executed + plan_run: The completed PlanRunV2 output: The final output from the plan execution """ before_tool_call: ( - Callable[[Tool, dict[str, Any], PlanRun, Step], Clarification | None] | None + Callable[[Tool, dict[str, Any], PlanRunV2, StepV2], Clarification | None] | None ) = None """Called before the tool is called. @@ -107,21 +108,21 @@ class ExecutionHooks(BaseModel): tool: The tool about to be called args: The args for the tool call. These are mutable and so can be modified in place as required. - plan_run: The current plan run - step: The step being executed + plan_run: The current PlanRunV2 + step: The StepV2 being executed Returns: Clarification | None: A clarification to raise, or None to proceed with the tool call """ - after_tool_call: Callable[[Tool, Any, PlanRun, Step], Clarification | None] | None = None + after_tool_call: Callable[[Tool, Any, PlanRunV2, StepV2], Clarification | None] | None = None """Called after the tool is called. Args: tool: The tool that was called output: The output returned from the tool call - plan_run: The current plan run - step: The step being executed + plan_run: The current PlanRunV2 + step: The StepV2 being executed Returns: Clarification | None: A clarification to raise, or None to proceed. If a clarification @@ -135,8 +136,8 @@ class ExecutionHooks(BaseModel): def clarify_on_all_tool_calls( tool: Tool, args: dict[str, Any], - plan_run: PlanRun, - step: Step, + plan_run: PlanRunV2, + step: StepV2, ) -> Clarification | None: """Raise a clarification to check the user is happy with all tool calls before proceeding. @@ -152,7 +153,7 @@ def clarify_on_all_tool_calls( def clarify_on_tool_calls( tool: str | Tool | list[str] | list[Tool], -) -> Callable[[Tool, dict[str, Any], PlanRun, Step], Clarification | None]: +) -> Callable[[Tool, dict[str, Any], PlanRunV2, StepV2], Clarification | None]: """Return a hook that raises a clarification before calls to the specified tool. Args: @@ -185,8 +186,8 @@ def clarify_on_tool_calls( def _clarify_on_tool_call_hook( tool: Tool, args: dict[str, Any], - plan_run: PlanRun, - step: Step, # noqa: ARG001 + plan_run: PlanRunV2, + step: StepV2, # noqa: ARG001 tool_ids: list[str] | None, ) -> Clarification | None: """Raise a clarification to check the user is happy with all tool calls before proceeding.""" @@ -217,8 +218,8 @@ def _clarify_on_tool_call_hook( return None -def log_step_outputs(plan: Plan, plan_run: PlanRun, step: Step, output: Output) -> None: # noqa: ARG001 +def log_step_outputs(plan: PlanV2, plan_run: PlanRunV2, step: StepV2, output: Output) -> None: # noqa: ARG001 """Log the output of a step in the plan.""" logger().info( - f"Step with task {step.task} using tool {step.tool_id} completed with result: {output}" + f"Step {step.step_name} completed with result: {output}" ) diff --git a/portia/open_source_tools/browser_tool.py b/portia/open_source_tools/browser_tool.py index b4f266d26..eb0e13c8d 100644 --- a/portia/open_source_tools/browser_tool.py +++ b/portia/open_source_tools/browser_tool.py @@ -615,10 +615,18 @@ def __init__(self, api_key: str | None = None, project_id: str | None = None) -> def step_complete(self, ctx: ToolRunContext) -> None: """Call when the step is complete closes the session to persist context.""" # Only clean up the session on the final browser tool call - if any( - step.tool_id == "browser_tool" - for step in ctx.plan.steps[ctx.plan_run.current_step_index + 1 :] - ): + # Check if there are more browser tool steps coming up + from portia.builder.invoke_tool_step import InvokeToolStep + + has_more_browser_steps = False + for step in ctx.plan.steps[ctx.plan_run.current_step_index + 1 :]: + if isinstance(step, InvokeToolStep): + tool_id = step.tool if isinstance(step.tool, str) else step.tool.id + if tool_id == "browser_tool": + has_more_browser_steps = True + break + + if has_more_browser_steps: return session_id = ctx.end_user.get_additional_data("bb_session_id") diff --git a/portia/plan_run.py b/portia/plan_run.py index 016619cd7..78bbeab8d 100644 --- a/portia/plan_run.py +++ b/portia/plan_run.py @@ -17,6 +17,8 @@ from __future__ import annotations +from typing import TYPE_CHECKING, Any + from pydantic import BaseModel, ConfigDict, Field from portia.clarification import ( @@ -25,9 +27,14 @@ ClarificationListType, ) from portia.common import PortiaEnum +from portia.config import Config +from portia.end_user import EndUser from portia.execution_agents.output import LocalDataValue, Output from portia.prefixed_uuid import PlanRunUUID, PlanUUID +if TYPE_CHECKING: + from portia.builder.plan_v2 import PlanV2 + class PlanRunState(PortiaEnum): """The current state of the Plan Run. @@ -228,3 +235,148 @@ def from_plan_run(cls, plan_run: PlanRun) -> ReadOnlyPlanRun: plan_run_inputs=plan_run.plan_run_inputs, structured_output_schema=plan_run.structured_output_schema, ) + + +class StepOutputValue(BaseModel): + """Value that can be referenced by name in a PlanV2 run. + + Attributes: + value: The referenced value. + description: Description of the referenced value. + step_name: The name of the step that produced this value. + step_num: The step number that produced this value. + + """ + + value: Any = Field(description="The referenced value.") + description: str = Field(description="Description of the referenced value.", default="") + step_name: str = Field(description="The name of the step that produced this value.") + step_num: int = Field(description="The step number that produced this value.") + + +class PlanRunV2(BaseModel): + """A V2 plan run represents a running instance of a PlanV2. + + This is the native representation for PlanV2 execution, replacing the legacy + PlanRun structure with a cleaner, more consistent data model. + + Attributes: + id: A unique ID for this plan run. + state: The current state of the plan run. + current_step_index: The current step that is being executed. + plan: The PlanV2 being executed. + end_user: The end user executing the plan. + step_output_values: List of outputs from completed steps. + final_output: The final consolidated output of the plan run. + plan_run_inputs: Dict mapping plan input names to their values. + config: The Portia configuration. + + """ + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + id: PlanRunUUID = Field( + default_factory=PlanRunUUID, + description="A unique ID for this plan run.", + ) + state: PlanRunState = Field( + default=PlanRunState.NOT_STARTED, + description="The current state of the plan run.", + ) + current_step_index: int = Field( + default=0, + description="The current step that is being executed.", + ) + plan: PlanV2 = Field( + description="The PlanV2 being executed.", + ) + end_user: EndUser = Field( + description="The end user executing the plan.", + ) + step_output_values: list[StepOutputValue] = Field( + default_factory=list, + description="List of outputs from completed steps.", + ) + final_output: Output | None = Field( + default=None, + description="The final consolidated output of the plan run.", + ) + plan_run_inputs: dict[str, LocalDataValue] = Field( + default_factory=dict, + description="Dict mapping plan input names to their values.", + ) + config: Config = Field( + description="The Portia configuration.", + ) + clarifications: ClarificationListType = Field( + default_factory=list, + description="Clarifications raised during this plan run.", + ) + + def get_outstanding_clarifications(self) -> ClarificationListType: + """Return all outstanding clarifications. + + Returns: + ClarificationListType: A list of outstanding clarifications that have not been resolved. + + """ + return [ + clarification + for clarification in self.clarifications + if not clarification.resolved + ] + + def get_clarifications_for_step(self, step: int | None = None) -> ClarificationListType: + """Return clarifications for the given step. + + Args: + step: The step to get clarifications for. Defaults to current step. + + Returns: + ClarificationListType: A list of clarifications for the given step. + + """ + if step is None: + step = self.current_step_index + return [ + clarification + for clarification in self.clarifications + if clarification.step == step + ] + + def get_clarification_for_step( + self, category: ClarificationCategory, step: int | None = None + ) -> Clarification | None: + """Return a clarification of the given category for the given step if it exists. + + Args: + step: The step to get a clarification for. Defaults to current step. + category: The category of the clarification to get. + + Returns: + Clarification | None: The clarification if found, None otherwise. + + """ + if step is None: + step = self.current_step_index + return next( + ( + clarification + for clarification in self.clarifications + if clarification.step == step and clarification.category == category + ), + None, + ) + + def __str__(self) -> str: + """Return the string representation of the PlanRunV2. + + Returns: + str: A string representation containing key run attributes. + + """ + return ( + f"PlanRunV2(id={self.id}, plan_id={self.plan.id}, " + f"state={self.state}, current_step_index={self.current_step_index}, " + f"final_output={'set' if self.final_output else 'unset'})" + ) diff --git a/portia/run_context.py b/portia/run_context.py index f8881ea17..d7d7c966f 100644 --- a/portia/run_context.py +++ b/portia/run_context.py @@ -1,43 +1,27 @@ """Context for a PlanV2 run.""" -from typing import Any - from pydantic import BaseModel, ConfigDict, Field -from portia.builder.plan_v2 import PlanV2 -from portia.config import Config -from portia.end_user import EndUser from portia.execution_hooks import ExecutionHooks -from portia.plan import Plan -from portia.plan_run import PlanRun +from portia.plan_run import PlanRunV2 from portia.storage import Storage from portia.telemetry.telemetry_service import BaseProductTelemetry from portia.tool import ToolRunContext from portia.tool_registry import ToolRegistry -class StepOutputValue(BaseModel): - """Value that can be referenced by name.""" - - value: Any = Field(description="The referenced value.") - description: str = Field(description="Description of the referenced value.", default="") - step_name: str = Field(description="The name of the referenced value.") - step_num: int = Field(description="The step number of the referenced value.") +class RunContext(BaseModel): + """Context for executing a PlanV2. + This is the main context object used throughout PlanV2 execution. + It contains the PlanRunV2 instance along with all services needed + for execution (storage, tools, hooks, etc.). -class RunContext(BaseModel): - """Data that is returned from a step.""" + """ model_config = ConfigDict(arbitrary_types_allowed=True) - plan: PlanV2 = Field(description="The Portia plan being executed.") - legacy_plan: Plan = Field(description="The legacy plan representation.") - plan_run: PlanRun = Field(description="The current plan run instance.") - end_user: EndUser = Field(description="The end user executing the plan.") - step_output_values: list[StepOutputValue] = Field( - default_factory=list, description="Outputs set by the step." - ) - config: Config = Field(description="The Portia config.") + plan_run: PlanRunV2 = Field(description="The current plan run instance.") storage: Storage = Field(description="The Portia storage.") tool_registry: ToolRegistry = Field(description="The Portia tool registry.") execution_hooks: ExecutionHooks = Field(description="The Portia execution hooks.") @@ -46,9 +30,6 @@ class RunContext(BaseModel): def get_tool_run_ctx(self) -> ToolRunContext: """Get the tool run context.""" return ToolRunContext( - end_user=self.end_user, plan_run=self.plan_run, - plan=self.legacy_plan, - config=self.config, clarifications=self.plan_run.get_clarifications_for_step(), ) diff --git a/portia/tool.py b/portia/tool.py index 0dd3cdbdc..6935fda38 100644 --- a/portia/tool.py +++ b/portia/tool.py @@ -39,6 +39,7 @@ model_validator, ) +from portia.builder.plan_v2 import PlanV2 from portia.clarification import ( ActionClarification, Clarification, @@ -58,8 +59,7 @@ from portia.execution_agents.output import LocalDataValue, Output from portia.logger import logger from portia.mcp_session import McpClientConfig, get_mcp_session -from portia.plan import Plan -from portia.plan_run import PlanRun +from portia.plan_run import PlanRunV2 from portia.templates.render import render_template """MAX_TOOL_DESCRIPTION_LENGTH is limited to stop overflows in the planner context window.""" @@ -69,22 +69,36 @@ class ToolRunContext(BaseModel): """Context passed to tools when running. + This context provides tools with access to the current plan run state, + including the plan being executed, user information, configuration, and + any clarifications that need to be handled. + Attributes: - plan_run(PlanRun): The run the tool run is part of. - plan(Plan): The plan the tool run is part of. - config(Config): The config for the SDK as a whole. - clarifications(ClarificationListType): Relevant clarifications for this tool plan_run. + plan_run: The V2 plan run the tool run is part of. + clarifications: Relevant clarifications for this tool execution. """ - model_config = ConfigDict(extra="forbid") + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) - end_user: EndUser - plan_run: PlanRun - plan: Plan - config: Config + plan_run: PlanRunV2 clarifications: ClarificationListType + @property + def end_user(self) -> EndUser: + """Get the end user from the plan run.""" + return self.plan_run.end_user + + @property + def config(self) -> Config: + """Get the config from the plan run.""" + return self.plan_run.config + + @property + def plan(self) -> PlanV2: + """Get the plan from the plan run.""" + return self.plan_run.plan + class _ArgsSchemaPlaceholder(BaseModel): """Placeholder ArgsSchema for tools that take no arguments."""