polos-sdk 0.1.0__py3-none-any.whl

polos/middleware/hook.py

@@ -0,0 +1,164 @@
+"""Hook decorator for lifecycle hooks.
+
+Hooks are callables that can intercept and modify execution at various lifecycle points.
+They are executed within a workflow execution context and support durable execution.
+
+Hooks have a specific signature: (ctx: WorkflowContext, hook_context: HookContext) -> HookResult
+"""
+
+import inspect
+from collections.abc import Callable
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+from ..types.types import AgentConfig, Step
+
+
+class HookAction(Enum):
+    """Action a hook can take after execution."""
+
+    CONTINUE = "continue"
+    FAIL = "fail"
+
+
+class HookContext(BaseModel):
+    """Context available to hooks.
+
+    This context is passed to hooks and contains information about
+    the current execution state.
+    """
+
+    # Immutable identifiers
+    workflow_id: str
+    session_id: str | None = None
+    user_id: str | None = None
+    agent_config: AgentConfig | None = None  # Not available to agent's on_start and on_end hooks
+
+    # Current state
+    steps: list[Step] = []  # All previous steps
+
+    # For workflow/tool hooks
+    current_tool: str | None = None
+    current_payload: dict[str, Any] | BaseModel | None = None
+    current_output: dict[str, Any] | BaseModel | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert hook context to dictionary for serialization."""
+        return self.model_dump(mode="json")
+
+    @classmethod
+    def from_dict(cls, data: Any) -> "HookContext":
+        """Create HookContext from dictionary."""
+        if isinstance(data, HookContext):
+            return data
+        if isinstance(data, dict):
+            return cls.model_validate(data)
+        raise TypeError(f"Cannot create HookContext from {type(data)}")
+
+
+class HookResult(BaseModel):
+    """Result from a hook execution.
+
+    Hooks return this to indicate what action to take and any modifications
+    to apply to the execution state.
+    """
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    action: HookAction = HookAction.CONTINUE
+
+    # Optional modifications
+    modified_payload: dict[str, Any] | None = None
+    modified_output: Any | None = None
+
+    # For FAIL action
+    error_message: str | None = None
+
+    @classmethod
+    def continue_with(cls, **modifications) -> "HookResult":
+        """Continue with optional modifications.
+
+        Args:
+            **modifications: Can include modified_payload, modified_output
+
+        Returns:
+            HookResult with CONTINUE action and modifications
+        """
+        return cls(action=HookAction.CONTINUE, **modifications)
+
+    @classmethod
+    def fail(cls, message: str) -> "HookResult":
+        """Fail processing with an error message.
+
+        Args:
+            message: Error message to return
+
+        Returns:
+            HookResult with FAIL action
+        """
+        return cls(action=HookAction.FAIL, error_message=message)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert hook result to dictionary for serialization."""
+        return self.model_dump(mode="json")
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "HookResult":
+        """Create HookResult from dictionary."""
+        return cls.model_validate(data)
+
+
+def _validate_hook_signature(func: Callable) -> None:
+    """Validate that hook function has correct signature.
+
+    Expected: (ctx: WorkflowContext, hook_context: HookContext) -> HookResult
+
+    Raises:
+        TypeError: If signature is invalid
+    """
+    sig = inspect.signature(func)
+    params = list(sig.parameters.values())
+
+    # Must have exactly 2 parameters: ctx and hook_context
+    if len(params) != 2:
+        raise TypeError(
+            f"Hook function '{func.__name__}' must have exactly 2 parameters: "
+            f"(ctx: WorkflowContext, hook_context: HookContext). Got {len(params)} parameters."
+        )
+
+
+def hook(func: Callable | None = None):
+    """
+    Decorator to mark a function as a hook.
+
+    Hook functions must have the signature:
+        (ctx: WorkflowContext, hook_context: HookContext) -> HookResult
+
+    Usage:
+        @hook
+        def my_hook(ctx: WorkflowContext, hook_context: HookContext) -> HookResult:
+            return HookResult.continue_with()
+
+    Args:
+        func: The function to decorate (when used as @hook)
+
+    Returns:
+        The function itself (validated)
+
+    Raises:
+        TypeError: If function signature is invalid
+    """
+
+    def decorator(f: Callable) -> Callable:
+        # Validate function signature
+        _validate_hook_signature(f)
+        return f
+
+    # Handle @hook (without parentheses) - the function is passed as the first argument
+    if func is not None:
+        return decorator(func)
+
+    # Handle @hook() - return decorator
+    return decorator

polos/middleware/hook_executor.py

