planar 0.9.2__py3-none-any.whl → 0.10.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(
@@ -289,6 +290,7 @@ class Agent[
                     tool_result = await self.as_step_if_durable(
                         tool_fn,
                         step_type=StepType.TOOL_CALL,
+                        display_name=tool_call.name,
                     )(**tool_call.arguments)
                     logger.info(
                         "tool executed by agent",

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/ai/test_agent_tool_step_display.py

@@ -0,0 +1,78 @@
+import os
+from unittest.mock import patch
+
+from sqlmodel import col, select
+
+from planar.ai import models as m
+from planar.ai.agent import Agent
+from planar.ai.pydantic_ai import ModelRunResponse
+from planar.workflows.decorators import workflow
+from planar.workflows.execution import execute
+from planar.workflows.models import StepType, WorkflowStep
+
+
+async def test_agent_tool_step_has_display_name(session):
+    async def add(a: int, b: int) -> int:
+        return a + b
+
+    # Prepare mocked model responses: first triggers a tool call, then returns final content
+    first = ModelRunResponse[str](
+        response=m.CompletionResponse[str](
+            content=None,
+            tool_calls=[
+                m.ToolCall(id="call_1", name="add", arguments={"a": 2, "b": 3})
+            ],
+            text_content="",
+            reasoning_content=None,
+        ),
+        extra_turns_used=0,
+    )
+    second = ModelRunResponse[str](
+        response=m.CompletionResponse[str](
+            content="5",
+            tool_calls=[],
+            text_content="5",
+            reasoning_content=None,
+        ),
+        extra_turns_used=0,
+    )
+
+    responses = [first, second]
+
+    async def fake_model_run(*args, **kwargs):
+        assert responses, "No more fake responses configured"
+        return responses.pop(0)
+
+    # Patch the model run to avoid any network/model dependency
+    # Use unittest.mock.patch context managers to ensure cleanup
+    with (
+        patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=False),
+        patch("planar.ai.agent.model_run", side_effect=fake_model_run),
+    ):
+        agent = Agent[str, str, None](
+            name="test_agent",
+            system_prompt="",
+            user_prompt="",
+            model="openai:gpt-4o-mini",
+            tools=[add],
+            max_turns=3,
+        )
+
+        @workflow()
+        async def run_agent():
+            result = await agent("please add")
+            return result.output
+
+        wf = await run_agent.start()
+        result = await execute(wf)
+    assert result == "5"
+
+    steps = (
+        await session.exec(select(WorkflowStep).order_by(col(WorkflowStep.step_id)))
+    ).all()
+    # Ensure there is a tool call step with the display name set to the tool name
+    tool_steps = [s for s in steps if s.step_type == StepType.TOOL_CALL]
+    assert tool_steps, "Expected at least one TOOL_CALL step recorded"
+    assert any(s.display_name == "add" for s in tool_steps), (
+        f"Expected a TOOL_CALL step with display_name 'add', got {[s.display_name for s in tool_steps]}"
+    )

planar/app.py

@@ -28,6 +28,7 @@ from planar.routers import (
     create_workflow_router,
 )
 from planar.routers.agents_router import create_agent_router
+from planar.routers.dataset_router import create_dataset_router
 from planar.routers.entity_router import create_entities_router
 from planar.routers.object_config_router import create_object_config_router
 from planar.routers.rule import create_rule_router
@@ -129,6 +130,10 @@ class PlanarApp:
             create_human_task_routes(),
             prefix="/human-tasks",
         )
+        self.router_v1.include_router(
+            create_dataset_router(),
+            prefix="/datasets",
+        )
 
         self.router_v1.include_router(
             create_info_router(

planar/cli.py

@@ -175,9 +175,17 @@ def run_command(
 
         try:
             result = subprocess.run(
-                ["uv", "run", str(app_path)], env=os.environ.copy(), check=True
+                ["uv", "run", str(app_path)], env=os.environ.copy(), check=False
             )
+            if result.returncode != 0:
+                typer.echo(
+                    f"Error running script: Process exited with code {result.returncode}",
+                    err=True,
+                )
             raise typer.Exit(code=result.returncode)
+        except typer.Exit:
+            # Re-raise typer.Exit without modification
+            raise
         except subprocess.CalledProcessError as e:
             typer.echo(f"Error running script: {e}", err=True)
             raise typer.Exit(code=e.returncode)

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