Agent: agent/task-20251021-131626

portia/__init__.py

@@ -93,22 +93,25 @@
 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 import PlanBuilder, PlanContext, PlanInput, PlanUUID, Step, Variable
 from portia.plan_run import PlanRun, PlanRunState
 
+# Type alias for backward compatibility - Plan now points to PlanV2
+Plan = PlanV2
+
 # 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__ = [

portia/builder/plan_v2.py

@@ -3,16 +3,19 @@
 from __future__ import annotations
 
 import uuid
-from typing import Self
+from typing import TYPE_CHECKING, Self
 
 from pydantic import BaseModel, Field, model_validator
 
 from portia.builder.reference import default_step_name
 from portia.builder.step_v2 import StepV2
 from portia.logger import logger
-from portia.plan import Plan, PlanContext, PlanInput
+from portia.plan import PlanContext, PlanInput
 from portia.prefixed_uuid import PlanUUID
 
+if TYPE_CHECKING:
+    from portia.plan import Plan
+
 
 class PlanV2(BaseModel):
     """An ordered collection of executable steps that can be executed by Portia.

portia/cost_estimator.py

@@ -23,9 +23,8 @@
 from portia.end_user import EndUser
 from portia.logger import logger
 from portia.open_source_tools.llm_tool import LLMTool
-from portia.plan import Plan, PlanContext, Step
 from portia.plan_run import PlanRun, PlanRunState
 from portia.token_check import estimate_tokens
 from portia.tool import ToolRunContext
 
 if TYPE_CHECKING:
@@ -120,7 +119,7 @@
         self.config = config or Config.from_default()
         self.estimation_tool = LLMTool(structured_output_schema=LLMEstimationResult)
 
-    def plan_estimate(self, plan: Plan | PlanV2) -> PlanCostEstimate:
+    def plan_estimate(self, plan: PlanV2) -> PlanCostEstimate:
         """Estimate the cost of running a plan.
 
         Args:
@@ -132,11 +131,9 @@
         """
         logger().info("Starting cost estimation for plan")
 
-        if isinstance(plan, Plan):
-            return self._estimate_v1_plan(plan)
         return self._estimate_v2_plan(plan)
 
-    async def aplan_estimate(self, plan: Plan | PlanV2) -> PlanCostEstimate:
+    async def aplan_estimate(self, plan: PlanV2) -> PlanCostEstimate:
         """Asynchronously estimate the cost of running a plan.
 
         Args:
