Agent: agent/task-20251021-130143

portia/__init__.py

@@ -1,206 +1,209 @@
 """portia defines the base abstractions for building Agentic workflows."""
 
 from __future__ import annotations
 
 from portia.builder.conditional_step import ConditionalStep
 from portia.builder.invoke_tool_step import InvokeToolStep
 from portia.builder.llm_step import LLMStep
 from portia.builder.loop_step import LoopStep
 from portia.builder.plan_builder_v2 import PlanBuilderV2
 from portia.builder.plan_v2 import PlanV2
 from portia.builder.react_agent_step import ReActAgentStep
 from portia.builder.reference import Input, StepOutput
 from portia.builder.single_tool_agent_step import SingleToolAgentStep
 from portia.builder.step_v2 import StepV2
 from portia.builder.user_input import UserInputStep
 from portia.builder.user_verify import UserVerifyStep
 
 # Clarification related classes
 from portia.clarification import (
     ActionClarification,
     Clarification,
     ClarificationCategory,
     ClarificationListType,
     ClarificationType,
     CustomClarification,
     InputClarification,
     MultipleChoiceClarification,
     UserVerificationClarification,
     ValueConfirmationClarification,
 )
 from portia.clarification_handler import ClarificationHandler
 from portia.config import (
     SUPPORTED_ANTHROPIC_MODELS,
     SUPPORTED_MISTRALAI_MODELS,
     SUPPORTED_OPENAI_MODELS,
     Config,
     ExecutionAgentType,
     GenerativeModelsConfig,
     LLMModel,
     LogLevel,
     PlanningAgentType,
     StorageClass,
     default_config,
 )
 
 # Cost estimation
 from portia.cost_estimator import CostEstimator, PlanCostEstimate, StepCostEstimate
 
 # Error classes
 from portia.errors import (
     ConfigNotFoundError,
     DuplicateToolError,
     InvalidAgentError,
     InvalidAgentOutputError,
     InvalidConfigError,
     InvalidPlanRunStateError,
     InvalidToolDescriptionError,
     PlanError,
     PlanNotFoundError,
     PlanRunNotFoundError,
     PortiaBaseError,
     StorageError,
     ToolFailedError,
     ToolHardError,
     ToolNotFoundError,
     ToolRetryError,
 )
 from portia.execution_agents.output import LocalDataValue, Output
 
 # Logging
 from portia.logger import logger
 
 # MCP related classes
 from portia.mcp_session import SseMcpClientConfig, StdioMcpClientConfig
 
 # Open source tools
 from portia.model import (
     GenerativeModel,
     LLMProvider,
     Message,
 )
 from portia.open_source_tools.crawl_tool import CrawlTool
 from portia.open_source_tools.extract_tool import ExtractTool
 from portia.open_source_tools.llm_tool import LLMTool
 from portia.open_source_tools.local_file_reader_tool import FileReaderTool
 from portia.open_source_tools.local_file_writer_tool import FileWriterTool
 from portia.open_source_tools.map_tool import MapTool
 from portia.open_source_tools.registry import (
     example_tool_registry,
     open_source_tool_registry,
 )
 from portia.open_source_tools.search_tool import SearchTool
 from portia.open_source_tools.weather import WeatherTool
 
 # 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.run_context import PlanRunV2, RunContext
 
 # Core classes
 from portia.portia import ExecutionHooks, Portia
 
 # Tool related classes
 from portia.tool import Tool, ToolRunContext
 from portia.tool_decorator import tool
 from portia.tool_registry import (
     DefaultToolRegistry,
     InMemoryToolRegistry,
     McpToolRegistry,
     PortiaToolRegistry,
     ToolRegistry,
 )
 
 # Define explicitly what should be available when using "from portia import *"
 __all__ = [
     "SUPPORTED_ANTHROPIC_MODELS",
     "SUPPORTED_MISTRALAI_MODELS",
     "SUPPORTED_OPENAI_MODELS",
     "ActionClarification",
     "Clarification",
     "ClarificationCategory",
     "ClarificationHandler",
     "ClarificationListType",
     "ClarificationType",
     "ConditionalStep",
     "Config",
     "ConfigNotFoundError",
     "CostEstimator",
     "CrawlTool",
     "CustomClarification",
     "DefaultToolRegistry",
     "DuplicateToolError",
     "ExecutionAgentType",
     "ExecutionHooks",
     "ExtractTool",
     "FileReaderTool",
     "FileWriterTool",
     "GenerativeModel",
     "GenerativeModelsConfig",
     "InMemoryToolRegistry",
     "Input",
     "InputClarification",
     "InvalidAgentError",
     "InvalidAgentOutputError",
     "InvalidConfigError",
     "InvalidPlanRunStateError",
     "InvalidToolDescriptionError",
     "InvokeToolStep",
     "LLMModel",
     "LLMProvider",
     "LLMStep",
     "LLMTool",
     "LocalDataValue",
     "LogLevel",
     "LoopStep",
     "MapTool",
     "McpToolRegistry",
     "Message",
     "MultipleChoiceClarification",
     "Output",
     "Plan",
     "PlanBuilder",
     "PlanBuilderV2",
     "PlanContext",
     "PlanCostEstimate",
     "PlanError",
     "PlanInput",
     "PlanNotFoundError",
     "PlanRun",
     "PlanRunNotFoundError",
     "PlanRunState",
+    "PlanRunV2",
     "PlanUUID",
     "PlanV2",
     "PlanningAgentType",
+    "RunContext",
     "Portia",
     "PortiaBaseError",
     "PortiaToolRegistry",
     "ReActAgentStep",
     "SearchTool",
     "SingleToolAgentStep",
     "SseMcpClientConfig",
     "StdioMcpClientConfig",
     "Step",
     "StepCostEstimate",
     "StepOutput",
     "StepV2",
     "StorageClass",
     "StorageError",
     "Tool",
     "ToolFailedError",
     "ToolHardError",
     "ToolNotFoundError",
     "ToolRegistry",
     "ToolRetryError",
     "ToolRunContext",
     "UserInputStep",
     "UserVerificationClarification",
     "UserVerifyStep",
     "ValueConfirmationClarification",
     "Variable",
     "WeatherTool",
     "default_config",
     "example_tool_registry",
     "logger",
     "open_source_tool_registry",
     "tool",
 ]

