planar 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

planar/ai/agent_base.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import abc
+from dataclasses import dataclass, field
+from typing import (
+    Any,
+    Callable,
+    Coroutine,
+    Type,
+    cast,
+    overload,
+)
+
+from pydantic import BaseModel
+
+from planar.ai.models import AgentConfig, AgentEventEmitter, AgentRunResult
+from planar.logging import get_logger
+from planar.modeling.field_helpers import JsonSchema
+from planar.utils import P, R, T, U
+from planar.workflows import as_step
+from planar.workflows.models import StepType
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class AgentBase[
+    # TODO: add `= str` default when we upgrade to 3.13
+    TInput: BaseModel | str,
+    TOutput: BaseModel | str,
+](abc.ABC):
+    """An LLM-powered agent that can be called directly within workflows."""
+
+    name: str
+    system_prompt: str
+    output_type: Type[TOutput] | None = None
+    input_type: Type[TInput] | None = None
+    user_prompt: str = ""
+    tools: list[Callable] = field(default_factory=list)
+    max_turns: int = 2
+    model_parameters: dict[str, Any] = field(default_factory=dict)
+    event_emitter: AgentEventEmitter | None = None
+    durable: bool = True
+
+    # TODO: move here to serialize to frontend
+    #
+    # built_in_vars: Dict[str, str] = field(default_factory=lambda: {
+    #     "datetime_now": datetime.datetime.now().isoformat(),
+    #     "date_today": datetime.date.today().isoformat(),
+    # })
+
+    def __post_init__(self):
+        if self.input_type:
+            if (
+                not issubclass(self.input_type, BaseModel)
+                and self.input_type is not str
+            ):
+                raise ValueError(
+                    "input_type must be 'str' or a subclass of a Pydantic model"
+                )
+        if self.max_turns < 1:
+            raise ValueError("Max_turns must be greater than or equal to 1.")
+        if self.tools and self.max_turns <= 1:
+            raise ValueError(
+                "For tool calling to work, max_turns must be greater than 1."
+            )
+
+    def input_schema(self) -> JsonSchema | None:
+        if self.input_type is None:
+            return None
+        if self.input_type is str:
+            return None
+        assert issubclass(self.input_type, BaseModel), (
+            "input_type must be a subclass of BaseModel or str"
+        )
+        return self.input_type.model_json_schema()
+
+    def output_schema(self) -> JsonSchema | None:
+        if self.output_type is None:
+            return None
+        if self.output_type is str:
+            return None
+        assert issubclass(self.output_type, BaseModel), (
+            "output_type must be a subclass of BaseModel or str"
+        )
+        return self.output_type.model_json_schema()
+
+    @overload
+    async def __call__(
+        self: "AgentBase[TInput, str]",
+        input_value: TInput,
+    ) -> AgentRunResult[str]: ...
+
+    @overload
+    async def __call__(
+        self: "AgentBase[TInput, TOutput]",
+        input_value: TInput,
+    ) -> AgentRunResult[TOutput]: ...
+
+    def as_step_if_durable(
+        self,
+        func: Callable[P, Coroutine[T, U, R]],
+        step_type: StepType,
+        display_name: str | None = None,
+        return_type: Type[R] | None = None,
+    ) -> Callable[P, Coroutine[T, U, R]]:
+        if not self.durable:
+            return func
+        return as_step(
+            func,
+            step_type=step_type,
+            display_name=display_name or self.name,
+            return_type=return_type,
+        )
+
+    async def __call__(
+        self,
+        input_value: TInput,
+    ) -> AgentRunResult[Any]:
+        if self.input_type is not None and not isinstance(input_value, self.input_type):
+            raise ValueError(
+                f"Input value must be of type {self.input_type}, but got {type(input_value)}"
+            )
+        elif not isinstance(input_value, (str, BaseModel)):
+            # Should not happen based on type constraints, but just in case
+            # user does not have type checking enabled
+            raise ValueError(
+                "Input value must be a string or a Pydantic model if input_type is not provided"
+            )
+
+        if self.output_type is None:
+            run_step = self.as_step_if_durable(
+                self.run_step,
+                step_type=StepType.AGENT,
+                display_name=self.name,
+                return_type=AgentRunResult[str],
+            )
+        else:
+            run_step = self.as_step_if_durable(
+                self.run_step,
+                step_type=StepType.AGENT,
+                display_name=self.name,
+                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)
+
+    @abc.abstractmethod
+    async def run_step(
+        self,
+        input_value: TInput,
+    ) -> AgentRunResult[TOutput]: ...
+
+    @abc.abstractmethod
+    def get_model_str(self) -> str: ...
+
+    def to_config(self) -> AgentConfig:
+        return AgentConfig(
+            system_prompt=self.system_prompt,
+            user_prompt=self.user_prompt,
+            model=self.get_model_str(),
+            max_turns=self.max_turns,
+            model_parameters=self.model_parameters,
+        )

