Skip to main content

portia.storage

Storage classes for managing the saving and retrieval of plans, runs, and tool calls.

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.

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, 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.

Each storage class handles the following tasks:

  • Sending and receiving data to its respective storage medium - memory, file system, or API.
  • Validating responses from storage and raising errors when necessary.
  • Handling exceptions and re-raising them as custom StorageError exceptions to provide more informative error handling.

PlanStorage Objects​

class PlanStorage(ABC)

Abstract base class for storing and retrieving plans.

Subclasses must implement the methods to save and retrieve plans.

Methods:

save_plan(self, plan: Plan) -> None: Save a plan. get_plan(self, plan_id: PlanUUID) -> Plan: Get a plan by ID.

save_plan​

@abstractmethod
def save_plan(plan: Plan) -> None

Save a plan.

Arguments:

  • plan Plan - The Plan object to save.

Raises:

  • NotImplementedError - If the method is not implemented.

get_plan​

@abstractmethod
def get_plan(plan_id: PlanUUID) -> Plan

Retrieve a plan by its ID.

Arguments:

  • plan_id PlanUUID - The UUID of the plan to retrieve.

Returns:

  • Plan - The Plan object associated with the provided plan_id.

Raises:

  • NotImplementedError - If the method is not implemented.

PlanRunListResponse Objects​

class PlanRunListResponse(BaseModel)

Response for the get_plan_runs operation. Can support pagination.

RunStorage Objects​

class RunStorage(ABC)

Abstract base class for storing and retrieving runs.

Subclasses must implement the methods to save and retrieve PlanRuns.

Methods:

save_plan_run(self, run: Run) -> None: Save a PlanRun. get_plan_run(self, plan_run_id: PlanRunUUID) -> PlanRun: Get PlanRun by ID. get_plan_runs(self, run_state: RunState | None = None, page=int | None = None) -> PlanRunListResponse: Return runs that match the given run_state

save_plan_run​

@abstractmethod
def save_plan_run(plan_run: PlanRun) -> None

Save a PlanRun.

Arguments:

  • plan_run PlanRun - The Run object to save.

Raises:

  • NotImplementedError - If the method is not implemented.

get_plan_run​

@abstractmethod
def get_plan_run(plan_run_id: PlanRunUUID) -> PlanRun

Retrieve PlanRun by its ID.

Arguments:

  • plan_run_id RunUUID - The UUID of the run to retrieve.

Returns:

  • Run - The Run object associated with the provided plan_run_id.

Raises:

  • NotImplementedError - If the method is not implemented.

get_plan_runs​

@abstractmethod
def get_plan_runs(run_state: PlanRunState | None = None,
                  page: int | None = None) -> PlanRunListResponse

List runs by their state.

Arguments:

  • run_state RunState | None - Optionally filter runs by their state.
  • page int | None - Optional pagination data

Returns:

  • list[Run] - A list of Run objects that match the given state.

Raises:

  • NotImplementedError - If the method is not implemented.

AdditionalStorage Objects​

class AdditionalStorage(ABC)

Abstract base class for additional storage.

Subclasses must implement the methods.

Methods:

save_tool_call(self, tool_call: ToolCallRecord) -> None: Save a tool_call.

save_tool_call​

@abstractmethod
def save_tool_call(tool_call: ToolCallRecord) -> None

Save a ToolCall.

Arguments:

  • tool_call ToolCallRecord - The ToolCallRecord object to save.

Raises:

  • NotImplementedError - If the method is not implemented.

LogAdditionalStorage Objects​

class LogAdditionalStorage(AdditionalStorage)

AdditionalStorage that logs calls rather than persisting them.

Useful for storages that don't care about tool_calls etc.

save_tool_call​

def save_tool_call(tool_call: ToolCallRecord) -> None

Log the tool call.

Arguments:

  • tool_call ToolCallRecord - The ToolCallRecord object to log.

Storage Objects​

class Storage(PlanStorage, RunStorage, LogAdditionalStorage)

Combined base class for Plan Run + Additional storages.

InMemoryStorage Objects​

class InMemoryStorage(PlanStorage, RunStorage, LogAdditionalStorage)

Simple storage class that keeps plans + runs in memory.

Tool Calls are logged via the LogAdditionalStorage.

__init__​

def __init__() -> None

Initialize Storage.

save_plan​

def save_plan(plan: Plan) -> None

Add plan to dict.

Arguments:

  • plan Plan - The Plan object to save.

get_plan​

def get_plan(plan_id: PlanUUID) -> Plan

Get plan from dict.

Arguments:

  • plan_id PlanUUID - The UUID of the plan to retrieve.

Returns:

  • Plan - The Plan object associated with the provided plan_id.

Raises:

  • PlanNotFoundError - If the plan is not found.

save_plan_run​

def save_plan_run(plan_run: PlanRun) -> None

Add run to dict.

Arguments:

  • plan_run PlanRun - The Run object to save.

get_plan_run​

def get_plan_run(plan_run_id: PlanRunUUID) -> PlanRun