@@ -148,27 +145,6 @@
         """
         return await asyncio.to_thread(self.plan_estimate, plan)
 
-    def _estimate_v1_plan(self, plan: Plan) -> PlanCostEstimate:
-        """Estimate costs for a V1 plan."""
-        step_estimates = []
-        total_cost = 0.0
-
-        execution_model = self.config.get_execution_model()
-        model_name = execution_model.model_name if execution_model else "gpt-4o"
-
-        for step in plan.steps:
-            estimate = self._estimate_v1_step_cost(step, model_name)
-            step_estimates.append(estimate)
-            total_cost += estimate.estimated_cost
-
-        return PlanCostEstimate(
-            total_estimated_cost=total_cost,
-            step_estimates=step_estimates,
-            model_used=model_name,
-            methodology=self._get_methodology_explanation(),
-            limitations=self._get_limitations_explanation(),
-        )
-
     def _estimate_v2_plan(self, plan: PlanV2) -> PlanCostEstimate:
         """Estimate costs for a V2 plan."""
         step_estimates = []
@@ -190,66 +166,6 @@
             limitations=self._get_limitations_explanation(),
         )
 
-    def _estimate_v1_step_cost(self, step: Step, model_name: str) -> StepCostEstimate:
-        """Estimate cost for a V1 plan step."""
-        step_name = f"Step {step.output}"
-        step_type = "V1Step"
-        task = step.task
-
-        tools = step.tool_id or ""
-        input_context_tokens = estimate_tokens(task) + 500  # Base context estimation
-
-        estimation_result = self._get_llm_estimation(
-            "ExecutionAgentStep", task, model_name, tools, input_context_tokens
-        )
-
-        base_cost = self._calculate_cost(
-            estimation_result["estimated_input_tokens"] * estimation_result["number_of_llm_calls"],
-            estimation_result["estimated_output_tokens"] * estimation_result["number_of_llm_calls"],
-            model_name,
-        )
-
-        has_condition = step.condition is not None
-        introspection_cost = 0.0
-
-        if has_condition:
-            introspection_model = self.config.get_introspection_model()
-            if introspection_model is not None:
-                introspection_model_name = introspection_model.model_name
-                introspection_cost = self._calculate_introspection_cost(
-                    introspection_model_name
-                )  # pragma: no cover
-            else:
-                default_model = self.config.get_default_model()  # pragma: no cover
-                introspection_model_name = (
-                    default_model.model_name if default_model else "gpt-4o"
-                )  # pragma: no cover
-                introspection_cost = self._calculate_introspection_cost(
-                    introspection_model_name
-                )  # pragma: no cover
-
-        total_cost = base_cost + introspection_cost
-
-        return StepCostEstimate(
-            step_name=step_name,
-            step_type=step_type,
-            estimated_input_tokens=estimation_result["estimated_input_tokens"]
-            * estimation_result["number_of_llm_calls"],
-            estimated_output_tokens=estimation_result["estimated_output_tokens"]
-            * estimation_result["number_of_llm_calls"],
-            estimated_cost=total_cost,
-            cost_breakdown={"execution": base_cost, "introspection": introspection_cost},
-            explanation=(
-                f"LLM-driven estimation for V1 step: "
-                f"{estimation_result['estimated_input_tokens']} input x "
-                f"{estimation_result['number_of_llm_calls']} calls + "
-                f"{estimation_result['estimated_output_tokens']} output x "
-                f"{estimation_result['number_of_llm_calls']} calls"
-            ),
-            has_condition=has_condition,
-            introspection_cost=introspection_cost,
-        )
-
     def _estimate_v2_step_cost(self, step: StepV2, model_name: str) -> StepCostEstimate:
         """Estimate cost for a V2 plan step using LLM-driven estimation."""
         step_type = type(step).__name__
@@ -334,9 +250,8 @@
         self, step_type: str, task: str, model_name: str, tools: str, input_context_tokens: int
     ) -> dict[str, Any]:
         """Use LLM to estimate token usage for a step."""
-        plan = Plan(
-            plan_context=PlanContext(query="cost_estimation", tool_ids=[]),
-            plan_inputs=[],
+        plan = PlanV2(
+            label="cost_estimation",
             steps=[],
         )
 

portia/execution_agents/base_execution_agent.py

@@ -3,25 +3,26 @@
 The BaseAgent class is the base class that all agents must extend.
 """
 
 from __future__ import annotations
 
 import asyncio
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Literal
 
 from langchain_core.messages import ToolMessage
 from langgraph.graph import END, MessagesState
 
 from portia.execution_agents.context import StepInput, build_context
 from portia.execution_agents.execution_utils import (
     MAX_RETRIES,
     AgentNode,
     is_clarification,
     is_soft_tool_error,
 )
 from portia.execution_agents.output import LocalDataValue
+from portia.builder.plan_v2 import PlanV2
 from portia.logger import logger
-from portia.plan import Plan, ReadOnlyStep, Step
+from portia.plan import ReadOnlyStep, Step
 from portia.plan_run import PlanRun, ReadOnlyPlanRun
 from portia.telemetry.telemetry_service import ProductTelemetry
 