planar/ai/agent_utils.py

@@ -1,8 +1,4 @@
-import asyncio
 import inspect
-import json
-from collections.abc import AsyncGenerator
-from enum import Enum
 from typing import (
     Any,
     Callable,
@@ -20,77 +16,16 @@ from planar.ai.models import (
 from planar.files.models import PlanarFile
 from planar.logging import get_logger
 from planar.object_config import ConfigurableObjectType, ObjectConfigurationIO
-from planar.utils import utc_now
 from planar.workflows import step
 
 logger = get_logger(__name__)
 
 
-class AgentEventType(str, Enum):
-    """Valid event types that can be emitted by an Agent."""
+class ModelSpec(BaseModel):
+    """Pydantic model for AI model specifications."""
 
-    RESPONSE = "response"
-    TOOL_RESPONSE = "tool_response"
-    COMPLETED = "completed"
-    ERROR = "error"
-    THINK = "think"
-    TEXT = "text"
-
-
-class AgentEvent:
-    def __init__(
-        self,
-        event_type: AgentEventType,
-        data: BaseModel | str | None,
-    ):
-        self.event_type = event_type
-        self.data = data
-        self.timestamp = utc_now().isoformat()
-
-
-class AgentEventEmitter:
-    def __init__(self):
-        self.queue: asyncio.Queue[AgentEvent] = asyncio.Queue()
-
-    def emit(self, event_type: AgentEventType, data: BaseModel | str | None):
-        event = AgentEvent(event_type, data)
-        self.queue.put_nowait(event)
-
-    async def get_events(self) -> AsyncGenerator[str, None]:
-        while True:
-            event = await self.queue.get()
-
-            if isinstance(event.data, BaseModel):
-                data = {
-                    "data": event.data.model_dump(),
-                    "event_type": event.event_type,
-                }
-            else:
-                data = {
-                    "data": event.data,
-                    "event_type": event.event_type,
-                }
-
-            yield f"data: {json.dumps(data)}\n\n"
-
-            self.queue.task_done()
-
-            if event.event_type in (AgentEventType.COMPLETED, AgentEventType.ERROR):
-                break
-
-    def is_empty(self) -> bool:
-        """Check if the queue is empty."""
-        return self.queue.empty()
-
-
-# Define JsonData type as a union of valid JSON values
-JsonData = str | int | float | bool | None | dict[str, Any] | list[Any]
-
-
-class ToolCallResult(BaseModel):
-    tool_call_id: str
-    tool_call_name: str
-    content: BaseModel | JsonData
+    model_id: str
+    parameters: dict[str, Any] = {}
 
 
 def extract_files_from_model(

planar/ai/models.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum
 from typing import (
     Annotated,
     Any,
@@ -7,6 +8,7 @@ from typing import (
     List,
     Literal,
     Optional,
+    Protocol,
     TypeVar,
     Union,
 )
@@ -88,10 +90,23 @@ class ToolMessage(ModelMessage):
     tool_call_id: str  # ID of the tool call this is responding to
 
 
+# Define JsonData type as a union of valid JSON values
+JsonData = str | int | float | bool | None | dict[str, Any] | list[Any]
+
+
+class ToolCallResult(BaseModel):
+    tool_call_id: str
+    tool_call_name: str
+    content: BaseModel | JsonData
+
+
 class CompletionResponse[T: BaseModel | str](BaseModel):
     """Response object that may contain content or tool calls."""
 
     content: Optional[T] = None  # Content as str or parsed Pydantic model
+    text_content: Optional[str] = (
+        None  # Optional text content, if separate from structured output
+    )
     reasoning_content: Optional[str] = None  # Optional reasoning content
     tool_calls: Optional[List[ToolCall]] = None  # List of tool calls, if any
 
@@ -138,3 +153,18 @@ class AgentSerializeable(BaseModel):
 
     # TODO: actually fetch built_in_vars from agent object
     built_in_vars: dict[str, str] = Field(default_factory=dict)
+
+
+class AgentEventType(str, Enum):
+    """Valid event types that can be emitted by an Agent."""
+
+    RESPONSE = "response"
+    TOOL_RESPONSE = "tool_response"
+    COMPLETED = "completed"
+    ERROR = "error"
+    THINK = "think"
+    TEXT = "text"
+
+
+class AgentEventEmitter(Protocol):
+    def emit(self, event_type: AgentEventType, data: BaseModel | str | None): ...

planar/ai/pydantic_ai.py

@@ -1,8 +1,9 @@
 import base64
 import json
+import os
 import re
 import textwrap
-from typing import Any, Literal, Protocol, Type, cast
+from typing import Any, Type, cast
 
 from pydantic import BaseModel, ValidationError
 from pydantic_ai import BinaryContent
@@ -28,12 +29,13 @@ from pydantic_ai.messages import (
     UserContent,
     UserPromptPart,
 )
-from pydantic_ai.models import Model, ModelRequestParameters
+from pydantic_ai.models import KnownModelName, Model, ModelRequestParameters
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import ToolDefinition
 from pydantic_core import ErrorDetails
 
 from planar.ai import models as m
+from planar.files.models import PlanarFile
 from planar.logging import get_logger
 from planar.utils import partition
 
@@ -67,7 +69,65 @@ def format_validation_errors(errors: list[ErrorDetails], function: bool) -> str:
     return "\n".join(lines)
 
 
-async def build_file_map(messages: list[m.ModelMessage]) -> m.FileMap:
+async def openai_try_upload_file(
+    model: KnownModelName | Model, file: PlanarFile
+) -> m.FileIdContent | None:
+    # Currently pydanticAI doesn't support passing file_ids, but leaving the
+    # implementation here for when they add support.
+    return None
+
+    if file.content_type != "application/pdf":
+        # old implementation only does this for pdf files, so keep the behavior for now
+        return None
+
+    if isinstance(model, str) and not model.startswith("openai:"):
+        # not using openai provider
+        return None
+
+    try:
+        # make this code work with openai as optional dependency
+        from pydantic_ai.models.openai import OpenAIModel
+    except ImportError:
+        return None
+
+    if os.getenv("OPENAI_BASE_URL", None) is not None:
+        # cannot use OpenAI file upload if using a custom base url
+        return None
+
+    if (
+        isinstance(model, OpenAIModel)
+        and model.client.base_url.host != "api.openai.com"
+    ):
+        # same as above
+        return None
+
+    logger.debug("uploading pdf file to openai", filename=file.filename)
+
+    # use a separate AsyncClient instance since the model might be provided as a string
+    from openai import AsyncClient
+
+    client = AsyncClient()
+
+    # upload the file to the provider
+    openai_file = await client.files.create(
+        file=(
+            file.filename,
+            await file.get_content(),
+            file.content_type,
+        ),
+        purpose="user_data",
+    )
+    logger.info(
+        "uploaded pdf file to openai",
+        filename=file.filename,
+        openai_file_id=openai_file.id,
+    )
+    return m.FileIdContent(content=openai_file.id)
+
+
+async def build_file_map(
+    model: KnownModelName | Model, messages: list[m.ModelMessage]
+) -> m.FileMap:
     logger.debug("building file map", num_messages=len(messages))
     file_dict = {}
 
@@ -86,6 +146,12 @@ async def build_file_map(messages: list[m.ModelMessage]) -> m.FileMap:
                     content_type=file.content_type,
                 )
 
+                file_content_id = await openai_try_upload_file(model, file)
+                # TODO: add more `try_upload_file` implementations for other providers that support
+                if file_content_id is not None:
+                    file_dict[str(file.id)] = file_content_id
+                    continue
+
                 # For now we are not using uploaded files with Gemini, so convert all to base64
                 if file.content_type.startswith(
                     ("image/", "audio/", "video/", "application/pdf")
@@ -107,7 +173,9 @@ async def build_file_map(messages: list[m.ModelMessage]) -> m.FileMap:
     return m.FileMap(mapping=file_dict)
 
 
-async def prepare_messages(messages: list[m.ModelMessage]) -> list[Any]:
+async def prepare_messages(
+    model: KnownModelName | Model, messages: list[m.ModelMessage]
+) -> list[Any]:
     """Prepare messages from Planar representations into the format expected by PydanticAI.
 
     Args:
@@ -118,7 +186,7 @@ async def prepare_messages(messages: list[m.ModelMessage]) -> list[Any]:
         List of messages in PydanticAI format
     """
     pydantic_messages: list[ModelMessage] = []
-    file_map = await build_file_map(messages)
+    file_map = await build_file_map(model, messages)
 
     def append_request_part(part: ModelRequestPart):
         last = (
@@ -197,10 +265,6 @@ async def prepare_messages(messages: list[m.ModelMessage]) -> list[Any]:
     return pydantic_messages
 
 
-class StreamEventHandler(Protocol):
-    def emit(self, event: Literal["text", "think"], data: str) -> None: ...
-
-
 def setup_native_structured_output(
     request_params: ModelRequestParameters,
     output_type: Type[BaseModel],
@@ -262,12 +326,14 @@ def return_native_structured_output[TOutput: BaseModel](
         result = m.CompletionResponse(
             content=output_type.model_validate_json(content),
             tool_calls=final_tool_calls,
+            text_content=content,
             reasoning_content=thinking,
         )
         logger.info(
             "model run completed with structured output",
             content=result.content,
             reasoning_content=result.reasoning_content,
+            text_content=content,
             tool_calls=result.tool_calls,
         )
         return result
@@ -291,6 +357,7 @@ def return_tool_structured_output[TOutput: BaseModel](
         result = m.CompletionResponse(
             content=output_type.model_validate(final_result_tc.arguments),
             tool_calls=tool_calls,
+            text_content=content,
             reasoning_content=thinking,
         )
         logger.info(
@@ -315,12 +382,12 @@ class ModelRunResponse[TOutput: BaseModel | str](BaseModel):
 
 
 async def model_run[TOutput: BaseModel | str](
-    model: Model | str,
+    model: Model | KnownModelName,
     max_extra_turns: int,
     model_settings: dict[str, Any] | None = None,
     messages: list[m.ModelMessage] = [],
     tools: list[m.ToolDefinition] = [],
-    event_handler: StreamEventHandler | None = None,
+    event_handler: m.AgentEventEmitter | None = None,
     output_type: Type[TOutput] = str,
 ) -> ModelRunResponse[TOutput]:
     # assert that the caller doesn't provide a tool called "final_result"
@@ -350,11 +417,11 @@ async def model_run[TOutput: BaseModel | str](
 
     structured_output = issubclass(output_type, BaseModel)
 
-    def emit(event_type: Literal["text", "think"], content: str):
+    def emit(event_type: m.AgentEventType, content: str):
         if event_handler:
             event_handler.emit(event_type, content)
 
-    history = await prepare_messages(messages=messages)
+    history = await prepare_messages(model, messages=messages)
 
     if structured_output:
         if supports_native_structured_output:
@@ -383,10 +450,10 @@ async def model_run[TOutput: BaseModel | str](
                     case PartStartEvent(part=part):
                         response_parts.append(part)
                         if isinstance(part, TextPart):
-                            emit("text", part.content)
+                            emit(m.AgentEventType.TEXT, part.content)
                             text_buffer.append(part.content)
                         elif isinstance(part, ThinkingPart):
-                            emit("think", part.content)
+                            emit(m.AgentEventType.THINK, part.content)
                             think_buffer.append(part.content)
                         elif isinstance(part, ToolCallPart):
                             if current_tool_call is not None:
@@ -412,14 +479,14 @@ async def model_run[TOutput: BaseModel | str](
                         current = response_parts[-1]
                         if isinstance(delta, TextPartDelta):
                             assert isinstance(current, TextPart)
-                            emit("text", delta.content_delta)
+                            emit(m.AgentEventType.TEXT, delta.content_delta)
                             text_buffer.append(delta.content_delta)
                             current.content += delta.content_delta
                         elif (
                             isinstance(delta, ThinkingPartDelta) and delta.content_delta
                         ):
                             assert isinstance(current, ThinkingPart)
-                            emit("think", delta.content_delta)
+                            emit(m.AgentEventType.THINK, delta.content_delta)
                             think_buffer.append(delta.content_delta)
                             current.content += delta.content_delta
                         elif isinstance(delta, ToolCallPartDelta):
@@ -479,6 +546,7 @@ async def model_run[TOutput: BaseModel | str](
             return ModelRunResponse(
                 response=m.CompletionResponse(
                     tool_calls=final_tool_calls,
+                    text_content=content,
                     reasoning_content=thinking,
                 ),
                 extra_turns_used=extra_turns_used,
@@ -555,6 +623,7 @@ async def model_run[TOutput: BaseModel | str](
             m.CompletionResponse(
                 content=content,
                 tool_calls=final_tool_calls,
+                text_content=content,
                 reasoning_content=thinking,
             ),
         )

planar/ai/test_agent_serialization.py

@@ -56,7 +56,7 @@ def test_agent_with_tools():
         name="test_agent_with_tools",
         system_prompt="System with tools",
         user_prompt="User: {input}",
-        model="anthropic:claude-3-sonnet",
+        model="anthropic:claude-3-5-sonnet-latest",
         max_turns=5,
         tools=[test_tool],
     )

planar/app.py

@@ -37,7 +37,6 @@ from planar.security.authorization import PolicyService, policy_service_context
 from planar.session import config_var, session_context
 from planar.sse.proxy import SSEProxy
 from planar.workflows import (
-    Workflow,
     WorkflowNotification,
     WorkflowNotificationCallback,
     WorkflowOrchestrator,
@@ -169,13 +168,8 @@ class PlanarApp:
             return
 
         def on_workflow_notification(notification: WorkflowNotification):
-            workflow_id = (
-                notification.data.id
-                if isinstance(notification.data, Workflow)
-                else notification.data.workflow_id
-            )
             self.sse_proxy.push(
-                f"{notification.kind.value}:{workflow_id}",
+                f"{notification.kind.value}:{notification.workflow_id}",
                 notification.data.model_dump(mode="json"),
             )
 

planar/config.py

@@ -21,6 +21,7 @@ from pydantic import (
 )
 from sqlalchemy import URL, make_url
 
+from planar.data.config import DataConfig
 from planar.files.storage.config import LocalDirectoryConfig, StorageConfig
 from planar.logging import get_logger
 
@@ -225,6 +226,7 @@ class PlanarConfig(BaseModel):
     logging: dict[str, LoggerConfig] | None = None
     use_alembic: bool | None = True
     otel: OtelConfig | None = None
+    data: DataConfig | None = None
 
     # forbid extra keys in the config to prevent accidental misconfiguration
     model_config = ConfigDict(extra="forbid")

planar/data/__init__.py

@@ -0,0 +1,17 @@
+from typing import TYPE_CHECKING
+
+from planar.dependencies import lazy_exports
+
+lazy_exports(
+    __name__,
+    {
+        "PlanarDataset": (".dataset", "PlanarDataset"),
+    },
+)
+
+if TYPE_CHECKING:
+    from .dataset import PlanarDataset
+
+    __all__ = [
+        "PlanarDataset",
+    ]

planar/data/config.py

@@ -0,0 +1,49 @@
+"""Configuration for Planar data module."""
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from planar.files.storage.config import StorageConfig
+
+
+class DuckDBCatalogConfig(BaseModel):
+    """Configuration for DuckDB catalog backend."""
+
+    type: Literal["duckdb"]
+    path: str  # Path to .ducklake file
+
+
+class PostgresCatalogConfig(BaseModel):
+    """Configuration for PostgreSQL catalog backend."""
+
+    type: Literal["postgres"]
+    host: str | None = None
+    port: int | None = None
+    user: str | None = None
+    password: str | None = None
+    db: str
+
+
+class SQLiteCatalogConfig(BaseModel):
+    """Configuration for SQLite catalog backend."""
+
+    type: Literal["sqlite"]
+    path: str  # Path to .sqlite file
+
+
+# Discriminated union for catalog configurations
+CatalogConfig = Annotated[
+    DuckDBCatalogConfig | PostgresCatalogConfig | SQLiteCatalogConfig,
+    Field(discriminator="type"),
+]
+
+
+class DataConfig(BaseModel):
+    """Configuration for data features."""
+
+    catalog: CatalogConfig
+    storage: StorageConfig  # Reuse existing StorageConfig from files
+
+    # Optional settings
+    catalog_name: str = "planar_data"  # Default catalog name in Ducklake