agentpool 2.1.9__py3-none-any.whl

agentpool_config/conditions.py

@@ -0,0 +1,301 @@
+"""Condition configuration."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from datetime import timedelta
+from typing import TYPE_CHECKING, Annotated, Any, Literal, assert_never
+
+from pydantic import ConfigDict, Field, ImportString
+from schemez import Schema
+
+
+if TYPE_CHECKING:
+    from agentpool.talk.registry import EventContext
+
+
+class ConnectionCondition(Schema):
+    """Base class for connection control conditions."""
+
+    type: str = Field(init=False, title="Condition type")
+    """Discriminator for condition types."""
+
+    name: str | None = Field(
+        default=None,
+        examples=["exit_condition", "message_filter", "cost_check"],
+        title="Condition name",
+    )
+    """Optional name for the condition for referencing."""
+
+    model_config = ConfigDict(frozen=True)
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if condition is met."""
+        raise NotImplementedError
+
+
+class Jinja2Condition(ConnectionCondition):
+    """Triggers based on Jinja2 template."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Jinja2 Template Condition"})
+
+    type: Literal["jinja2"] = Field("jinja2", init=False)
+    """Jinja2 template-based condition."""
+
+    template: str = Field(
+        examples=[
+            "{{ ctx.stats.message_count > 10 }}",
+            "{{ 'error' in ctx.message.content }}",
+            "{{ ctx.stats.total_cost > 1.0 }}",
+        ],
+        title="Jinja2 template",
+    )
+    """Jinja2 template to evaluate."""
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        from jinjarope import Environment
+
+        from agentpool.utils.now import get_now
+
+        env = Environment(trim_blocks=True, lstrip_blocks=True, enable_async=True)
+        template = env.from_string(self.template)
+        result = await template.render_async(ctx=context, now=get_now())
+        return result.strip().lower() == "true" or bool(result)
+
+
+class WordMatchCondition(ConnectionCondition):
+    """Triggers when word/phrase is found in message."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Word Match Condition"})
+
+    type: Literal["word_match"] = Field("word_match", init=False)
+    """Word-comparison-based condition."""
+
+    words: list[str] = Field(
+        examples=[["error", "failed"], ["complete", "done", "finished"], ["urgent"]],
+        title="Words to match",
+    )
+    """Words or phrases to match in messages."""
+
+    case_sensitive: bool = Field(default=False, title="Case sensitive matching")
+    """Whether to match case-sensitively."""
+
+    mode: Literal["any", "all"] = Field(
+        default="any",
+        examples=["any", "all"],
+        title="Matching mode",
+    )
+    """Match mode:
+    - any: Trigger if any word matches
+    - all: Require all words to match
+    """
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Triggers if message contains specified words."""
+        text = str(context.message.content)
+        if not self.case_sensitive:
+            text = text.lower()
+            words = [w.lower() for w in self.words]
+        else:
+            words = self.words
+
+        matches = [w in text for w in words]
+        return all(matches) if self.mode == "all" else any(matches)
+
+
+class MessageCountCondition(ConnectionCondition):
+    """Triggers after N messages."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Message Count Condition"})
+
+    type: Literal["message_count"] = Field("message_count", init=False)
+    """Message-count-based condition."""
+
+    max_messages: int = Field(gt=0, examples=[10, 50, 100], title="Maximum message count")
+    """Maximum number of messages before triggering."""
+
+    count_mode: Literal["total", "per_agent"] = Field(
+        default="total",
+        examples=["total", "per_agent"],
+        title="Message counting mode",
+    )
+    """How to count messages:
+    - total: All messages in conversation
+    - per_agent: Messages from each agent separately
+    """
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if message count threshold is reached."""
+        if self.count_mode == "total":
+            return context.stats.message_count >= self.max_messages
+
+        # Count per agent
+        messages = [m for m in context.stats.messages if m.name == context.message.name]
+        return len(messages) >= self.max_messages
+
+
+class TimeCondition(ConnectionCondition):
+    """Triggers after time period."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Time-based Condition"})
+
+    type: Literal["time"] = Field("time", init=False)
+    """Time-based condition."""
+
+    duration: timedelta = Field(title="Duration threshold")
+    """How long the connection should stay active."""
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if time duration has elapsed."""
+        from agentpool.utils.now import get_now
+
+        elapsed = get_now() - context.stats.start_time
+        return elapsed >= self.duration
+
+
+class TokenThresholdCondition(ConnectionCondition):
+    """Triggers after token threshold is reached."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Token Threshold Condition"})
+
+    type: Literal["token_threshold"] = Field("token_threshold", init=False)
+    """Type discriminator."""
+
+    max_tokens: int = Field(gt=0, examples=[1000, 4000, 8000], title="Maximum token count")
+    """Maximum number of tokens allowed."""
+
+    count_type: Literal["total", "prompt", "completion"] = Field(
+        default="total",
+        examples=["total", "prompt", "completion"],
+        title="Token counting type",
+    )
+    """What tokens to count:
+    - total: All tokens used
+    - prompt: Only prompt tokens
+    - completion: Only completion tokens
+    """
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if token threshold is reached."""
+        if not context.message.cost_info:
+            return False
+
+        match self.count_type:
+            case "total":
+                return context.stats.token_count >= self.max_tokens
+            case "prompt":
+                return context.message.cost_info.token_usage.input_tokens >= self.max_tokens
+            case "completion":
+                return context.message.cost_info.token_usage.output_tokens >= self.max_tokens
+            case _ as unreachable:
+                assert_never(unreachable)
+
+
+class CostCondition(ConnectionCondition):
+    """Triggers when cost threshold is reached."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Cost Condition"})
+
+    type: Literal["cost"] = Field("cost", init=False)
+    """Cost-based condition."""
+
+    max_cost: float = Field(gt=0.0, examples=[1.0, 5.0, 10.0], title="Maximum cost threshold")
+    """Maximum cost in USD."""
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if cost limit is reached."""
+        return context.stats.total_cost >= self.max_cost
+
+
+class CostLimitCondition(ConnectionCondition):
+    """Triggers when cost limit is reached."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Cost Limit Condition"})
+
+    type: Literal["cost_limit"] = Field("cost_limit", init=False)
+    """Cost-limit condition."""
+
+    max_cost: float = Field(gt=0.0, examples=[1.0, 5.0, 10.0], title="Cost limit")
+    """Maximum cost in USD before triggering."""
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if cost limit is reached."""
+        if not context.message.cost_info:
+            return False
+        return float(context.message.cost_info.total_cost) >= self.max_cost
+
+
+class CallableCondition(ConnectionCondition):
+    """Custom predicate function."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "Callable Condition"})
+
+    type: Literal["callable"] = Field("callable", init=False)
+    """Condition based on an import path pointing to a predicate."""
+
+    predicate: ImportString[Callable[..., bool | Awaitable[bool]]] = Field(
+        examples=["mymodule.check_condition", "utils.predicates:is_urgent"],
+        title="Predicate function",
+    )
+    """Function to evaluate condition:
+    Args:
+        message: Current message being processed
+        stats: Current connection statistics
+    Returns:
+        Whether condition is met
+    """
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Execute predicate function."""
+        from agentpool.utils.inspection import execute
+
+        return await execute(self.predicate, context.message, context.stats)
+
+
+class AndCondition(ConnectionCondition):
+    """Require all conditions to be met."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "AND Condition"})
+
+    type: Literal["and"] = Field("and", init=False)
+    """Condition to AND-combine multiple conditions."""
+
+    conditions: list[ConnectionCondition] = Field(title="AND conditions")
+    """List of conditions to check."""
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if all conditions are met."""
+        results = [await c.check(context) for c in self.conditions]
+        return all(results)
+
+
+class OrCondition(ConnectionCondition):
+    """Require any condition to be met."""
+
+    model_config = ConfigDict(json_schema_extra={"title": "OR Condition"})
+
+    type: Literal["or"] = Field("or", init=False)
+    """Condition to OR-combine multiple conditions."""
+
+    conditions: list[ConnectionCondition] = Field(title="OR conditions")
+    """List of conditions to check."""
+
+    async def check(self, context: EventContext[Any]) -> bool:
+        """Check if any condition is met."""
+        results = [await c.check(context) for c in self.conditions]
+        return any(results)
+
+
+# Union type for condition validation
+Condition = Annotated[
+    WordMatchCondition
+    | MessageCountCondition
+    | TimeCondition
+    | TokenThresholdCondition
+    | CostLimitCondition
+    | CallableCondition
+    | Jinja2Condition
+    | AndCondition
+    | OrCondition,
+    Field(discriminator="type"),
+]

agentpool_config/converters.py

@@ -0,0 +1,30 @@
+"""Converter configuration."""
+
+from __future__ import annotations
+
+from docler.configs import ConverterConfig
+from pydantic import ConfigDict, Field
+from schemez import Schema
+
+
+class ConversionConfig(Schema):
+    """Global conversion configuration."""
+
+    providers: list[ConverterConfig] | None = Field(default=None, title="Converter providers")
+    """List of configured converter providers."""
+
+    default_provider: str | None = Field(
+        default=None,
+        examples=["markitdown", "youtube", "whisper_api"],
+        title="Default provider",
+    )
+    """Name of default provider for conversions."""
+
+    max_size: int | None = Field(
+        default=None,
+        examples=[1048576, 10485760, 52428800],
+        title="Global size limit",
+    )
+    """Global size limit for all converters."""
+
+    model_config = ConfigDict(frozen=True)

agentpool_config/durable.py

@@ -0,0 +1,331 @@
+"""Durable execution configuration models."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import Annotated, Literal
+
+from pydantic import ConfigDict, Field, field_validator
+from schemez import Schema
+
+from agentpool.utils.parse_time import parse_time_period
+
+
+class BaseDurableExecutionConfig(Schema):
+    """Base configuration for durable execution providers."""
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "x-icon": "octicon:sync-16",
+            "x-doc-title": "Durable Execution Configuration",
+            "title": "Base Durable Execution Configuration",
+        }
+    )
+
+    type: str = Field(init=False, title="Provider type")
+    """Durable execution provider type discriminator."""
+
+
+class TemporalActivityConfig(Schema):
+    """Temporal activity configuration options.
+
+    Maps to temporalio.workflow.ActivityConfig.
+    """
+
+    start_to_close_timeout: str | timedelta | None = Field(
+        default="60s",
+        title="Start to close timeout",
+        examples=["60s", "5m", "1h"],
+    )
+    """Maximum time for a single activity execution attempt."""
+
+    schedule_to_close_timeout: str | timedelta | None = Field(
+        default=None,
+        title="Schedule to close timeout",
+        examples=["5m", "1h"],
+    )
+    """Total time for all activity attempts including retries."""
+
+    schedule_to_start_timeout: str | timedelta | None = Field(
+        default=None,
+        title="Schedule to start timeout",
+        examples=["30s", "2m"],
+    )
+    """Maximum time activity can wait in task queue."""
+
+    heartbeat_timeout: str | timedelta | None = Field(
+        default=None,
+        title="Heartbeat timeout",
+        examples=["10s", "30s"],
+    )
+    """Maximum time between activity heartbeats."""
+
+    @field_validator(
+        "start_to_close_timeout",
+        "schedule_to_close_timeout",
+        "schedule_to_start_timeout",
+        "heartbeat_timeout",
+        mode="before",
+    )
+    @classmethod
+    def parse_duration(cls, v: str | timedelta | None) -> timedelta | None:
+        """Parse string duration to timedelta."""
+        if v is None or isinstance(v, timedelta):
+            return v
+        return parse_time_period(v)
+
+
+class TemporalRetryPolicy(Schema):
+    """Temporal retry policy configuration.
+
+    Maps to temporalio.common.RetryPolicy.
+    """
+
+    model_config = ConfigDict(json_schema_extra={"title": "Temporal Retry Policy Configuration"})
+
+    initial_interval: str | timedelta = Field(
+        default="1s",
+        title="Initial retry interval",
+        examples=["1s", "5s", "10s"],
+    )
+    """Initial backoff interval."""
+
+    backoff_coefficient: float = Field(
+        default=2.0,
+        ge=1.0,
+        title="Backoff coefficient",
+    )
+    """Coefficient for exponential backoff."""
+
+    maximum_interval: str | timedelta | None = Field(
+        default=None,
+        title="Maximum interval",
+        examples=["5m", "10m"],
+    )
+    """Maximum backoff interval."""
+
+    maximum_attempts: int = Field(
+        default=0,
+        ge=0,
+        title="Maximum attempts",
+    )
+    """Maximum number of attempts (0 = unlimited)."""
+
+    non_retryable_error_types: list[str] = Field(
+        default_factory=list,
+        title="Non-retryable error types",
+    )
+    """Error type names that should not be retried."""
+
+    @field_validator("initial_interval", "maximum_interval", mode="before")
+    @classmethod
+    def parse_duration(cls, v: str | timedelta | None) -> timedelta | None:
+        """Parse string duration to timedelta."""
+        if v is None or isinstance(v, timedelta):
+            return v
+        return parse_time_period(v)
+
+
+class TemporalDurableConfig(BaseDurableExecutionConfig):
+    """Configuration for Temporal durable execution.
+
+    Wraps agents in TemporalAgent for workflow-based execution with
+    automatic retries and state persistence.
+    """
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "x-icon": "simple-icons:temporal",
+            "x-doc-title": "Temporal Durable Execution",
+            "title": "Temporal Durable Execution Configuration",
+        }
+    )
+
+    type: Literal["temporal"] = Field("temporal", init=False)
+    """Temporal durable execution provider."""
+
+    activity_config: TemporalActivityConfig = Field(
+        default_factory=TemporalActivityConfig,
+        title="Base activity configuration",
+    )
+    """Base Temporal activity config for all activities."""
+
+    model_activity_config: TemporalActivityConfig | None = Field(
+        default=None,
+        title="Model activity configuration",
+    )
+    """Activity config for model request activities (merged with base)."""
+
+    retry_policy: TemporalRetryPolicy | None = Field(
+        default=None,
+        title="Retry policy",
+    )
+    """Retry policy for activities."""
+
+    toolset_activity_config: dict[str, TemporalActivityConfig] | None = Field(
+        default=None,
+        title="Toolset activity configurations",
+    )
+    """Per-toolset activity configs keyed by toolset ID."""
+
+    tool_activity_config: dict[str, dict[str, TemporalActivityConfig | bool]] | None = Field(
+        default=None,
+        title="Tool activity configurations",
+    )
+    """Per-tool activity configs: {toolset_id: {tool_name: config | False}}."""
+
+
+class PrefectTaskConfig(Schema):
+    """Prefect task configuration options.
+
+    Maps to Prefect task decorator options.
+    """
+
+    retries: int = Field(
+        default=0,
+        ge=0,
+        title="Maximum retries",
+    )
+    """Maximum number of retries for the task."""
+
+    retry_delay_seconds: float | list[float] = Field(
+        default=1.0,
+        title="Retry delay",
+    )
+    """Delay between retries (single value or list for custom backoff)."""
+
+    timeout_seconds: float | None = Field(
+        default=None,
+        title="Task timeout",
+    )
+    """Maximum time in seconds for task completion."""
+
+    persist_result: bool = Field(
+        default=True,
+        title="Persist result",
+    )
+    """Whether to persist task results."""
+
+    log_prints: bool = Field(
+        default=False,
+        title="Log prints",
+    )
+    """Whether to capture print statements in task logs."""
+
+
+class PrefectDurableConfig(BaseDurableExecutionConfig):
+    """Configuration for Prefect durable execution.
+
+    Wraps agents in PrefectAgent for flow-based execution with
+    automatic task tracking and observability.
+    """
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "x-icon": "simple-icons:prefect",
+            "x-doc-title": "Prefect Durable Execution",
+        }
+    )
+
+    type: Literal["prefect"] = Field("prefect", init=False)
+    """Prefect durable execution provider."""
+
+    model_task_config: PrefectTaskConfig | None = Field(
+        default=None,
+        title="Model task configuration",
+    )
+    """Task config for model request tasks."""
+
+    mcp_task_config: PrefectTaskConfig | None = Field(
+        default=None,
+        title="MCP task configuration",
+    )
+    """Task config for MCP server tasks."""
+
+    tool_task_config: PrefectTaskConfig | None = Field(
+        default=None,
+        title="Default tool task configuration",
+    )
+    """Default task config for tool calls."""
+
+    tool_task_config_by_name: dict[str, PrefectTaskConfig | None] | None = Field(
+        default=None,
+        title="Per-tool task configurations",
+    )
+    """Per-tool task configs (None value disables task wrapping)."""
+
+    event_stream_handler_task_config: PrefectTaskConfig | None = Field(
+        default=None,
+        title="Event handler task configuration",
+    )
+    """Task config for event stream handler."""
+
+
+class DBOSStepConfig(Schema):
+    """DBOS step configuration options.
+
+    Maps to DBOS step decorator options.
+    """
+
+    retries_allowed: bool = Field(
+        default=True,
+        title="Retries allowed",
+    )
+    """Whether retries are allowed for this step."""
+
+    interval_seconds: float = Field(
+        default=1.0,
+        gt=0,
+        title="Retry interval",
+    )
+    """Base interval between retries in seconds."""
+
+    max_attempts: int = Field(
+        default=3,
+        ge=1,
+        title="Maximum attempts",
+    )
+    """Maximum number of attempts."""
+
+    backoff_rate: float = Field(
+        default=2.0,
+        ge=1.0,
+        title="Backoff rate",
+    )
+    """Multiplier for exponential backoff."""
+
+
+class DBOSDurableConfig(BaseDurableExecutionConfig):
+    """Configuration for DBOS durable execution.
+
+    Wraps agents in DBOSAgent for step-based execution with
+    automatic checkpointing and recovery.
+    """
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "x-icon": "octicon:database-16",
+            "x-doc-title": "DBOS Durable Execution",
+        }
+    )
+
+    type: Literal["dbos"] = Field("dbos", init=False)
+    """DBOS durable execution provider."""
+
+    model_step_config: DBOSStepConfig | None = Field(
+        default=None,
+        title="Model step configuration",
+    )
+    """Step config for model request steps."""
+
+    mcp_step_config: DBOSStepConfig | None = Field(
+        default=None,
+        title="MCP step configuration",
+    )
+    """Step config for MCP server steps."""
+
+
+DurableExecutionConfig = Annotated[
+    TemporalDurableConfig | PrefectDurableConfig | DBOSDurableConfig,
+    Field(discriminator="type"),
+]