@@ -51,7 +52,7 @@
 
     def __init__(
         self,
-        plan: Plan,
+        plan: PlanV2,
         plan_run: PlanRun,
         config: Config,
         end_user: EndUser,
@@ -67,7 +68,7 @@
         of the execute_sync method.
 
         Args:
-            plan (Plan): The plan containing the steps.
+            plan (PlanV2): The plan containing the steps.
             plan_run (PlanRun): The run that contains the step and related data.
             config (Config): The configuration settings for the agent.
             end_user (EndUser): The end user for the execution.
@@ -85,11 +86,13 @@
         self.execution_hooks = execution_hooks
         self.telemetry = ProductTelemetry()
         self.new_clarifications: list[Clarification] = []
+        # Convert PlanV2 to legacy steps for execution
+        self._legacy_steps = [step.to_legacy_step(plan) for step in plan.steps]
 
     @property
     def step(self) -> Step:
         """Get the current step from the plan."""
-        return self.plan.steps[self.plan_run.current_step_index]
+        return self._legacy_steps[self.plan_run.current_step_index]
 
     @abstractmethod
     def execute_sync(self) -> Output:

portia/execution_agents/default_execution_agent.py

@@ -4,36 +4,37 @@
 in completing tasks.
 """
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Literal
 
 from langchain_core.prompts import (
     ChatPromptTemplate,
     HumanMessagePromptTemplate,
     SystemMessagePromptTemplate,
 )
 from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.prebuilt import ToolNode
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from portia.clarification import Clarification, InputClarification
 from portia.errors import InvalidAgentError, InvalidPlanRunStateError
 from portia.execution_agents.base_execution_agent import BaseExecutionAgent
 from portia.execution_agents.context import StepInput  # noqa: TC001
 from portia.execution_agents.execution_utils import (
     MAX_RETRIES,
     AgentNode,
     get_arg_value_with_templating,
     process_output,
     template_in_required_inputs,
     tool_call_or_end,
 )
 from portia.execution_agents.memory_extraction import MemoryExtractionStep
 from portia.execution_agents.utils.step_summarizer import StepSummarizer
+from portia.builder.plan_v2 import PlanV2
 from portia.logger import logger
 from portia.model import GenerativeModel, Message
-from portia.plan import Plan, ReadOnlyStep
+from portia.plan import ReadOnlyStep
 from portia.plan_run import PlanRun, ReadOnlyPlanRun
 from portia.telemetry.views import ExecutionAgentUsageTelemetryEvent, ToolCallTelemetryEvent
 from portia.tool import Tool, ToolRunContext

portia/execution_agents/one_shot_agent.py

@@ -31,9 +31,10 @@
 )
 from portia.execution_agents.memory_extraction import MemoryExtractionStep
 from portia.execution_agents.utils.step_summarizer import StepSummarizer
+from portia.builder.plan_v2 import PlanV2
 from portia.logger import logger
 from portia.model import GenerativeModel
-from portia.plan import Plan, ReadOnlyStep
+from portia.plan import ReadOnlyStep
 from portia.plan_run import PlanRun, ReadOnlyPlanRun
 from portia.telemetry.views import ExecutionAgentUsageTelemetryEvent, ToolCallTelemetryEvent
 from portia.tool import Tool, ToolRunContext

portia/execution_hooks.py

@@ -18,7 +18,8 @@
 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.builder.plan_v2 import PlanV2
+from portia.plan import Step
 from portia.plan_run import PlanRun
 from portia.tool import Tool
 
@@ -53,7 +54,7 @@ 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, PlanRun, Step], BeforeStepExecutionOutcome] | None = None
     """Called before executing each step.
 
     Args:
@@ -66,7 +67,7 @@ class ExecutionHooks(BaseModel):
             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, PlanRun, Step, Output], None] | None = None
     """Called after executing each step.
 
     When there's an error, this is called with the error as the output value.
@@ -78,15 +79,15 @@ class ExecutionHooks(BaseModel):
         output: The output from the step execution
     """
 