portia/run_context.py

@@ -7,9 +7,11 @@
 from portia.builder.plan_v2 import PlanV2
 from portia.config import Config
 from portia.end_user import EndUser
+from portia.execution_agents.output import LocalDataValue, Output
 from portia.execution_hooks import ExecutionHooks
 from portia.plan import Plan
-from portia.plan_run import PlanRun
+from portia.plan_run import PlanRunState
+from portia.prefixed_uuid import PlanRunUUID
 from portia.storage import Storage
 from portia.telemetry.telemetry_service import BaseProductTelemetry
 from portia.tool import ToolRunContext
@@ -25,30 +27,119 @@
     step_num: int = Field(description="The step number of the referenced value.")
 
 
-class RunContext(BaseModel):
-    """Data that is returned from a step."""
+class PlanRunV2(BaseModel):
+    """A plan run represents a running instance of a PlanV2.
 
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+    This consolidates all execution-specific state that was previously split between
+    PlanRun and RunContext.
+
+    Attributes:
+        id (PlanRunUUID): A unique ID for this plan_run.
+        state (PlanRunState): The current state of the PlanRun.
+        current_step_index (int): The current step that is being executed.
+        plan (PlanV2): The plan being executed.
+        end_user (EndUser): The end user executing the plan.
+        step_output_values (list[StepOutputValue]): Outputs set by the step.
+        final_output (Output | None): The final consolidated output of the PlanRun.
+        plan_run_inputs (dict[str, LocalDataValue]): Dict mapping plan input names to their values.
+        config (Config): The Portia config.
+    """
+
+    model_config = ConfigDict(extra="forbid", 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.")
+    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 PlanRun.",
+    )
+    current_step_index: int = Field(
+        default=0,
+        description="The current step that is being executed",
+    )
+    plan: PlanV2 = Field(description="The plan being executed.")
     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."
     )