@@ -0,0 +1,104 @@
+"""Hook execution infrastructure for lifecycle hooks.
+
+Hooks are executed within a workflow execution context and support durable execution.
+"""
+
+from collections.abc import Callable
+
+from ..core.context import WorkflowContext
+from ..core.workflow import _execution_context
+from .hook import HookAction, HookContext, HookResult
+
+
+def _get_function_identifier(func: Callable, index: int) -> str:
+    """Get a unique identifier for a function call.
+
+    Uses function name if available, otherwise falls back to index.
+    """
+    if hasattr(func, "__name__") and func.__name__ != "<lambda>":
+        return func.__name__
+    return f"hook_{index}"
+
+
+async def execute_hooks(
+    hook_name: str,
+    hooks: list[Callable],
+    hook_context: HookContext,
+    ctx: WorkflowContext,
+) -> HookResult:
+    """
+    Execute a list of hooks sequentially and return the combined result.
+
+    Hooks are executed within a workflow execution context. Each hook execution:
+    1. Checks for cached result (for durable execution)
+    2. If cached, returns cached result
+    3. If not cached, executes hook and stores result
+
+    Each hook can:
+    - Return CONTINUE to proceed to the next hook
+    - Return STOP to stop execution and return a value
+    - Return ERROR to stop execution with an error
+
+    Modifications from hooks are accumulated and applied in order.
+
+    Args:
+        hooks: List of hook callables (functions decorated with @hook)
+        hook_context: Context to pass to hooks
+        ctx: WorkflowContext for the current execution
+
+    Returns:
+        HookResult with action and any modifications
+
+    Raises:
+        ValueError: If not executed within a workflow execution context
+    """
+    if not hooks:
+        return HookResult.continue_with()
+
+    # Check we're in a workflow execution context
+    exec_context = _execution_context.get()
+    if not exec_context or not exec_context.get("execution_id"):
+        raise ValueError("Hooks must be executed within a workflow or agent")
+
+    # Accumulated modifications
+    modified_payload = hook_context.current_payload.copy() if hook_context.current_payload else {}
+    modified_output = hook_context.current_output.copy() if hook_context.current_output else {}
+
+    # Execute hooks sequentially
+    for index, hook_func in enumerate(hooks):
+        # Get function identifier for durable execution
+        func_id = _get_function_identifier(hook_func, index)
+        hook_result = await ctx.step.run(
+            f"{hook_name}.{func_id}.{index}", hook_func, ctx, hook_context
+        )
+
+        # Ensure result is HookResult
+        if not isinstance(hook_result, HookResult):
+            hook_result = HookResult.fail(
+                f"Hook '{func_id}' returned invalid result type: "
+                f"{type(hook_result)}. Expected HookResult."
+            )
+
+        # Apply modifications
+        if hook_result.modified_payload is not None:
+            modified_payload.update(hook_result.modified_payload)
+
+        if hook_result.modified_output is not None:
+            modified_output.update(hook_result.modified_output)
+
+        # Update hook_context with accumulated modifications for next hook
+        hook_context.current_payload = modified_payload
+        hook_context.current_output = modified_output
+
+        # Check action
+        if hook_result.action == HookAction.FAIL:
+            # Fail execution - return error
+            return hook_result
+
+        # CONTINUE - proceed to next hook
+
+    # All hooks completed with CONTINUE - return accumulated modifications
+    return HookResult.continue_with(
+        modified_payload=modified_payload,
+        modified_output=modified_output,
+    )

polos/runtime/batch.py

@@ -0,0 +1,87 @@
+"""Batch workflow triggering utilities."""
+
+from typing import Any
+
+from ..agents.agent import AgentRunConfig
+from ..core.workflow import _execution_context
+from ..types.types import BatchWorkflowInput
+from .client import ExecutionHandle, PolosClient
+
+
+async def batch_invoke(
+    client: PolosClient,
+    workflows: list[BatchWorkflowInput],
+    session_id: str | None = None,
+    user_id: str | None = None,
+) -> list[ExecutionHandle]:
+    """Invoke multiple different workflows in a single batch and return handles immediately.
+
+    This function cannot be called from within a workflow or agent.
+    Use step.batch_invoke() to call workflows from within workflows.
+
+    Args:
+        client: PolosClient instance
+        workflows: List of BatchWorkflowInput objects with 'id' (workflow_id string)
+            and 'payload' (dict or Pydantic model)
+        session_id: Optional session ID
+        user_id: Optional user ID
+
+    Returns:
+        List of ExecutionHandle objects for the submitted workflows
+
+    Example:
+        handles = await batch_invoke([
+            BatchWorkflowInput(id="workflow-1", payload={"foo": "bar"}),
+            BatchWorkflowInput(id="workflow-2", payload={"baz": 42}),
+        ])
+    """
+    # Check if we're in an execution context - fail if we are
+    if _execution_context.get() is not None:
+        raise RuntimeError(
+            "batch_invoke() cannot be called from within a workflow or agent. "
+            "Use step.batch_invoke() to call workflows from within workflows."
+        )
+
+    return await client.batch_invoke(workflows, session_id=session_id, user_id=user_id)
+
+
+async def batch_agent_invoke(
+    client: PolosClient,
+    agents: list[AgentRunConfig],
+) -> list[ExecutionHandle]:
+    """
+    Invoke multiple agents in parallel and return execution handles.
+
+    This helper is intended for use with Agent.with_input(), which returns
+    AgentRunConfig instances.
+
+    Args:
+        client: PolosClient instance
+        agents: List of AgentRunConfig instances
+
+    Example:
+        handles = await batch_agent_invoke([
+            grammar_agent.with_input("Check this"),
+            tone_agent.with_input("Check this too"),
+        ])
+    """
+    workflows: list[BatchWorkflowInput] = []
+    for config in agents:
+        payload: dict[str, Any] = {
+            "input": config.input,
+            "streaming": config.streaming,
+            "session_id": config.session_id,
+            "conversation_id": config.conversation_id,
+            "user_id": config.user_id,
+            **config.kwargs,
+        }
+        workflows.append(
+            BatchWorkflowInput(
+                id=config.agent.id,
+                payload=payload,
+                initial_state=config.initial_state,
+                run_timeout_seconds=config.run_timeout_seconds,
+            )
+        )
+
+    return await batch_invoke(client, workflows)