-    before_plan_run: Callable[[Plan, PlanRun], None] | None = None
+    before_plan_run: Callable[[PlanV2, PlanRun], 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
     """
 
-    after_plan_run: Callable[[Plan, PlanRun, Output], None] | None = None
+    after_plan_run: Callable[[PlanV2, PlanRun, 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
@@ -217,7 +218,7 @@ 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: PlanRun, step: Step, 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}"

portia/introspection_agents/default_introspection_agent.py

@@ -9,14 +9,14 @@
 from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
 from langchain.schema import SystemMessage
 
+from portia.builder.plan_v2 import PlanV2
 from portia.config import Config
 from portia.introspection_agents.introspection_agent import (
     BaseIntrospectionAgent,
     PreStepIntrospection,
 )
 from portia.logger import logger
 from portia.model import Message
-from portia.plan import Plan
 from portia.plan_run import PlanRun
 from portia.storage import AgentMemory, AgentMemoryValue
 
@@ -100,11 +100,13 @@ def __init__(self, config: Config, agent_memory: AgentMemory) -> None:
 
     def pre_step_introspection(
         self,
-        plan: Plan,
+        plan: PlanV2,
         plan_run: PlanRun,
     ) -> PreStepIntrospection:
         """Ask the LLM whether to continue, skip or fail the plan_run."""
-        introspection_condition = plan.steps[plan_run.current_step_index].condition
+        # Convert PlanV2 steps to legacy Step objects for introspection
+        legacy_steps = [step.to_legacy_step(plan) for step in plan.steps]
+        introspection_condition = legacy_steps[plan_run.current_step_index].condition
 
         memory_outputs = [
             self.agent_memory.get_plan_run_output(output.output_name, plan_run.id)
@@ -126,28 +128,30 @@ def pre_step_introspection(
                     prev_step_outputs=plan_run.outputs.model_dump_json(),
                     plan_run_inputs=plan_run.plan_run_inputs,
                     memory_outputs=memory_outputs,
-                    query=plan.plan_context.query,
-                    condition=plan.steps[plan_run.current_step_index].condition,
+                    query=plan.label,
+                    condition=legacy_steps[plan_run.current_step_index].condition,
                     current_step_idex=plan_run.current_step_index + 1,
                     total_steps_count=len(plan.steps),
-                    plan=self._get_plan_steps_pretty(plan),
+                    plan=self._get_plan_steps_pretty(legacy_steps),
                 )
             ],
         )
 
-    def _get_plan_steps_pretty(self, plan: Plan) -> str:
+    def _get_plan_steps_pretty(self, steps: list) -> str:
         """Get the pretty print representation of the plan steps."""
         return "\n".join(
-            [f"Step {i+1}: {step.pretty_print()}" for i, step in enumerate(plan.steps)]
+            [f"Step {i+1}: {step.pretty_print()}" for i, step in enumerate(steps)]
         )
 
     async def apre_step_introspection(
         self,
-        plan: Plan,
+        plan: PlanV2,
         plan_run: PlanRun,
     ) -> PreStepIntrospection:
         """pre_step_introspection is introspection run before a plan happens.."""
-        introspection_condition = plan.steps[plan_run.current_step_index].condition
+        # Convert PlanV2 steps to legacy Step objects for introspection
+        legacy_steps = [step.to_legacy_step(plan) for step in plan.steps]
+        introspection_condition = legacy_steps[plan_run.current_step_index].condition
 
         memory_outputs = [
             await self.agent_memory.aget_plan_run_output(output.output_name, plan_run.id)
@@ -169,11 +173,11 @@ async def apre_step_introspection(
                     prev_step_outputs=plan_run.outputs.model_dump_json(),
                     plan_run_inputs=plan_run.plan_run_inputs,
                     memory_outputs=memory_outputs,
-                    query=plan.plan_context.query,
-                    condition=plan.steps[plan_run.current_step_index].condition,
+                    query=plan.label,
+                    condition=legacy_steps[plan_run.current_step_index].condition,
                     current_step_idex=plan_run.current_step_index + 1,
                     total_steps_count=len(plan.steps),
-                    plan=self._get_plan_steps_pretty(plan),
+                    plan=self._get_plan_steps_pretty(legacy_steps),
                 )
             ],
         )