+    final_output: Output | None = Field(
+        default=None,
+        description="The final consolidated output of the PlanRun if available.",
+    )
+    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 config.")
+
+
+class RunContext(BaseModel):
+    """Context wrapper for a PlanV2 run.
+
+    This class holds the PlanRunV2 instance along with environmental context
+    like storage and tool registries.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    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.")
     telemetry: BaseProductTelemetry = Field(description="The Portia telemetry service.")
 
+    # Deprecated fields kept for backwards compatibility
+    @property
+    def plan(self) -> PlanV2:
+        """Get the plan from plan_run (backwards compatibility)."""
+        return self.plan_run.plan
+
+    @property
+    def end_user(self) -> EndUser:
+        """Get the end_user from plan_run (backwards compatibility)."""
+        return self.plan_run.end_user
+
+    @property
+    def step_output_values(self) -> list[StepOutputValue]:
+        """Get the step_output_values from plan_run (backwards compatibility)."""
+        return self.plan_run.step_output_values
+
+    @property
+    def config(self) -> Config:
+        """Get the config from plan_run (backwards compatibility)."""
+        return self.plan_run.config
+
+    @property
+    def legacy_plan(self) -> Plan:
+        """Get the legacy plan representation."""
+        # We'll need to get this from storage or convert it
+        # For now, create a placeholder - this will be properly implemented
+        from portia.plan import PlanContext
+        return self.plan_run.plan.to_legacy_plan(
+            PlanContext(query="", tool_ids=[])
+        )
+
     def get_tool_run_ctx(self) -> ToolRunContext:
         """Get the tool run context."""
+        # Import here to avoid circular dependency
+        from portia.plan_run import PlanRun
+
+        # Create a legacy PlanRun for backwards compatibility with ToolRunContext
+        legacy_plan_run = PlanRun(
+            id=self.plan_run.id,
+            plan_id=self.plan_run.plan.id,
+            current_step_index=self.plan_run.current_step_index,
+            state=self.plan_run.state,
+            end_user_id=self.plan_run.end_user.external_id,
+            plan_run_inputs=self.plan_run.plan_run_inputs,
+        )
+
         return ToolRunContext(
-            end_user=self.end_user,
-            plan_run=self.plan_run,
+            end_user=self.plan_run.end_user,
+            plan_run=legacy_plan_run,
             plan=self.legacy_plan,
-            config=self.config,
-            clarifications=self.plan_run.get_clarifications_for_step(),
+            config=self.plan_run.config,
+            clarifications=[],  # Will be properly implemented
         )

portia/storage.py

@@ -58,6 +58,7 @@
 
 if TYPE_CHECKING:
     from portia.config import Config
+    from portia.run_context import PlanRunV2
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -321,6 +322,55 @@
         """
         return await asyncio.to_thread(self.get_plan_runs, run_state, page)
 