Get run from dict.

Arguments:

  • plan_run_id PlanRunUUID - The UUID of the PlanRun to retrieve.

Returns:

  • PlanRun - The PlanRun object associated with the provided plan_run_id.

Raises:

  • PlanRunNotFoundError - If the PlanRun is not found.

get_plan_runs​

def get_plan_runs(run_state: PlanRunState | None = None,
                  page: int | None = None) -> PlanRunListResponse

Get run from dict.

Arguments:

  • run_state RunState | None - Optionally filter runs by their state.
  • page int | None - Optional pagination data which is not used for in memory storage.

Returns:

  • list[Run] - A list of Run objects that match the given state.

DiskFileStorage Objects​

class DiskFileStorage(PlanStorage, RunStorage, LogAdditionalStorage)

Disk-based implementation of the Storage interface.

Stores serialized Plan and Run objects as JSON files on disk.

__init__​

def __init__(storage_dir: str | None) -> None

Set storage dir.

Arguments:

  • storage_dir str | None - Optional directory for storing files.

save_plan​

def save_plan(plan: Plan) -> None

Save a Plan object to the storage.

Arguments:

  • plan Plan - The Plan object to save.

get_plan​

def get_plan(plan_id: PlanUUID) -> Plan

Retrieve a Plan object by its ID.

Arguments:

  • plan_id PlanUUID - The ID of the Plan to retrieve.

Returns:

  • Plan - The retrieved Plan object.

Raises:

  • PlanNotFoundError - If the Plan is not found or validation fails.

save_plan_run​

def save_plan_run(plan_run: PlanRun) -> None

Save PlanRun object to the storage.

Arguments:

  • plan_run PlanRun - The Run object to save.

get_plan_run​

def get_plan_run(plan_run_id: PlanRunUUID) -> PlanRun

Retrieve PlanRun object by its ID.

Arguments:

  • plan_run_id RunUUID - The ID of the Run to retrieve.

Returns:

  • Run - The retrieved Run object.

Raises:

  • RunNotFoundError - If the Run is not found or validation fails.

get_plan_runs​

def get_plan_runs(run_state: PlanRunState | None = None,
                  page: int | None = None) -> PlanRunListResponse

Find all plan runs in storage that match state.

Arguments:

  • run_state RunState | None - Optionally filter runs by their state.
  • page int | None - Optional pagination data which is not used for in memory storage.

Returns:

  • list[Run] - A list of Run objects that match the given state.

PortiaCloudStorage Objects​

class PortiaCloudStorage(Storage)

Save plans, runs and tool calls to portia cloud.

__init__​

def __init__(config: Config) -> None

Initialize the PortiaCloudStorage instance.

Arguments:

  • config Config - The configuration containing API details for Portia Cloud.

check_response​

def check_response(response: httpx.Response) -> None

Validate the response from Portia API.

Arguments:

  • response httpx.Response - The response from the Portia API to check.

Raises:

  • StorageError - If the response from the Portia API indicates an error.

save_plan​

def save_plan(plan: Plan) -> None

Save a plan to Portia Cloud.

Arguments:

  • plan Plan - The Plan object to save to the cloud.

Raises:

  • StorageError - If the request to Portia Cloud fails.

get_plan​

def get_plan(plan_id: PlanUUID) -> Plan

Retrieve a plan from Portia Cloud.

Arguments:

  • plan_id PlanUUID - The ID of the plan to retrieve.

Returns:

  • Plan - The Plan object retrieved from Portia Cloud.

Raises:

  • StorageError - If the request to Portia Cloud fails or the plan does not exist.

save_plan_run​

def save_plan_run(plan_run: PlanRun) -> None

Save PlanRun to Portia Cloud.

Arguments:

  • plan_run PlanRun - The Run object to save to the cloud.

Raises:

  • StorageError - If the request to Portia Cloud fails.

get_plan_run​

def get_plan_run(plan_run_id: PlanRunUUID) -> PlanRun

Retrieve PlanRun from Portia Cloud.

Arguments:

  • plan_run_id RunUUID - The ID of the run to retrieve.

Returns:

  • Run - The Run object retrieved from Portia Cloud.

Raises:

  • StorageError - If the request to Portia Cloud fails or the run does not exist.

get_plan_runs​

def get_plan_runs(run_state: PlanRunState | None = None,
                  page: int | None = None) -> PlanRunListResponse

Find all runs in storage that match state.

Arguments:

  • run_state RunState | None - Optionally filter runs by their state.
  • page int | None - Optional pagination data which is not used for in memory storage.

Returns:

  • list[Run] - A list of Run objects retrieved from Portia Cloud.

Raises:

  • StorageError - If the request to Portia Cloud fails.

save_tool_call​

def save_tool_call(tool_call: ToolCallRecord) -> None

Save a tool call to Portia Cloud.

Arguments:

  • tool_call ToolCallRecord - The ToolCallRecord object to save to the cloud.

Raises:

  • StorageError - If the request to Portia Cloud fails.