diff --git a/portia/__init__.py b/portia/__init__.py index 50dbb49f0..76f189aa5 100644 --- a/portia/__init__.py +++ b/portia/__init__.py @@ -93,9 +93,12 @@ 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 diff --git a/portia/builder/plan_v2.py b/portia/builder/plan_v2.py index 278cb916c..51854fd79 100644 --- a/portia/builder/plan_v2.py +++ b/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. diff --git a/portia/cost_estimator.py b/portia/cost_estimator.py index 461443ebe..5826b59a2 100644 --- a/portia/cost_estimator.py +++ b/portia/cost_estimator.py @@ -23,7 +23,6 @@ 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 @@ -120,7 +119,7 @@ def __init__(self, config: Config | None = None) -> None: 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 @@ def plan_estimate(self, plan: Plan | PlanV2) -> PlanCostEstimate: """ 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 @@ async def aplan_estimate(self, plan: Plan | PlanV2) -> PlanCostEstimate: """ 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 @@ def _estimate_v2_plan(self, plan: PlanV2) -> PlanCostEstimate: 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 @@ def _get_llm_estimation( 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=[], ) diff --git a/portia/execution_agents/base_execution_agent.py b/portia/execution_agents/base_execution_agent.py index 025bc5e08..2e3cde0ae 100644 --- a/portia/execution_agents/base_execution_agent.py +++ b/portia/execution_agents/base_execution_agent.py @@ -20,8 +20,9 @@ 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 @@ class BaseExecutionAgent: def __init__( self, - plan: Plan, + plan: PlanV2, plan_run: PlanRun, config: Config, end_user: EndUser, @@ -67,7 +68,7 @@ def __init__( 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 @@ def __init__( 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: diff --git a/portia/execution_agents/default_execution_agent.py b/portia/execution_agents/default_execution_agent.py index fd200847f..748ffedbc 100644 --- a/portia/execution_agents/default_execution_agent.py +++ b/portia/execution_agents/default_execution_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, 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 diff --git a/portia/execution_agents/one_shot_agent.py b/portia/execution_agents/one_shot_agent.py index 65cc9cab3..e1d78c279 100644 --- a/portia/execution_agents/one_shot_agent.py +++ b/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 diff --git a/portia/execution_hooks.py b/portia/execution_hooks.py index 5b8d2fa13..7a7c12ccd 100644 --- a/portia/execution_hooks.py +++ b/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,7 +79,7 @@ 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: @@ -86,7 +87,7 @@ class ExecutionHooks(BaseModel): 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}" diff --git a/portia/introspection_agents/default_introspection_agent.py b/portia/introspection_agents/default_introspection_agent.py index b9fa5f80d..450d201b3 100644 --- a/portia/introspection_agents/default_introspection_agent.py +++ b/portia/introspection_agents/default_introspection_agent.py @@ -9,6 +9,7 @@ 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, @@ -16,7 +17,6 @@ ) 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), ) ], ) diff --git a/portia/introspection_agents/introspection_agent.py b/portia/introspection_agents/introspection_agent.py index 195a5b25c..1571e403c 100644 --- a/portia/introspection_agents/introspection_agent.py +++ b/portia/introspection_agents/introspection_agent.py @@ -5,9 +5,9 @@ from pydantic import BaseModel, Field +from portia.builder.plan_v2 import PlanV2 from portia.common import PortiaEnum from portia.config import Config -from portia.plan import Plan from portia.plan_run import PlanRun from portia.storage import AgentMemory @@ -61,7 +61,7 @@ def __init__(self, config: Config, agent_memory: AgentMemory) -> None: @abstractmethod def pre_step_introspection( self, - plan: Plan, + plan: PlanV2, plan_run: PlanRun, ) -> PreStepIntrospection: """pre_step_introspection is introspection run before a plan happens..""" @@ -69,7 +69,7 @@ def pre_step_introspection( async def apre_step_introspection( self, - plan: Plan, + plan: PlanV2, plan_run: PlanRun, ) -> PreStepIntrospection: """pre_step_introspection is introspection run before a plan happens..""" diff --git a/portia/plan.py b/portia/plan.py index 9a708ff69..bb03582d6 100644 --- a/portia/plan.py +++ b/portia/plan.py @@ -408,154 +408,6 @@ def serialize_tool_ids(self, tool_ids: list[str]) -> list[str]: return sorted(tool_ids) -class Plan(BaseModel): - """A plan represents a series of steps that an agent should follow to execute the query. - - A plan defines the entire sequence of steps required to process a query and generate a result. - It also includes the context in which the plan was created. - - Args: - id (PlanUUID): A unique ID for the plan. - plan_context (PlanContext): The context for when the plan was created. - steps (list[Step]): The set of steps that make up the plan. - inputs (list[PlanInput]): The inputs required by the plan. - - """ - - model_config = ConfigDict(extra="forbid") - id: PlanUUID = Field( - default_factory=PlanUUID, - description="The ID of the plan.", - ) - plan_context: PlanContext = Field(description="The context for when the plan was created.") - steps: list[Step] = Field(description="The set of steps to solve the query.") - plan_inputs: list[PlanInput] = Field( - default=[], - description="The inputs required by the plan.", - ) - structured_output_schema: type[BaseModel] | None = Field( - default=None, - exclude=True, - description="The optional structured output schema for the query.", - ) - - def __str__(self) -> str: - """Return the string representation of the plan. - - Returns: - str: A string representation of the plan's ID, context, and steps. - - """ - return ( - f"PlanModel(id={self.id!r}," - f"plan_context={self.plan_context!r}, " - f"steps={self.steps!r}, " - f"inputs={self.plan_inputs!r}" - ) - - @classmethod - def from_response(cls, response_json: dict) -> Plan: - """Create a plan from a response. - - Args: - response_json (dict): The response from the API. - - Returns: - Plan: The plan. - - """ - return cls( - id=PlanUUID.from_string(response_json["id"]), - plan_context=PlanContext( - query=response_json["query"], - tool_ids=response_json["tool_ids"], - ), - steps=[Step.model_validate(step) for step in response_json["steps"]], - plan_inputs=[ - PlanInput.model_validate(input_) for input_ in response_json.get("plan_inputs", []) - ], - ) - - def pretty_print(self) -> str: - """Return the pretty print representation of the plan. - - Returns: - str: A pretty print representation of the plan's ID, context, and steps. - - """ - portia_tools = [tool for tool in self.plan_context.tool_ids if tool.startswith("portia:")] - other_tools = [ - tool for tool in self.plan_context.tool_ids if not tool.startswith("portia:") - ] - tools_summary = f"{len(portia_tools)} portia tools, {len(other_tools)} other tools" - - inputs_section = "" - if self.plan_inputs: - inputs_section = ( - "Inputs:\n " - + "\n ".join([input_.pretty_print() for input_ in self.plan_inputs]) - + "\n" - ) - - return ( - f"Task: {self.plan_context.query}\n" - f"Tools Available Summary: {tools_summary}\n" - f"{inputs_section}" - f"Steps:\n" - + "\n".join([step.pretty_print() for step in self.steps]) - + ( - f"\nStructured Output Schema: {self.structured_output_schema.__name__}" - if self.structured_output_schema - else "" - ) - ) - - @model_validator(mode="after") - def validate_plan(self) -> Self: - """Validate the plan. - - Checks that all outputs + conditions are unique. - - Returns: - Plan: The validated plan. - - """ - outputs = [step.output + (step.condition or "") for step in self.steps] - if len(outputs) != len(set(outputs)): - raise ValueError("Outputs + conditions must be unique") - - # Validate plan input names are unique - input_names = [input_.name for input_ in self.plan_inputs] - if len(input_names) != len(set(input_names)): - raise ValueError("Plan input names must be unique") - - return self - - -class ReadOnlyPlan(Plan): - """A read-only copy of a plan, passed to agents for reference. - - This class provides a non-modifiable view of a plan instance, - ensuring that agents can access plan details without altering them. - """ - - model_config = ConfigDict(frozen=True, extra="forbid") - - @classmethod - def from_plan(cls, plan: Plan) -> ReadOnlyPlan: - """Create a read-only plan from a normal plan. - - Args: - plan (Plan): The original plan instance to create a read-only copy from. - - Returns: - ReadOnlyPlan: A new read-only instance of the provided plan. - - """ - return cls( - id=plan.id, - plan_context=plan.plan_context, - steps=plan.steps, - plan_inputs=plan.plan_inputs, - structured_output_schema=plan.structured_output_schema, - ) +# Legacy Plan and ReadOnlyPlan classes have been removed. +# Plan is now an alias for PlanV2 (defined in portia/__init__.py). +# Import PlanV2 from portia.builder.plan_v2 for the new plan format. diff --git a/portia/planning_agents/base_planning_agent.py b/portia/planning_agents/base_planning_agent.py index 14b00f0f2..02ddcfcd0 100644 --- a/portia/planning_agents/base_planning_agent.py +++ b/portia/planning_agents/base_planning_agent.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, ConfigDict, Field +from portia.builder.plan_v2 import PlanV2 from portia.plan import Plan, PlanInput, Step if TYPE_CHECKING: @@ -51,7 +52,7 @@ def generate_steps_or_error( query: str, tool_list: list[Tool], end_user: EndUser, - examples: list[Plan] | None = None, + examples: list[PlanV2 | Plan] | None = None, plan_inputs: list[PlanInput] | None = None, ) -> StepsOrError: """Generate a list of steps for the given query. @@ -63,7 +64,7 @@ def generate_steps_or_error( query (str): The user query to generate a list of steps for. tool_list (list[Tool]): A list of tools available for the plan. end_user (EndUser): The end user for this plan - examples (list[Plan] | None): Optional list of example plans to guide the PlanningAgent. + examples (list[PlanV2 | Plan] | None): Optional list of example plans to guide the PlanningAgent. plan_inputs (list[PlanInput] | None): Optional list of PlanInput objects defining the inputs required for the plan. @@ -78,7 +79,7 @@ async def agenerate_steps_or_error( query: str, tool_list: list[Tool], end_user: EndUser, - examples: list[Plan] | None = None, + examples: list[PlanV2 | Plan] | None = None, plan_inputs: list[PlanInput] | None = None, ) -> StepsOrError: """Generate a list of steps for the given query asynchronously. @@ -90,7 +91,7 @@ async def agenerate_steps_or_error( query (str): The user query to generate a list of steps for. tool_list (list[Tool]): A list of tools available for the plan. end_user (EndUser): The end user for this plan - examples (list[Plan] | None): Optional list of example plans to guide the PlanningAgent. + examples (list[PlanV2 | Plan] | None): Optional list of example plans to guide the PlanningAgent. plan_inputs (list[PlanInput] | None): Optional list of PlanInput objects defining the inputs required for the plan. diff --git a/portia/planning_agents/context.py b/portia/planning_agents/context.py index c84877118..352cdbe67 100644 --- a/portia/planning_agents/context.py +++ b/portia/planning_agents/context.py @@ -9,6 +9,7 @@ from portia.templates.render import render_template if TYPE_CHECKING: + from portia.builder.plan_v2 import PlanV2 from portia.end_user import EndUser from portia.plan import Plan, PlanInput from portia.tool import Tool @@ -18,7 +19,7 @@ def render_prompt_insert_defaults( query: str, tool_list: list[Tool], end_user: EndUser, - examples: list[Plan] | None = None, + examples: list[PlanV2 | Plan] | None = None, plan_inputs: list[PlanInput] | None = None, previous_errors: list[str] | None = None, ) -> str: diff --git a/portia/planning_agents/default_planning_agent.py b/portia/planning_agents/default_planning_agent.py index 6080da252..119eb0687 100644 --- a/portia/planning_agents/default_planning_agent.py +++ b/portia/planning_agents/default_planning_agent.py @@ -14,6 +14,7 @@ from portia.planning_agents.context import render_prompt_insert_defaults if TYPE_CHECKING: + from portia.builder.plan_v2 import PlanV2 from portia.config import Config from portia.end_user import EndUser from portia.plan import Plan, PlanInput, Step @@ -80,7 +81,7 @@ def generate_steps_or_error( query: str, tool_list: list[Tool], end_user: EndUser, - examples: list[Plan] | None = None, + examples: list[PlanV2 | Plan] | None = None, plan_inputs: list[PlanInput] | None = None, ) -> StepsOrError: """Generate a plan or error using an LLM from a query and a list of tools.""" @@ -171,7 +172,7 @@ async def agenerate_steps_or_error( query: str, tool_list: list[Tool], end_user: EndUser, - examples: list[Plan] | None = None, + examples: list[PlanV2 | Plan] | None = None, plan_inputs: list[PlanInput] | None = None, ) -> StepsOrError: """Generate a plan or error using an LLM from a query and a list of tools.""" diff --git a/portia/portia.py b/portia/portia.py index a25c2d84c..db13af071 100644 --- a/portia/portia.py +++ b/portia/portia.py @@ -73,7 +73,10 @@ ) from portia.logger import logger, logger_manager, truncate_message from portia.open_source_tools.llm_tool import LLMTool -from portia.plan import Plan, PlanContext, PlanInput, PlanUUID, ReadOnlyPlan, ReadOnlyStep, Step +from portia.plan import PlanContext, PlanInput, PlanUUID, ReadOnlyStep, Step + +if TYPE_CHECKING: + from portia.plan import Plan, ReadOnlyPlan from portia.plan_run import PlanRun, PlanRunState, PlanRunUUID, ReadOnlyPlanRun from portia.planning_agents.default_planning_agent import DefaultPlanningAgent from portia.run_context import RunContext, StepOutputValue @@ -201,7 +204,7 @@ def run( self, query: str, tools: list[Tool] | list[str] | None = None, - example_plans: Sequence[Plan | PlanUUID | str] | None = None, + example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None = None, end_user: str | EndUser | None = None, plan_run_inputs: list[PlanInput] | list[dict[str, str]] | dict[str, str] | None = None, structured_output_schema: type[BaseModel] | None = None, @@ -215,8 +218,8 @@ def run( query (str): The query to be executed. tools (list[Tool] | list[str] | None): List of tools to use for the query. If not provided all tools in the registry will be used. - example_plans (Sequence[Plan | PlanUUID | str] | None): Optional list of example - plans or plan IDs. This can include Plan objects, PlanUUID objects, + example_plans (Sequence[Plan | PlanV2 | PlanUUID | str] | None): Optional list of example + plans or plan IDs. This can include Plan objects, PlanV2 objects, PlanUUID objects, or plan ID strings (starting with "plan-"). Plan IDs will be loaded from storage. If not provided, a default set of example plans will be used. end_user (str | EndUser | None = None): The end user for this plan run. @@ -267,7 +270,7 @@ async def arun( self, query: str, tools: list[Tool] | list[str] | None = None, - example_plans: Sequence[Plan | PlanUUID | str] | None = None, + example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None = None, end_user: str | EndUser | None = None, plan_run_inputs: list[PlanInput] | list[dict[str, str]] | dict[str, str] | None = None, structured_output_schema: type[BaseModel] | None = None, @@ -281,8 +284,8 @@ async def arun( query (str): The query to be executed. tools (list[Tool] | list[str] | None): List of tools to use for the query. If not provided all tools in the registry will be used. - example_plans (Sequence[Plan | PlanUUID | str] | None): Optional list of example - plans or plan IDs. This can include Plan objects, PlanUUID objects, + example_plans (Sequence[Plan | PlanV2 | PlanUUID | str] | None): Optional list of example + plans or plan IDs. This can include Plan objects, PlanV2 objects, PlanUUID objects, or plan ID strings (starting with "plan-"). Plan IDs will be loaded from storage. If not provided, a default set of example plans will be used. end_user (str | EndUser | None = None): The end user for this plan run. @@ -366,21 +369,21 @@ def plan( self, query: str, tools: list[Tool] | list[str] | None = None, - example_plans: Sequence[Plan | PlanUUID | str] | None = None, + example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None = None, end_user: str | EndUser | None = None, plan_inputs: list[PlanInput] | list[dict[str, str]] | list[str] | None = None, structured_output_schema: type[BaseModel] | None = None, use_cached_plan: bool = False, - ) -> Plan: + ) -> Plan | PlanV2: """Plans how to do the query given the set of tools and any examples. Args: query (str): The query to generate the plan for. tools (list[Tool] | list[str] | None): List of tools to use for the query. If not provided all tools in the registry will be used. - example_plans (Sequence[Plan | PlanUUID | str] | None): Optional list of example + example_plans (Sequence[Plan | PlanV2 | PlanUUID | str] | None): Optional list of example plans or plan IDs. - This can include Plan objects, PlanUUID objects, or plan ID strings + This can include Plan objects, PlanV2 objects, PlanUUID objects, or plan ID strings (starting with "plan-"). Plan IDs will be loaded from storage. If not provided, a default set of example plans will be used. end_user (str | EndUser | None = None): The optional end user for this plan. @@ -397,7 +400,7 @@ def plan( use_cached_plan (bool): Whether to use a cached plan if it exists. Returns: - Plan: The plan for executing the query. + Plan | PlanV2: The plan for executing the query. Raises: PlanError: If there is an error while generating the plan. @@ -429,19 +432,19 @@ def plan( ) def _resolve_example_plans( - self, example_plans: Sequence[Plan | PlanUUID | str] | None - ) -> list[Plan] | None: + self, example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None + ) -> list[Plan | PlanV2] | None: """Resolve example plans from Plan objects, PlanUUIDs and planID strings. Args: - example_plans (Sequence[Plan | PlanUUID | str] | None): List of example plans or + example_plans (Sequence[Plan | PlanV2 | PlanUUID | str] | None): List of example plans or plan IDs. - - Plan objects are used directly + - Plan and PlanV2 objects are used directly - PlanUUID objects are loaded from storage - String objects must be plan ID strings (starting with "plan-") Returns: - list[Plan] | None: List of resolved Plan objects, or None if input was None. + list[Plan | PlanV2] | None: List of resolved Plan or PlanV2 objects, or None if input was None. Raises: PlanNotFoundError: If a plan ID cannot be found in storage. @@ -460,19 +463,19 @@ def _resolve_example_plans( return resolved_plans async def _aresolve_example_plans( - self, example_plans: Sequence[Plan | PlanUUID | str] | None - ) -> list[Plan] | None: + self, example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None + ) -> list[Plan | PlanV2] | None: """Resolve example plans from Plan objects, PlanUUIDs and planID strings. Args: - example_plans (Sequence[Plan | PlanUUID | str] | None): List of example plans or + example_plans (Sequence[Plan | PlanV2 | PlanUUID | str] | None): List of example plans or plan IDs. - - Plan objects are used directly + - Plan and PlanV2 objects are used directly - PlanUUID objects are loaded from storage - String objects must be plan ID strings (starting with "plan-") Returns: - list[Plan] | None: List of resolved Plan objects, or None if input was None. + list[Plan | PlanV2] | None: List of resolved Plan or PlanV2 objects, or None if input was None. Raises: PlanNotFoundError: If a plan ID cannot be found in storage. @@ -490,44 +493,44 @@ async def _aresolve_example_plans( return resolved_plans - def _resolve_single_example_plan(self, example_plan: Plan | PlanUUID | str) -> Plan: + def _resolve_single_example_plan(self, example_plan: Plan | PlanV2 | PlanUUID | str) -> Plan | PlanV2: """Resolve a single example plan from various input types.""" - if isinstance(example_plan, Plan): + if isinstance(example_plan, (Plan, PlanV2)): return example_plan if isinstance(example_plan, PlanUUID): return self._load_plan_by_uuid(example_plan) if isinstance(example_plan, str): return self._resolve_string_example_plan(example_plan) raise TypeError( - f"Invalid example plan type: {type(example_plan)}. Expected Plan, PlanUUID, or str." + f"Invalid example plan type: {type(example_plan)}. Expected Plan, PlanV2, PlanUUID, or str." ) - async def _aresolve_single_example_plan(self, example_plan: Plan | PlanUUID | str) -> Plan: - if isinstance(example_plan, Plan): + async def _aresolve_single_example_plan(self, example_plan: Plan | PlanV2 | PlanUUID | str) -> Plan | PlanV2: + if isinstance(example_plan, (Plan, PlanV2)): return example_plan if isinstance(example_plan, PlanUUID): return await self._aload_plan_by_uuid(example_plan) if isinstance(example_plan, str): return await self._aresolve_string_example_plan(example_plan) raise TypeError( - f"Invalid example plan type: {type(example_plan)}. Expected Plan, PlanUUID, or str." + f"Invalid example plan type: {type(example_plan)}. Expected Plan, PlanV2, PlanUUID, or str." ) - def _load_plan_by_uuid(self, plan_uuid: PlanUUID) -> Plan: + def _load_plan_by_uuid(self, plan_uuid: PlanUUID) -> Plan | PlanV2: """Load a plan from storage by UUID.""" try: return self.storage.get_plan(plan_uuid) except Exception as e: raise PlanNotFoundError(plan_uuid) from e - async def _aload_plan_by_uuid(self, plan_uuid: PlanUUID) -> Plan: + async def _aload_plan_by_uuid(self, plan_uuid: PlanUUID) -> Plan | PlanV2: """Load a plan from storage by UUID asynchronously.""" try: return await self.storage.aget_plan(plan_uuid) except Exception as e: raise PlanNotFoundError(plan_uuid) from e - def _resolve_string_example_plan(self, example_plan: str) -> Plan: + def _resolve_string_example_plan(self, example_plan: str) -> Plan | PlanV2: """Resolve a string example plan - must be a plan ID string.""" # Only support plan ID strings, not query strings if not example_plan.startswith("plan-"): @@ -542,7 +545,7 @@ def _resolve_string_example_plan(self, example_plan: str) -> Plan: except Exception as e: raise PlanNotFoundError(plan_uuid) from e - async def _aresolve_string_example_plan(self, example_plan: str) -> Plan: + async def _aresolve_string_example_plan(self, example_plan: str) -> Plan | PlanV2: """Resolve a string example plan - must be a plan ID string.""" # Only support plan ID strings, not query strings if not example_plan.startswith("plan-"): @@ -561,19 +564,19 @@ async def aplan( self, query: str, tools: list[Tool] | list[str] | None = None, - example_plans: Sequence[Plan | PlanUUID | str] | None = None, + example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None = None, end_user: str | EndUser | None = None, plan_inputs: list[PlanInput] | list[dict[str, str]] | list[str] | None = None, structured_output_schema: type[BaseModel] | None = None, use_cached_plan: bool = False, - ) -> Plan: + ) -> Plan | PlanV2: """Plans how to do the query given the set of tools and any examples asynchronously. Args: query (str): The query to generate the plan for. tools (list[Tool] | list[str] | None): List of tools to use for the query. If not provided all tools in the registry will be used. - example_plans (list[Plan] | None): Optional list of example plans. If not + example_plans (list[Plan | PlanV2] | None): Optional list of example plans. If not provide a default set of example plans will be used. end_user (str | EndUser | None = None): The optional end user for this plan. plan_inputs (list[PlanInput] | list[dict[str, str]] | list[str] | None): Optional list @@ -589,7 +592,7 @@ async def aplan( use_cached_plan (bool): Whether to use a cached plan if it exists. Returns: - Plan: The plan for executing the query. + Plan | PlanV2: The plan for executing the query. Raises: PlanError: If there is an error while generating the plan. @@ -624,12 +627,12 @@ def _plan( self, query: str, tools: list[Tool] | list[str] | None = None, - example_plans: Sequence[Plan | PlanUUID | str] | None = None, + example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None = None, end_user: str | EndUser | None = None, plan_inputs: list[PlanInput] | list[dict[str, str]] | list[str] | None = None, structured_output_schema: type[BaseModel] | None = None, use_cached_plan: bool = False, - ) -> Plan: + ) -> Plan | PlanV2: """Implement synchronous planning logic. This is used when we're already in an event loop and can't use asyncio.run(). @@ -657,7 +660,7 @@ def _plan( use_cached_plan (bool): Whether to use a cached plan if it exists. Returns: - Plan: The plan for executing the query. + Plan | PlanV2: The plan for executing the query. Raises: PlanError: If there is an error while generating the plan. @@ -726,12 +729,12 @@ async def _aplan( self, query: str, tools: list[Tool] | list[str] | None = None, - example_plans: Sequence[Plan | PlanUUID | str] | None = None, + example_plans: Sequence[Plan | PlanV2 | PlanUUID | str] | None = None, end_user: str | EndUser | None = None, plan_inputs: list[PlanInput] | list[dict[str, str]] | list[str] | None = None, structured_output_schema: type[BaseModel] | None = None, use_cached_plan: bool = False, - ) -> Plan: + ) -> Plan | PlanV2: """Async implementation of planning logic. This is the core async implementation that both sync and async methods use. @@ -759,7 +762,7 @@ async def _aplan( use_cached_plan (bool): Whether to use a cached plan if it exists. Returns: - Plan: The plan for executing the query. + Plan | PlanV2: The plan for executing the query. Raises: PlanError: If there is an error while generating the plan. diff --git a/portia/run_context.py b/portia/run_context.py index f8881ea17..9f838b77d 100644 --- a/portia/run_context.py +++ b/portia/run_context.py @@ -8,7 +8,6 @@ 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.storage import Storage from portia.telemetry.telemetry_service import BaseProductTelemetry @@ -31,7 +30,6 @@ class RunContext(BaseModel): 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( @@ -48,7 +46,7 @@ def get_tool_run_ctx(self) -> ToolRunContext: return ToolRunContext( end_user=self.end_user, plan_run=self.plan_run, - plan=self.legacy_plan, + plan=self.plan, config=self.config, clarifications=self.plan_run.get_clarifications_for_step(), ) diff --git a/portia/storage.py b/portia/storage.py index cfdd55658..6fd8fdab8 100644 --- a/portia/storage.py +++ b/portia/storage.py @@ -1,20 +1,23 @@ -"""Storage classes for managing the saving and retrieval of plans, runs, and tool calls. +"""Storage classes for managing the saving and retrieval of runs, and tool calls. + +DEPRECATED: Plan storage has been removed. The plan is the code. This module defines a set of storage classes that provide different backends for saving, retrieving, -and managing plans, runs, and tool calls. These storage classes include both in-memory and -file-based storage, as well as integration with the Portia Cloud API. Each class is responsible -for handling interactions with its respective storage medium, including validating responses -and raising appropriate exceptions when necessary. +and managing runs and tool calls. Plan storage methods are deprecated and raise NotImplementedError. +These storage classes include both in-memory and file-based storage, as well as integration with +the Portia Cloud API. Each class is responsible for handling interactions with its respective +storage medium, including validating responses and raising appropriate exceptions when necessary. Classes: - Storage (Base Class): A base class that defines common interfaces for all storage types, - ensuring consistent methods for saving and retrieving plans, runs, and tool calls. - - InMemoryStorage: An in-memory implementation of the `Storage` class for storing plans, - runs, and tool calls in a temporary, volatile storage medium. - - FileStorage: A file-based implementation of the `Storage` class for storing plans, runs, + ensuring consistent methods for saving and retrieving runs and tool calls. + Plan-related methods are deprecated. + - InMemoryStorage: An in-memory implementation of the `Storage` class for storing runs + and tool calls in a temporary, volatile storage medium. + - FileStorage: A file-based implementation of the `Storage` class for storing runs and tool calls as local files in the filesystem. - PortiaCloudStorage: A cloud-based implementation of the `Storage` class that interacts with - the Portia Cloud API to save and retrieve plans, runs, and tool call records. + the Portia Cloud API to save and retrieve runs and tool call records. Each storage class handles the following tasks: - Sending and receiving data to its respective storage medium - memory, file system, or API. @@ -65,15 +68,18 @@ class PlanStorage(ABC): """Abstract base class for storing and retrieving plans. + DEPRECATED: Plan storage is deprecated. The plan is the code. + These methods are maintained for backward compatibility but should not be used. + Subclasses must implement the methods to save and retrieve plans. Methods: save_plan(self, plan: Plan) -> None: - Save a plan. + Save a plan. DEPRECATED - raises NotImplementedError. get_plan(self, plan_id: PlanUUID) -> Plan: - Get a plan by ID. + Get a plan by ID. DEPRECATED - raises NotImplementedError. plan_exists(self, plan_id: PlanUUID) -> bool: - Check if a plan exists without raising an error. + Check if a plan exists without raising an error. DEPRECATED - raises NotImplementedError. """ @@ -81,19 +87,25 @@ class PlanStorage(ABC): def save_plan(self, plan: Plan) -> None: """Save a plan. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: plan (Plan): The Plan object to save. Raises: - NotImplementedError: If the method is not implemented. + NotImplementedError: Always raised - plan storage is deprecated. """ - raise NotImplementedError("save_plan is not implemented") + raise NotImplementedError("save_plan is deprecated - the plan is the code") @abstractmethod def get_plan(self, plan_id: PlanUUID) -> Plan: """Retrieve a plan by its ID. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: plan_id (PlanUUID): The UUID of the plan to retrieve. @@ -101,25 +113,34 @@ def get_plan(self, plan_id: PlanUUID) -> Plan: Plan: The Plan object associated with the provided plan_id. Raises: - NotImplementedError: If the method is not implemented. + NotImplementedError: Always raised - plan storage is deprecated. """ - raise NotImplementedError("get_plan is not implemented") + raise NotImplementedError("get_plan is deprecated - the plan is the code") @abstractmethod def get_plan_by_query(self, query: str) -> Plan: """Get a plan by query. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: query (str): The query to get a plan for. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - raise NotImplementedError("get_plan_by_query is not implemented") + raise NotImplementedError("get_plan_by_query is deprecated - the plan is the code") @abstractmethod def plan_exists(self, plan_id: PlanUUID) -> bool: """Check if a plan exists without raising an error. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: plan_id (PlanUUID): The UUID of the plan to check. @@ -127,14 +148,17 @@ def plan_exists(self, plan_id: PlanUUID) -> bool: bool: True if the plan exists, False otherwise. Raises: - NotImplementedError: If the method is not implemented. + NotImplementedError: Always raised - plan storage is deprecated. """ - raise NotImplementedError("plan_exists is not implemented") + raise NotImplementedError("plan_exists is deprecated - the plan is the code") def get_similar_plans(self, query: str, threshold: float = 0.5, limit: int = 10) -> list[Plan]: """Get similar plans to the query. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: query (str): The query to get similar plans for. threshold (float): The threshold for similarity. @@ -144,58 +168,85 @@ def get_similar_plans(self, query: str, threshold: float = 0.5, limit: int = 10) list[Plan]: The list of similar plans. Raises: - NotImplementedError: If the method is not implemented. + NotImplementedError: Always raised - plan storage is deprecated. """ - raise NotImplementedError("get_similar_plans is not implemented") # pragma: no cover + raise NotImplementedError("get_similar_plans is deprecated - the plan is the code") # pragma: no cover async def asave_plan(self, plan: Plan) -> None: """Save a plan asynchronously using threaded execution. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: plan (Plan): The Plan object to save. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - await asyncio.to_thread(self.save_plan, plan) + raise NotImplementedError("asave_plan is deprecated - the plan is the code") async def aget_plan(self, plan_id: PlanUUID) -> Plan: """Retrieve a plan by its ID asynchronously using threaded execution. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: plan_id (PlanUUID): The UUID of the plan to retrieve. Returns: Plan: The Plan object associated with the provided plan_id. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - return await asyncio.to_thread(self.get_plan, plan_id) + raise NotImplementedError("aget_plan is deprecated - the plan is the code") async def aget_plan_by_query(self, query: str) -> Plan: """Get a plan by query asynchronously using threaded execution. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: query (str): The query to get a plan for. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - return await asyncio.to_thread(self.get_plan_by_query, query) + raise NotImplementedError("aget_plan_by_query is deprecated - the plan is the code") async def aplan_exists(self, plan_id: PlanUUID) -> bool: """Check if a plan exists without raising an error asynchronously using threaded execution. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: plan_id (PlanUUID): The UUID of the plan to check. Returns: bool: True if the plan exists, False otherwise. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - return await asyncio.to_thread(self.plan_exists, plan_id) + raise NotImplementedError("aplan_exists is deprecated - the plan is the code") async def aget_similar_plans( self, query: str, threshold: float = 0.5, limit: int = 10 ) -> list[Plan]: """Get similar plans to the query asynchronously using threaded execution. + DEPRECATED: Plan storage is deprecated. The plan is the code. + This method should no longer be used. + Args: query (str): The query to get similar plans for. threshold (float): The threshold for similarity. @@ -204,8 +255,11 @@ async def aget_similar_plans( Returns: list[Plan]: The list of similar plans. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - return await asyncio.to_thread(self.get_similar_plans, query, threshold, limit) + raise NotImplementedError("aget_similar_plans is deprecated - the plan is the code") class PlanRunListResponse(BaseModel): @@ -517,19 +571,20 @@ def log_tool_call(tool_call: ToolCallRecord) -> None: class InMemoryStorage(Storage): - """Simple storage class that keeps plans + runs in memory. + """Simple storage class that keeps runs in memory. + + DEPRECATED: Plan storage has been removed. The plan is the code. + Plan-related methods will raise NotImplementedError. Tool Calls are logged via the LogAdditionalStorage. """ - plans: dict[PlanUUID, Plan] runs: dict[PlanRunUUID, PlanRun] outputs: defaultdict[PlanRunUUID, dict[str, LocalDataValue]] end_users: dict[str, EndUser] def __init__(self) -> None: """Initialize Storage.""" - self.plans = {} self.runs = {} self.outputs = defaultdict(dict) self.end_users = {} @@ -537,15 +592,22 @@ def __init__(self) -> None: def save_plan(self, plan: Plan) -> None: """Add plan to dict. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan (Plan): The Plan object to save. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - self.plans[plan.id] = plan + raise NotImplementedError("save_plan is deprecated - the plan is the code") def get_plan(self, plan_id: PlanUUID) -> Plan: """Get plan from dict. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The UUID of the plan to retrieve. @@ -553,37 +615,41 @@ def get_plan(self, plan_id: PlanUUID) -> Plan: Plan: The Plan object associated with the provided plan_id. Raises: - PlanNotFoundError: If the plan is not found. + NotImplementedError: Always raised - plan storage is deprecated. """ - if plan_id in self.plans: - return self.plans[plan_id] - raise PlanNotFoundError(plan_id) + raise NotImplementedError("get_plan is deprecated - the plan is the code") def get_plan_by_query(self, query: str) -> Plan: """Get a plan by query. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: query (str): The query to get a plan for. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - plan: Plan | None = None - for plan in self.plans.values(): - if plan.plan_context.query == query: - return plan - raise StorageError(f"No plan found for query: {query}") + raise NotImplementedError("get_plan_by_query is deprecated - the plan is the code") def plan_exists(self, plan_id: PlanUUID) -> bool: """Check if a plan exists in memory. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The UUID of the plan to check. Returns: bool: True if the plan exists, False otherwise. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - return plan_id in self.plans + raise NotImplementedError("plan_exists is deprecated - the plan is the code") def save_plan_run(self, plan_run: PlanRun) -> None: """Add run to dict. @@ -721,7 +787,10 @@ def get_end_user(self, external_id: str) -> EndUser | None: class DiskFileStorage(Storage): """Disk-based implementation of the Storage interface. - Stores serialized Plan and Run objects as JSON files on disk. + DEPRECATED: Plan storage has been removed. The plan is the code. + Plan-related methods will raise NotImplementedError. + + Stores serialized Run objects as JSON files on disk. """ def __init__(self, storage_dir: str | None) -> None: @@ -748,11 +817,11 @@ def _ensure_storage(self, file_path: str | None = None) -> None: Path(self.storage_dir, file_path).parent.mkdir(parents=True, exist_ok=True) def _write(self, file_path: str, content: BaseModel) -> None: - """Write a serialized Plan or Run to a JSON file. + """Write a serialized Run to a JSON file. Args: file_path (str): Path of the file to write. - content (BaseModel): The Plan or Run object to serialize. + content (BaseModel): The Run object to serialize. """ self._ensure_storage(file_path) # Ensure storage directory exists @@ -781,15 +850,22 @@ def _read(self, file_name: str, model: type[T]) -> T: def save_plan(self, plan: Plan) -> None: """Save a Plan object to the storage. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan (Plan): The Plan object to save. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - self._write(f"{plan.id}.json", plan) + raise NotImplementedError("save_plan is deprecated - the plan is the code") def get_plan(self, plan_id: PlanUUID) -> Plan: """Retrieve a Plan object by its ID. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The ID of the Plan to retrieve. @@ -797,48 +873,41 @@ def get_plan(self, plan_id: PlanUUID) -> Plan: Plan: The retrieved Plan object. Raises: - PlanNotFoundError: If the Plan is not found or validation fails. + NotImplementedError: Always raised - plan storage is deprecated. """ - try: - return self._read(f"{plan_id}.json", Plan) - except (ValidationError, FileNotFoundError) as e: - raise PlanNotFoundError(plan_id) from e + raise NotImplementedError("get_plan is deprecated - the plan is the code") def get_plan_by_query(self, query: str) -> Plan: """Get a plan by query. - This method will return the first plan that matches the query. This is not always the most - recent plan. + DEPRECATED: Plan storage is deprecated. The plan is the code. Args: query (str): The query to get a plan for. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - # Get all plan files and sort by modification time (newest first) - # Using st_mtime for cross-platform compatibility - plan_files = [ - f - for f in Path(self.storage_dir).iterdir() - if f.is_file() and f.name.startswith(PLAN_UUID_PREFIX) - ] - for f in plan_files: - plan = self._read(f.name, Plan) - if plan.plan_context.query == query: - return plan - raise StorageError(f"No plan found for query: {query}") + raise NotImplementedError("get_plan_by_query is deprecated - the plan is the code") def plan_exists(self, plan_id: PlanUUID) -> bool: """Check if a plan exists on disk. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The UUID of the plan to check. Returns: bool: True if the plan exists, False otherwise. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - return Path(self.storage_dir, f"{plan_id}.json").exists() + raise NotImplementedError("plan_exists is deprecated - the plan is the code") def save_plan_run(self, plan_run: PlanRun) -> None: """Save PlanRun object to the storage. @@ -975,7 +1044,11 @@ def get_end_user(self, external_id: str) -> EndUser | None: class PortiaCloudStorage(Storage): - """Save plans, runs and tool calls to portia cloud.""" + """Save runs and tool calls to portia cloud. + + DEPRECATED: Plan storage has been removed. The plan is the code. + Plan-related methods will raise NotImplementedError. + """ DEFAULT_MAX_CACHE_SIZE = 20 @@ -1082,65 +1155,36 @@ def check_response(self, response: httpx.Response) -> None: def save_plan(self, plan: Plan) -> None: """Save a plan to Portia Cloud. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan (Plan): The Plan object to save to the cloud. Raises: - StorageError: If the request to Portia Cloud fails. + NotImplementedError: Always raised - plan storage is deprecated. """ - try: - response = self.client.post( - url="/api/v0/plans/", - json={ - "id": str(plan.id), - "query": plan.plan_context.query, - "tool_ids": plan.plan_context.tool_ids, - "steps": [step.model_dump(mode="json") for step in plan.steps], - "plan_inputs": [ - {**input_.model_dump(mode="json"), "description": input_.description} - for input_ in plan.plan_inputs - ], - }, - ) - except Exception as e: - raise StorageError(e) from e - else: - self.check_response(response) + raise NotImplementedError("save_plan is deprecated - the plan is the code") async def asave_plan(self, plan: Plan) -> None: """Save a plan to Portia Cloud. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan (Plan): The Plan object to save to the cloud. Raises: - StorageError: If the request to Portia Cloud fails. + NotImplementedError: Always raised - plan storage is deprecated. """ - try: - async with self.client_builder.async_client() as client: - response = await client.post( - url="/api/v0/plans/", - json={ - "id": str(plan.id), - "query": plan.plan_context.query, - "tool_ids": plan.plan_context.tool_ids, - "steps": [step.model_dump(mode="json") for step in plan.steps], - "plan_inputs": [ - {**input_.model_dump(mode="json"), "description": input_.description} - for input_ in plan.plan_inputs - ], - }, - ) - except Exception as e: - raise StorageError(e) from e - else: - self.check_response(response) + raise NotImplementedError("asave_plan is deprecated - the plan is the code") def get_plan(self, plan_id: PlanUUID) -> Plan: """Retrieve a plan from Portia Cloud. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The ID of the plan to retrieve. @@ -1148,23 +1192,16 @@ def get_plan(self, plan_id: PlanUUID) -> Plan: Plan: The Plan object retrieved from Portia Cloud. Raises: - StorageError: If the request to Portia Cloud fails or the plan does not exist. + NotImplementedError: Always raised - plan storage is deprecated. """ - try: - response = self.client.get( - url=f"/api/v0/plans/{plan_id}/", - ) - except Exception as e: - raise StorageError(e) from e - else: - self.check_response(response) - response_json = response.json() - return Plan.from_response(response_json) + raise NotImplementedError("get_plan is deprecated - the plan is the code") async def aget_plan(self, plan_id: PlanUUID) -> Plan: """Retrieve a plan from Portia Cloud. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The ID of the plan to retrieve. @@ -1172,89 +1209,72 @@ async def aget_plan(self, plan_id: PlanUUID) -> Plan: Plan: The Plan object retrieved from Portia Cloud. Raises: - StorageError: If the request to Portia Cloud fails or the plan does not exist. + NotImplementedError: Always raised - plan storage is deprecated. """ - try: - async with self.client_builder.async_client() as client: - response = await client.get( - url=f"/api/v0/plans/{plan_id}/", - ) - except Exception as e: - raise StorageError(e) from e - else: - self.check_response(response) - response_json = response.json() - return Plan.from_response(response_json) + raise NotImplementedError("aget_plan is deprecated - the plan is the code") def get_plan_by_query(self, query: str) -> Plan: """Get a plan by query. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: query (str): The query to get a plan for. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - try: - plans = self.get_similar_plans(query, threshold=1.0, limit=1) - except Exception as e: - raise StorageError(e) from e - if not plans: - raise StorageError(f"No plan found for query: {query}") - return plans[0] + raise NotImplementedError("get_plan_by_query is deprecated - the plan is the code") async def aget_plan_by_query(self, query: str) -> Plan: """Get a plan by query asynchronously using threaded execution. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: query (str): The query to get a plan for. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - try: - plans = await self.aget_similar_plans(query, threshold=1.0, limit=1) - except Exception as e: - raise StorageError(e) from e - if not plans: - raise StorageError(f"No plan found for query: {query}") - return plans[0] + raise NotImplementedError("aget_plan_by_query is deprecated - the plan is the code") def plan_exists(self, plan_id: PlanUUID) -> bool: """Check if a plan exists in Portia Cloud. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The UUID of the plan to check. Returns: bool: True if the plan exists, False otherwise. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - try: - response = self.client.get( - url=f"/api/v0/plans/{plan_id}/", - ) - except Exception: # noqa: BLE001 - return False - else: - return response.is_success + raise NotImplementedError("plan_exists is deprecated - the plan is the code") async def aplan_exists(self, plan_id: PlanUUID) -> bool: """Check if a plan exists in Portia Cloud. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: plan_id (PlanUUID): The UUID of the plan to check. Returns: bool: True if the plan exists, False otherwise. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - try: - async with self.client_builder.async_client() as client: - response = await client.get( - url=f"/api/v0/plans/{plan_id}/", - ) - except Exception: # noqa: BLE001 - return False - else: - return response.is_success + raise NotImplementedError("aplan_exists is deprecated - the plan is the code") def save_plan_run(self, plan_run: PlanRun) -> None: """Save PlanRun to Portia Cloud. @@ -1774,6 +1794,8 @@ async def aget_plan_run_output( def get_similar_plans(self, query: str, threshold: float = 0.5, limit: int = 5) -> list[Plan]: """Get similar plans to the query. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: query (str): The query to get similar plans for. threshold (float): The threshold for similarity. @@ -1782,27 +1804,19 @@ def get_similar_plans(self, query: str, threshold: float = 0.5, limit: int = 5) Returns: list[Plan]: The list of similar plans. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - try: - response = self.client.post( - "/api/v0/plans/embeddings/search/", - json={ - "query": query, - "threshold": threshold, - "limit": limit, - }, - ) - self.check_response(response) - results = response.json() - return [Plan.from_response(result) for result in results] - except Exception as e: - raise StorageError(e) from e + raise NotImplementedError("get_similar_plans is deprecated - the plan is the code") async def aget_similar_plans( self, query: str, threshold: float = 0.5, limit: int = 5 ) -> list[Plan]: """Get similar plans to the query. + DEPRECATED: Plan storage is deprecated. The plan is the code. + Args: query (str): The query to get similar plans for. threshold (float): The threshold for similarity. @@ -1811,22 +1825,11 @@ async def aget_similar_plans( Returns: list[Plan]: The list of similar plans. + Raises: + NotImplementedError: Always raised - plan storage is deprecated. + """ - try: - async with self.client_builder.async_client() as client: - response = await client.post( - url="/api/v0/plans/embeddings/search/", - json={ - "query": query, - "threshold": threshold, - "limit": limit, - }, - ) - self.check_response(response) - results = response.json() - return [Plan.from_response(result) for result in results] - except Exception as e: - raise StorageError(e) from e + raise NotImplementedError("aget_similar_plans is deprecated - the plan is the code") def save_end_user(self, end_user: EndUser) -> EndUser: """Save an end_user to Portia Cloud. diff --git a/portia/templates/default_planning_agent.xml.jinja b/portia/templates/default_planning_agent.xml.jinja index ae809e883..4b790ecee 100644 --- a/portia/templates/default_planning_agent.xml.jinja +++ b/portia/templates/default_planning_agent.xml.jinja @@ -37,15 +37,15 @@ Use the following information to construct a plan in response to the following u - {{example.plan_context.tool_ids | safe }} + {% if example.plan_context is defined %}{{example.plan_context.tool_ids | safe }}{% else %}{{ example.steps | map(attribute='tool_id') | select | list | unique | sort }}{% endif %} - {{example.plan_context.query}} + {% if example.plan_context is defined %}{{example.plan_context.query}}{% else %}{{example.label}}{% endif %} [{% for step in example.steps %} - {{step.model_dump(exclude_none=True) | tojson}}, + {% if step.to_legacy_step is defined %}{{step.to_legacy_step(example).model_dump(exclude_none=True) | tojson}},{% else %}{{step.model_dump(exclude_none=True) | tojson}},{% endif %} {% endfor %}] {% endfor %} diff --git a/portia/tool.py b/portia/tool.py index 0dd3cdbdc..0f718e51b 100644 --- a/portia/tool.py +++ b/portia/tool.py @@ -58,7 +58,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.builder.plan_v2 import PlanV2 from portia.plan_run import PlanRun from portia.templates.render import render_template @@ -71,7 +71,7 @@ class ToolRunContext(BaseModel): Attributes: plan_run(PlanRun): The run the tool run is part of. - plan(Plan): The plan the tool run is part of. + plan(PlanV2): 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. @@ -81,7 +81,7 @@ class ToolRunContext(BaseModel): end_user: EndUser plan_run: PlanRun - plan: Plan + plan: PlanV2 config: Config clarifications: ClarificationListType diff --git a/tests/unit/planning_agents/test_default_planning_agent.py b/tests/unit/planning_agents/test_default_planning_agent.py index c73d8b478..5b10d5479 100644 --- a/tests/unit/planning_agents/test_default_planning_agent.py +++ b/tests/unit/planning_agents/test_default_planning_agent.py @@ -8,6 +8,7 @@ import pytest +from portia.builder.plan_v2 import PlanV2 from portia.end_user import EndUser from portia.open_source_tools.llm_tool import LLMTool from portia.plan import Plan, PlanContext, PlanInput, Step, Variable @@ -64,7 +65,7 @@ def generate_steps_or_error( query: str, tool_list: list[Tool], end_user: EndUser, - examples: list[Plan] | None = None, + examples: list[PlanV2 | Plan] | None = None, plan_inputs: list[PlanInput] | None = None, # noqa: ARG002 ) -> StepsOrError: return super().generate_steps_or_error(query, tool_list, end_user, examples) # type: ignore # noqa: PGH003