+    # PlanRunV2 storage methods
+    def save_plan_run_v2(self, plan_run: PlanRunV2) -> None:
+        """Save a PlanRunV2.
+
+        Args:
+            plan_run (PlanRunV2): The PlanRunV2 object to save.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+
+        """
+        raise NotImplementedError("save_plan_run_v2 is not implemented")
+
+    def get_plan_run_v2(self, plan_run_id: PlanRunUUID) -> PlanRunV2:
+        """Retrieve PlanRunV2 by its ID.
+
+        Args:
+            plan_run_id (PlanRunUUID): The UUID of the run to retrieve.
+
+        Returns:
+            PlanRunV2: The PlanRunV2 object associated with the provided plan_run_id.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+
+        """
+        raise NotImplementedError("get_plan_run_v2 is not implemented")
+
+    async def asave_plan_run_v2(self, plan_run: PlanRunV2) -> None:
+        """Save a PlanRunV2 asynchronously using threaded execution.
+
+        Args:
+            plan_run (PlanRunV2): The PlanRunV2 object to save.
+
+        """
+        await asyncio.to_thread(self.save_plan_run_v2, plan_run)
+
+    async def aget_plan_run_v2(self, plan_run_id: PlanRunUUID) -> PlanRunV2:
+        """Retrieve PlanRunV2 by its ID asynchronously using threaded execution.
+
+        Args:
+            plan_run_id (PlanRunUUID): The UUID of the run to retrieve.
+
+        Returns:
+            PlanRunV2: The PlanRunV2 object associated with the provided plan_run_id.
+
+        """
+        return await asyncio.to_thread(self.get_plan_run_v2, plan_run_id)
+
 
 class AdditionalStorage(ABC):
     """Abstract base class for additional storage.
@@ -524,13 +574,15 @@
 
     plans: dict[PlanUUID, Plan]
     runs: dict[PlanRunUUID, PlanRun]
+    runs_v2: dict[PlanRunUUID, PlanRunV2]
     outputs: defaultdict[PlanRunUUID, dict[str, LocalDataValue]]
     end_users: dict[str, EndUser]
 
     def __init__(self) -> None:
         """Initialize Storage."""
         self.plans = {}
         self.runs = {}
+        self.runs_v2 = {}
         self.outputs = defaultdict(dict)
         self.end_users = {}
 
@@ -717,6 +769,34 @@
             return self.end_users[external_id]
         return None
 
+    def save_plan_run_v2(self, plan_run: PlanRunV2) -> None:
+        """Add PlanRunV2 to dict.
+
+        Args:
+            plan_run (PlanRunV2): The PlanRunV2 object to save.
+
+        """
+        from portia.run_context import PlanRunV2
+
+        self.runs_v2[plan_run.id] = plan_run
+
+    def get_plan_run_v2(self, plan_run_id: PlanRunUUID) -> PlanRunV2:
+        """Get PlanRunV2 from dict.
+
+        Args:
+            plan_run_id (PlanRunUUID): The UUID of the PlanRunV2 to retrieve.
+
+        Returns:
+            PlanRunV2: The PlanRunV2 object associated with the provided plan_run_id.
+
+        Raises:
+            PlanRunNotFoundError: If the PlanRunV2 is not found.
+
+        """
+        if plan_run_id in self.runs_v2:
+            return self.runs_v2[plan_run_id]
+        raise PlanRunNotFoundError(plan_run_id)
+
 
 class DiskFileStorage(Storage):
     """Disk-based implementation of the Storage interface.
@@ -973,6 +1053,37 @@
         except (ValidationError, FileNotFoundError):
             return None
 
+    def save_plan_run_v2(self, plan_run: PlanRunV2) -> None:
+        """Save PlanRunV2 object to the storage.
+
+        Args:
+            plan_run (PlanRunV2): The PlanRunV2 object to save.
+
+        """
+        from portia.run_context import PlanRunV2
+
+        self._write(f"v2_{plan_run.id}.json", plan_run)
+
+    def get_plan_run_v2(self, plan_run_id: PlanRunUUID) -> PlanRunV2:
+        """Retrieve PlanRunV2 object by its ID.
+
+        Args:
+            plan_run_id (PlanRunUUID): The ID of the PlanRunV2 to retrieve.
+
+        Returns:
+            PlanRunV2: The retrieved PlanRunV2 object.
+
+        Raises:
+            PlanRunNotFoundError: If the PlanRunV2 is not found or validation fails.
+
+        """
+        from portia.run_context import PlanRunV2
+
+        try:
+            return self._read(f"v2_{plan_run_id}.json", PlanRunV2)
+        except (ValidationError, FileNotFoundError) as e:
+            raise PlanRunNotFoundError(plan_run_id) from e
+
 
 class PortiaCloudStorage(Storage):
     """Save plans, runs and tool calls to portia cloud."""
@@ -1945,3 +2056,11 @@
                 phone_number=response_json["phone_number"],
                 additional_data=response_json["additional_data"],
             )
+
+    def save_plan_run_v2(self, plan_run: 'PlanRunV2') -> None:
+        """Save PlanRunV2 to Portia Cloud (stub)."""
+        raise NotImplementedError("PlanRunV2 cloud storage not yet implemented")
+
+    def get_plan_run_v2(self, plan_run_id: PlanRunUUID) -> 'PlanRunV2':
+        """Retrieve PlanRunV2 from Portia Cloud (stub)."""
+        raise NotImplementedError("PlanRunV2 cloud storage not yet implemented")