planar 0.5.0__py3-none-any.whl

planar/ai/agent_utils.py

@@ -0,0 +1,205 @@
+import asyncio
+import inspect
+import json
+from collections.abc import AsyncGenerator
+from enum import Enum
+from typing import (
+    Any,
+    Callable,
+    Dict,
+)
+
+from jinja2 import StrictUndefined, TemplateError
+from jinja2.sandbox import SandboxedEnvironment
+from pydantic import BaseModel, create_model
+
+from planar.ai.models import (
+    AgentConfig,
+    ToolDefinition,
+)
+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."""
+
+    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
+
+
+def extract_files_from_model(
+    model: BaseModel | str | None,
+) -> list[PlanarFile]:
+    """
+    Extract files from a Pydantic model. We extract any top-level or nested fields
+    that are of type `PlanarFile`, or are a list of `PlanarFile`.
+
+    Args:
+        model: The Pydantic model to extract files from.
+
+    Returns:
+        A list of PlanarFile objects.
+    """
+
+    if model is None:
+        return []
+
+    if isinstance(model, PlanarFile):
+        return [model]
+
+    if isinstance(model, str):
+        return []
+
+    files: list[PlanarFile] = []
+    for field_name in type(model).model_fields:
+        value = getattr(model, field_name)
+        match value:
+            case PlanarFile() as f:
+                files.append(f)
+            case BaseModel():
+                files.extend(extract_files_from_model(value))
+            case [*items]:
+                files.extend(item for item in items if isinstance(item, PlanarFile))
+            case _:
+                pass
+    return files
+
+
+# Jinja environment for safely rendering templates
+_JINJA_ENV = SandboxedEnvironment(undefined=StrictUndefined)
+
+
+def render_template(template: str, context: Dict[str, Any]) -> str:
+    """Render a template string using a sandboxed Jinja environment."""
+    try:
+        return _JINJA_ENV.from_string(template).render(context)
+    except TemplateError as exc:
+        logger.exception("error rendering jinja template")
+        raise ValueError(f"Error rendering prompt: {exc}") from exc
+
+
+agent_configuration = ObjectConfigurationIO(AgentConfig, ConfigurableObjectType.AGENT)
+
+
+@step(display_name="Agent Config")
+async def get_agent_config(agent_name: str, agent_config: AgentConfig) -> AgentConfig:
+    """
+    Retrieve agent configuration overrides from the database.
+
+    Args:
+        agent_name: Name of the agent instance.
+
+    Returns:
+        AgentOverrideConfig
+    """
+    logger.debug("getting agent config", agent_name=agent_name)
+    configs = await agent_configuration.read_configs_with_default(
+        agent_name, agent_config
+    )
+
+    active_config = next((config for config in configs if config.active), None)
+
+    if not active_config:
+        logger.warning("no active configuration found for agent", agent_name=agent_name)
+        raise ValueError(f"No active configuration found for agent {agent_name}")
+
+    logger.info(
+        "active configuration found for agent",
+        version=active_config.version,
+        agent_name=agent_name,
+    )
+    return active_config.data
+
+
+def create_tool_definition(tool_fn: Callable) -> ToolDefinition:
+    """Create a ToolDefinition from a function using a Pydantic model."""
+    sig = inspect.signature(tool_fn)
+    doc = inspect.getdoc(tool_fn) or ""
+    name = tool_fn.__name__
+
+    fields = {}
+    for param_name, param in sig.parameters.items():
+        param_type = (
+            param.annotation if param.annotation != inspect.Parameter.empty else Any
+        )
+        default_value = (
+            param.default if param.default != inspect.Parameter.empty else ...
+        )
+        fields[param_name] = (param_type, default_value)
+
+    model_name = f"{name.capitalize()}Parameters"
+    parameters_model = create_model(
+        model_name, __config__={"extra": "forbid"}, **fields
+    )
+
+    return ToolDefinition(
+        name=name,
+        description=doc,
+        parameters=parameters_model.model_json_schema(),
+    )

planar/ai/models.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from typing import (
+    Annotated,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+)
+
+from pydantic import BaseModel, Field
+
+from planar.files.models import PlanarFile
+from planar.modeling.field_helpers import JsonSchema
+from planar.object_config.object_config import (
+    ObjectConfigurationBase,
+)
+
+# Type variable for output_type, which can be str or a Pydantic model
+T = TypeVar("T", bound=Union[str, BaseModel])
+
+
+# Agent configuration overrides.
+# This model allows storing configurations that override the default
+# settings defined in Agent instances.
+class AgentConfig(BaseModel):
+    system_prompt: str
+    user_prompt: str = Field()
+    model: str = Field()
+    max_turns: int = Field()
+    model_parameters: Dict[str, Any] = Field(default_factory=dict)
+
+
+class ToolDefinition(BaseModel):
+    """Defines a tool that the model can call."""
+
+    name: str
+    description: str
+    # This is a json schema string
+    parameters: Dict[str, Any]
+
+
+class ToolCall(BaseModel):
+    """Represents a tool call made by the model."""
+
+    id: Optional[str] = None  # Optional ID, included by providers like OpenAI
+    name: str
+    arguments: Dict[str, Any]  # Arguments for the tool call, as a dictionary
+
+
+class ToolResponse(BaseModel):
+    """Represents a response to a tool call."""
+
+    tool_call_id: Optional[str] = None  # ID of the corresponding tool call
+    content: str  # String content of the tool response
+
+
+class ModelMessage(BaseModel):
+    """Base class for messages exchanged with LLM providers."""
+
+    content: Optional[str] = None
+
+
+class AssistantMessage(ModelMessage):
+    """Message from the assistant, may include tool calls."""
+
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+class UserMessage(ModelMessage):
+    """Message from the user."""
+
+    files: Optional[list[PlanarFile]] = None
+
+
+class SystemMessage(ModelMessage):
+    """System message that provides context/instructions."""
+
+    pass
+
+
+class ToolMessage(ModelMessage):
+    """Tool message containing a tool response."""
+
+    tool_call_id: str  # ID of the tool call this is responding to
+
+
+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
+    reasoning_content: Optional[str] = None  # Optional reasoning content
+    tool_calls: Optional[List[ToolCall]] = None  # List of tool calls, if any
+
+
+class AgentRunResult[TOutput: BaseModel | str](BaseModel):
+    output: TOutput
+
+
+class Base64Content(BaseModel):
+    type: Literal["base64"] = "base64"
+    content: str
+    content_type: str
+
+    def __repr__(self):
+        return f"Base64Content(content_type={self.content_type}, content={self.content[:10]}...)"
+
+
+class URLContent(BaseModel):
+    type: Literal["url"] = "url"
+    content: str
+
+
+class FileIdContent(BaseModel):
+    type: Literal["file_id"] = "file_id"
+    content: str
+
+
+FileContent = Annotated[
+    Union[Base64Content, URLContent, FileIdContent],
+    Field(discriminator="type"),
+]
+
+
+class FileMap(BaseModel):
+    mapping: dict[str, FileContent]
+
+
+class AgentSerializeable(BaseModel):
+    name: str
+    input_schema: JsonSchema | None = None
+    output_schema: JsonSchema | None = None
+    tool_definitions: list[dict[str, Any]]
+    configs: list[ObjectConfigurationBase[AgentConfig]]
+
+    # TODO: actually fetch built_in_vars from agent object
+    built_in_vars: dict[str, str] = Field(default_factory=dict)