planar 0.9.3__py3-none-any.whl → 0.11.0__py3-none-any.whl

planar/ai/agent.py

@@ -50,7 +50,8 @@ class AgentWorkflowNotifier(AgentEventEmitter):
 class Agent[
     TInput: BaseModel | str,
     TOutput: BaseModel | str,
-](AgentBase[TInput, TOutput]):
+    TDeps,
+](AgentBase[TInput, TOutput, TDeps]):
     model: models.KnownModelName | models.Model = "openai:gpt-4o"
 
     async def run_step(

planar/ai/agent_base.py

@@ -15,6 +15,7 @@ from pydantic import BaseModel
 from pydantic_ai.settings import ModelSettings
 
 from planar.ai.models import AgentConfig, AgentEventEmitter, AgentRunResult
+from planar.ai.state import delete_state, set_state
 from planar.logging import get_logger
 from planar.modeling.field_helpers import JsonSchema
 from planar.utils import P, R, T, U
@@ -29,6 +30,7 @@ class AgentBase[
     # TODO: add `= str` default when we upgrade to 3.13
     TInput: BaseModel | str,
     TOutput: BaseModel | str,
+    TState,
 ](abc.ABC):
     """An LLM-powered agent that can be called directly within workflows."""
 
@@ -45,6 +47,7 @@ class AgentBase[
     )
     event_emitter: AgentEventEmitter | None = None
     durable: bool = True
+    state_type: Type[TState] | None = None
 
     # TODO: move here to serialize to frontend
     #
@@ -91,14 +94,16 @@ class AgentBase[
 
     @overload
     async def __call__(
-        self: "AgentBase[TInput, str]",
+        self: "AgentBase[TInput, str, TState]",
         input_value: TInput,
+        state: TState | None = None,
     ) -> AgentRunResult[str]: ...
 
     @overload
     async def __call__(
-        self: "AgentBase[TInput, TOutput]",
+        self: "AgentBase[TInput, TOutput, TState]",
         input_value: TInput,
+        state: TState | None = None,
     ) -> AgentRunResult[TOutput]: ...
 
     def as_step_if_durable(
@@ -120,6 +125,7 @@ class AgentBase[
     async def __call__(
         self,
         input_value: TInput,
+        state: TState | None = None,
     ) -> AgentRunResult[Any]:
         if self.input_type is not None and not isinstance(input_value, self.input_type):
             raise ValueError(
@@ -147,9 +153,22 @@ class AgentBase[
                 return_type=AgentRunResult[self.output_type],
             )
 
-        result = await run_step(input_value=input_value)
-        # Cast the result to ensure type compatibility
-        return cast(AgentRunResult[TOutput], result)
+        if state is not None:
+            if self.state_type is None:
+                raise ValueError("state cannot be provided when state_type is not set")
+            if not isinstance(state, self.state_type):
+                raise ValueError(
+                    f"state must be of type {self.state_type}, but got {type(state)}"
+                )
+            set_state(cast(TState, state))
+
+        try:
+            result = await run_step(input_value=input_value)
+            # Cast the result to ensure type compatibility
+            return cast(AgentRunResult[TOutput], result)
+        finally:
+            if state is not None:
+                delete_state()
 
     @abc.abstractmethod
     async def run_step(

planar/ai/state.py

@@ -0,0 +1,17 @@
+from typing import Any, Type, cast
+
+from planar.task_local import TaskLocal
+
+data: TaskLocal[Any] = TaskLocal()
+
+
+def set_state(ctx: Any):
+    return data.set(ctx)
+
+
+def get_state[T](_: Type[T]) -> T:
+    return cast(T, data.get())
+
+
+def delete_state():
+    return data.clear()

planar/app.py

@@ -130,9 +130,26 @@ class PlanarApp:
             prefix="/human-tasks",
         )
 
+        if self.config.data:
+            try:
+                from planar.routers.dataset_router import create_dataset_router
+
+                self.router_v1.include_router(
+                    create_dataset_router(),
+                    prefix="/datasets",
+                )
+            except ImportError:
+                logger.error(
+                    "Data dependencies not installed. Ensure you install the `data` optional dependency in your project (planar[data])"
+                )
+                raise
+
         self.router_v1.include_router(
             create_info_router(
-                title=title or "Planar API", description=description or "Planar API"
+                title=title or "Planar API",
+                description=description or "Planar API",
+                config=self.config,
+                registry=self._object_registry,
             ),
             prefix="",
         )

planar/data/connection.py

@@ -0,0 +1,108 @@
+import asyncio
+
+import ibis
+from ibis.backends.duckdb import Backend as DuckDBBackend
+
+from planar.config import PlanarConfig
+from planar.data.config import (
+    DuckDBCatalogConfig,
+    PostgresCatalogConfig,
+    SQLiteCatalogConfig,
+)
+from planar.data.exceptions import DataError
+from planar.files.storage.config import LocalDirectoryConfig, S3Config
+from planar.logging import get_logger
+from planar.session import get_config
+
+logger = get_logger(__name__)
+
+
+async def _create_connection(config: PlanarConfig) -> DuckDBBackend:
+    """Create Ibis DuckDB connection with Ducklake."""
+    data_config = config.data
+    if not data_config:
+        raise DataError("Data configuration not found")
+
+    # Connect to DuckDB with Ducklake extension
+    con = await asyncio.to_thread(ibis.duckdb.connect, extensions=["ducklake"])
+
+    # Build Ducklake connection string based on catalog type
+    catalog_config = data_config.catalog
+
+    match catalog_config:
+        case DuckDBCatalogConfig():
+            metadata_path = catalog_config.path
+        case PostgresCatalogConfig():
+            # Use connection components to build postgres connection string
+            metadata_path = f"postgres:dbname={catalog_config.db}"
+            if catalog_config.host:
+                metadata_path += f" host={catalog_config.host}"
+            if catalog_config.port:
+                metadata_path += f" port={catalog_config.port}"
+            if catalog_config.user:
+                metadata_path += f" user={catalog_config.user}"
+            if catalog_config.password:
+                metadata_path += f" password={catalog_config.password}"
+        case SQLiteCatalogConfig():
+            metadata_path = f"sqlite:{catalog_config.path}"
+        case _:
+            raise ValueError(f"Unsupported catalog type: {catalog_config.type}")
+
+    try:
+        await asyncio.to_thread(con.raw_sql, "INSTALL ducklake")
+        match catalog_config.type:
+            case "sqlite":
+                await asyncio.to_thread(con.raw_sql, "INSTALL sqlite;")
+            case "postgres":
+                await asyncio.to_thread(con.raw_sql, "INSTALL postgres;")
+        logger.debug("installed Ducklake extensions", catalog_type=catalog_config.type)
+    except Exception as e:
+        raise DataError(f"Failed to install Ducklake extensions: {e}") from e
+
+    # Build ATTACH statement
+    attach_sql = f"ATTACH 'ducklake:{metadata_path}' AS planar_ducklake"
+
+    # Add data path from storage config
+    storage = data_config.storage
+    if isinstance(storage, LocalDirectoryConfig):
+        data_path = storage.directory
+    elif isinstance(storage, S3Config):
+        data_path = f"s3://{storage.bucket_name}/"
+    else:
+        # Generic fallback
+        data_path = getattr(storage, "path", None) or getattr(storage, "directory", ".")
+
+    ducklake_catalog = data_config.catalog_name
+    attach_sql += f" (DATA_PATH '{data_path}'"
+    if catalog_config.type != "sqlite":
+        attach_sql += f", METADATA_SCHEMA '{ducklake_catalog}'"
+    attach_sql += ");"
+
+    # Attach to Ducklake
+    try:
+        await asyncio.to_thread(con.raw_sql, attach_sql)
+    except Exception as e:
+        raise DataError(f"Failed to attach to Ducklake: {e}") from e
+
+    await asyncio.to_thread(con.raw_sql, "USE planar_ducklake;")
+    logger.debug(
+        "connection created",
+        catalog=ducklake_catalog,
+        catalog_type=catalog_config.type,
+        attach_sql=attach_sql,
+    )
+
+    return con
+
+
+async def _get_connection() -> DuckDBBackend:
+    """Get Ibis connection to Ducklake."""
+    config = get_config()
+
+    if not config.data:
+        raise DataError(
+            "Data configuration not found. Please configure 'data' in your planar.yaml"
+        )
+
+    # TODO: Add cached connection pooling or memoize the connection
+    return await _create_connection(config)

planar/data/dataset.py

@@ -6,14 +6,11 @@ from typing import Literal, Self
 import ibis
 import polars as pl
 import pyarrow as pa
-from ibis.backends.duckdb import Backend as DuckDBBackend
 from ibis.common.exceptions import TableNotFound
 from pydantic import BaseModel
 
-from planar.config import PlanarConfig
-from planar.files.storage.config import LocalDirectoryConfig, S3Config
+from planar.data.connection import _get_connection
 from planar.logging import get_logger
-from planar.session import get_config
 
 from .exceptions import DataError, DatasetAlreadyExistsError, DatasetNotFoundError
 
@@ -32,6 +29,8 @@ class PlanarDataset(BaseModel):
     # TODO: Add snapshot version: no version = latest, otherwise time travel on read operations
     # TODO: Add partition support? A Dataset representation could be a table with a partition column
 
+    is_planar_dataset: bool = True
+
     model_config = {"arbitrary_types_allowed": True}
     # TODO: Add serialization metadata to make clear this is a dataset reference
     # like EntityField.
@@ -68,7 +67,8 @@ class PlanarDataset(BaseModel):
 
     async def exists(self) -> bool:
         """Check if the dataset exists in Ducklake."""
-        con = await self._get_connection()
+        con = await _get_connection()
+
         try:
             # TODO: Query for the table name directly
             tables = await asyncio.to_thread(con.list_tables)
@@ -79,16 +79,16 @@ class PlanarDataset(BaseModel):
 
     async def write(
         self,
-        data: pl.DataFrame | ibis.Table | list | dict,
+        data: pl.DataFrame | pl.LazyFrame | ibis.Table | list | dict,
         mode: Literal["overwrite", "append"] = "append",
     ) -> None:
         """Write data to the dataset.
 
         Args:
-            data: Data to write (Polars DataFrame, PyArrow Table, or Ibis expression)
+            data: Data to write (Polars DataFrame/LazyFrame, PyArrow Table, or Ibis expression)
             mode: Write mode - "append" or "overwrite"
         """
-        con = await self._get_connection()
+        con = await _get_connection()
         overwrite = mode == "overwrite"
 
         try:
@@ -99,7 +99,7 @@ class PlanarDataset(BaseModel):
             else:
                 # TODO: Explore if workflow context can be used to set metadata
                 # on the snapshot version for lineage
-                if isinstance(data, pl.DataFrame):
+                if isinstance(data, (pl.DataFrame, pl.LazyFrame)):
                     await asyncio.to_thread(
                         con.insert,
                         self.name,
@@ -133,7 +133,7 @@ class PlanarDataset(BaseModel):
         Returns:
             Ibis table expression that can be further filtered using Ibis methods
         """
-        con = await self._get_connection()
+        con = await _get_connection()
 
         try:
             table = await asyncio.to_thread(con.table, self.name)
@@ -162,102 +162,9 @@ class PlanarDataset(BaseModel):
 
     async def delete(self) -> None:
         """Delete the dataset."""
-        con = await self._get_connection()
+        con = await _get_connection()
         try:
             await asyncio.to_thread(con.drop_table, self.name, force=True)
             logger.info("deleted dataset", dataset_name=self.name)
         except Exception as e:
             raise DataError(f"Failed to delete dataset: {e}") from e
-
-    async def _get_connection(self) -> DuckDBBackend:
-        """Get Ibis connection to Ducklake."""
-        config = get_config()
-
-        if not config.data:
-            raise DataError(
-                "Data configuration not found. Please configure 'data' in your planar.yaml"
-            )
-
-        # TODO: Add cached connection pooling or memoize the connection
-        return await self._create_connection(config)
-
-    async def _create_connection(self, config: PlanarConfig) -> DuckDBBackend:
-        """Create Ibis DuckDB connection with Ducklake."""
-        data_config = config.data
-        if not data_config:
-            raise DataError("Data configuration not found")
-
-        # Connect to DuckDB with Ducklake extension
-        con = await asyncio.to_thread(ibis.duckdb.connect, extensions=["ducklake"])
-
-        # Build Ducklake connection string based on catalog type
-        catalog_config = data_config.catalog
-
-        if catalog_config.type == "duckdb":
-            metadata_path = catalog_config.path
-        elif catalog_config.type == "postgres":
-            # Use connection components to build postgres connection string
-            pg = catalog_config
-            metadata_path = f"postgres:dbname={pg.db}"
-            if pg.host:
-                metadata_path += f" host={pg.host}"
-            if pg.port:
-                metadata_path += f" port={pg.port}"
-            if pg.user:
-                metadata_path += f" user={pg.user}"
-            if pg.password:
-                metadata_path += f" password={pg.password}"
-        elif catalog_config.type == "sqlite":
-            metadata_path = f"sqlite:{catalog_config.path}"
-        else:
-            raise ValueError(f"Unsupported catalog type: {catalog_config.type}")
-
-        try:
-            await asyncio.to_thread(con.raw_sql, "INSTALL ducklake")
-            match catalog_config.type:
-                case "sqlite":
-                    await asyncio.to_thread(con.raw_sql, "INSTALL sqlite;")
-                case "postgres":
-                    await asyncio.to_thread(con.raw_sql, "INSTALL postgres;")
-            logger.debug(
-                "installed Ducklake extensions", catalog_type=catalog_config.type
-            )
-        except Exception as e:
-            raise DataError(f"Failed to install Ducklake extensions: {e}") from e
-
-        # Build ATTACH statement
-        attach_sql = f"ATTACH 'ducklake:{metadata_path}' AS planar_ducklake"
-
-        # Add data path from storage config
-        storage = data_config.storage
-        if isinstance(storage, LocalDirectoryConfig):
-            data_path = storage.directory
-        elif isinstance(storage, S3Config):
-            data_path = f"s3://{storage.bucket_name}/"
-        else:
-            # Generic fallback
-            data_path = getattr(storage, "path", None) or getattr(
-                storage, "directory", "."
-            )
-
-        ducklake_catalog = data_config.catalog_name
-        attach_sql += f" (DATA_PATH '{data_path}'"
-        if catalog_config.type != "sqlite":
-            attach_sql += f", METADATA_SCHEMA '{ducklake_catalog}'"
-        attach_sql += ");"
-
-        # Attach to Ducklake
-        try:
-            await asyncio.to_thread(con.raw_sql, attach_sql)
-        except Exception as e:
-            raise DataError(f"Failed to attach to Ducklake: {e}") from e
-
-        await asyncio.to_thread(con.raw_sql, "USE planar_ducklake;")
-        logger.debug(
-            "connection created",
-            catalog=ducklake_catalog,
-            catalog_type=catalog_config.type,
-            attach_sql=attach_sql,
-        )
-
-        return con

planar/data/utils.py

@@ -0,0 +1,89 @@
+import asyncio
+from typing import TypedDict
+
+import ibis.expr.datatypes as dt
+from ibis.common.exceptions import TableNotFound
+
+from planar.data.connection import _get_connection
+from planar.data.dataset import PlanarDataset
+from planar.data.exceptions import DatasetNotFoundError
+from planar.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+# TODO: consider connection pooling or memoize the connection
+
+
+async def list_datasets(limit: int = 100, offset: int = 0) -> list[PlanarDataset]:
+    conn = await _get_connection()
+    tables = await asyncio.to_thread(conn.list_tables)
+    return [PlanarDataset(name=table) for table in tables]
+
+
+async def list_schemas() -> list[str]:
+    METADATA_SCHEMAS = [
+        "information_schema",
+        # FIXME: why is list_databases returning pg_catalog
+        # if the ducklake catalog is sqlite?
+        "pg_catalog",
+    ]
+
+    conn = await _get_connection()
+
+    # in ibis, "databases" are schemas in the traditional sense
+    # e.g. psql: schema == ibis: database
+    # https://ibis-project.org/concepts/backend-table-hierarchy
+    schemas = await asyncio.to_thread(conn.list_databases)
+
+    return [schema for schema in schemas if schema not in METADATA_SCHEMAS]
+
+
+async def get_dataset(dataset_name: str, schema_name: str = "main") -> PlanarDataset:
+    # TODO: add schema_name as a parameter
+
+    dataset = PlanarDataset(name=dataset_name)
+
+    if not await dataset.exists():
+        raise DatasetNotFoundError(f"Dataset {dataset_name} not found")
+
+    return dataset
+
+
+async def get_dataset_row_count(dataset_name: str) -> int:
+    conn = await _get_connection()
+
+    try:
+        value = await asyncio.to_thread(
+            lambda conn, dataset_name: conn.table(dataset_name).count().to_polars(),
+            conn,
+            dataset_name,
+        )
+
+        assert isinstance(value, int), "Scalar must be an integer"
+
+        return value
+    except TableNotFound:
+        raise  # re-raise the exception and allow the caller to handle it
+
+
+class DatasetMetadata(TypedDict):
+    schema: dict[str, dt.DataType]
+    row_count: int
+
+
+async def get_dataset_metadata(
+    dataset_name: str, schema_name: str
+) -> DatasetMetadata | None:
+    conn = await _get_connection()
+
+    try:
+        schema, row_count = await asyncio.gather(
+            asyncio.to_thread(conn.get_schema, dataset_name, database=schema_name),
+            get_dataset_row_count(dataset_name),
+        )
+
+        return DatasetMetadata(schema=schema.fields, row_count=row_count)
+
+    except TableNotFound:
+        return None

planar/db/alembic/env.py

@@ -1,5 +1,7 @@
+from functools import wraps
 from logging.config import fileConfig
 
+import alembic.ddl.base as alembic_base
 from alembic import context
 from sqlalchemy import Connection, engine_from_config, pool
 
@@ -48,6 +50,28 @@ def include_name(name, type_, _):
         return True
 
 
+sqlite_schema_translate_map = {PLANAR_SCHEMA: None}
+
+
+def schema_translate_wrapper(f):
+    @wraps(f)
+    def format_table_name_with_schema(compiler, name, schema):
+        # when on sqlite, we need to translate the schema to None
+        is_sqlite = compiler.dialect.name == "sqlite"
+        if is_sqlite:
+            translated_schema = sqlite_schema_translate_map.get(schema, schema)
+        else:
+            translated_schema = schema
+        return f(compiler, name, translated_schema)
+
+    return format_table_name_with_schema
+
+
+alembic_base.format_table_name = schema_translate_wrapper(
+    alembic_base.format_table_name
+)
+
+
 def run_migrations_online() -> None:
     """Run migrations in 'online' mode.
 
@@ -102,7 +126,7 @@ def run_migrations_online() -> None:
         config_dict = config.get_section(config.config_ini_section, {})
         url = config_dict["sqlalchemy.url"]
         is_sqlite = url.startswith("sqlite://")
-        translate_map = {PLANAR_SCHEMA: None} if is_sqlite else {}
+        translate_map = sqlite_schema_translate_map if is_sqlite else {}
         connectable = engine_from_config(
             config_dict,
             prefix="sqlalchemy.",

planar/files/storage/azure_blob.py

@@ -278,7 +278,7 @@ class AzureBlobStorage(Storage):
 
         elif self.auth_method.name == "AZURE_AD":
             # Generate a User Delegation SAS signed with a user delegation key
-            start_time = datetime.utcnow()
+            start_time = datetime.now(UTC)
             user_delegation_key = await self.client.get_user_delegation_key(
                 key_start_time=start_time, key_expiry_time=expiry_time
             )

planar/registry_items.py

@@ -47,6 +47,7 @@ class RegisteredWorkflow:
     input_schema: dict[str, Any]
     output_schema: dict[str, Any]
     pydantic_model: Type[BaseModel]
+    is_interactive: bool
 
     @staticmethod
     def from_workflow(workflow: "WorkflowWrapper") -> "RegisteredWorkflow":
@@ -63,4 +64,5 @@ class RegisteredWorkflow:
                 workflow.original_fn
             ),
             pydantic_model=create_pydantic_model_for_workflow(workflow),
+            is_interactive=workflow.is_interactive,
         )