agentrun-sdk 0.1.2__py3-none-any.whl

agentrun_sdk/hooks/registry.py

@@ -0,0 +1,247 @@
+"""Hook registry system for managing event callbacks in the Strands Agent SDK.
+
+This module provides the core infrastructure for the typed hook system, enabling
+composable extension of agent functionality through strongly-typed event callbacks.
+The registry manages the mapping between event types and their associated callback
+functions, supporting both individual callback registration and bulk registration
+via hook provider objects.
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generator, Generic, Protocol, Type, TypeVar
+
+if TYPE_CHECKING:
+    from ..agent import Agent
+
+
+@dataclass
+class HookEvent:
+    """Base class for all hook events.
+
+    Attributes:
+        agent: The agent instance that triggered this event.
+    """
+
+    agent: "Agent"
+
+    @property
+    def should_reverse_callbacks(self) -> bool:
+        """Determine if callbacks for this event should be invoked in reverse order.
+
+        Returns:
+            False by default. Override to return True for events that should
+            invoke callbacks in reverse order (e.g., cleanup/teardown events).
+        """
+        return False
+
+    def _can_write(self, name: str) -> bool:
+        """Check if the given property can be written to.
+
+        Args:
+            name: The name of the property to check.
+
+        Returns:
+            True if the property can be written to, False otherwise.
+        """
+        return False
+
+    def __post_init__(self) -> None:
+        """Disallow writes to non-approved properties."""
+        # This is needed as otherwise the class can't be initialized at all, so we trigger
+        # this after class initialization
+        super().__setattr__("_disallow_writes", True)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Prevent setting attributes on hook events.
+
+        Raises:
+            AttributeError: Always raised to prevent setting attributes on hook events.
+        """
+        #  Allow setting attributes:
+        #    - during init (when __dict__) doesn't exist
+        #    - if the subclass specifically said the property is writable
+        if not hasattr(self, "_disallow_writes") or self._can_write(name):
+            return super().__setattr__(name, value)
+
+        raise AttributeError(f"Property {name} is not writable")
+
+
+TEvent = TypeVar("TEvent", bound=HookEvent, contravariant=True)
+"""Generic for adding callback handlers - contravariant to allow adding handlers which take in base classes."""
+
+TInvokeEvent = TypeVar("TInvokeEvent", bound=HookEvent)
+"""Generic for invoking events - non-contravariant to enable returning events."""
+
+
+class HookProvider(Protocol):
+    """Protocol for objects that provide hook callbacks to an agent.
+
+    Hook providers offer a composable way to extend agent functionality by
+    subscribing to various events in the agent lifecycle. This protocol enables
+    building reusable components that can hook into agent events.
+
+    Example:
+        ```python
+        class MyHookProvider(HookProvider):
+            def register_hooks(self, registry: HookRegistry) -> None:
+                registry.add_callback(StartRequestEvent, self.on_request_start)
+                registry.add_callback(EndRequestEvent, self.on_request_end)
+
+        agent = Agent(hooks=[MyHookProvider()])
+        ```
+    """
+
+    def register_hooks(self, registry: "HookRegistry", **kwargs: Any) -> None:
+        """Register callback functions for specific event types.
+
+        Args:
+            registry: The hook registry to register callbacks with.
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        ...
+
+
+class HookCallback(Protocol, Generic[TEvent]):
+    """Protocol for callback functions that handle hook events.
+
+    Hook callbacks are functions that receive a single strongly-typed event
+    argument and perform some action in response. They should not return
+    values and any exceptions they raise will propagate to the caller.
+
+    Example:
+        ```python
+        def my_callback(event: StartRequestEvent) -> None:
+            print(f"Request started for agent: {event.agent.name}")
+        ```
+    """
+
+    def __call__(self, event: TEvent) -> None:
+        """Handle a hook event.
+
+        Args:
+            event: The strongly-typed event to handle.
+        """
+        ...
+
+
+class HookRegistry:
+    """Registry for managing hook callbacks associated with event types.
+
+    The HookRegistry maintains a mapping of event types to callback functions
+    and provides methods for registering callbacks and invoking them when
+    events occur.
+
+    The registry handles callback ordering, including reverse ordering for
+    cleanup events, and provides type-safe event dispatching.
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty hook registry."""
+        self._registered_callbacks: dict[Type, list[HookCallback]] = {}
+
+    def add_callback(self, event_type: Type[TEvent], callback: HookCallback[TEvent]) -> None:
+        """Register a callback function for a specific event type.
+
+        Args:
+            event_type: The class type of events this callback should handle.
+            callback: The callback function to invoke when events of this type occur.
+
+        Example:
+            ```python
+            def my_handler(event: StartRequestEvent):
+                print("Request started")
+
+            registry.add_callback(StartRequestEvent, my_handler)
+            ```
+        """
+        callbacks = self._registered_callbacks.setdefault(event_type, [])
+        callbacks.append(callback)
+
+    def add_hook(self, hook: HookProvider) -> None:
+        """Register all callbacks from a hook provider.
+
+        This method allows bulk registration of callbacks by delegating to
+        the hook provider's register_hooks method. This is the preferred
+        way to register multiple related callbacks.
+
+        Args:
+            hook: The hook provider containing callbacks to register.
+
+        Example:
+            ```python
+            class MyHooks(HookProvider):
+                def register_hooks(self, registry: HookRegistry):
+                    registry.add_callback(StartRequestEvent, self.on_start)
+                    registry.add_callback(EndRequestEvent, self.on_end)
+
+            registry.add_hook(MyHooks())
+            ```
+        """
+        hook.register_hooks(self)
+
+    def invoke_callbacks(self, event: TInvokeEvent) -> TInvokeEvent:
+        """Invoke all registered callbacks for the given event.
+
+        This method finds all callbacks registered for the event's type and
+        invokes them in the appropriate order. For events with should_reverse_callbacks=True,
+        callbacks are invoked in reverse registration order. Any exceptions raised by callback
+        functions will propagate to the caller.
+
+        Args:
+            event: The event to dispatch to registered callbacks.
+
+        Returns:
+            The event dispatched to registered callbacks.
+
+        Example:
+            ```python
+            event = StartRequestEvent(agent=my_agent)
+            registry.invoke_callbacks(event)
+            ```
+        """
+        for callback in self.get_callbacks_for(event):
+            callback(event)
+
+        return event
+
+    def has_callbacks(self) -> bool:
+        """Check if the registry has any registered callbacks.
+
+        Returns:
+            True if there are any registered callbacks, False otherwise.
+
+        Example:
+            ```python
+            if registry.has_callbacks():
+                print("Registry has callbacks registered")
+            ```
+        """
+        return bool(self._registered_callbacks)
+
+    def get_callbacks_for(self, event: TEvent) -> Generator[HookCallback[TEvent], None, None]:
+        """Get callbacks registered for the given event in the appropriate order.
+
+        This method returns callbacks in registration order for normal events,
+        or reverse registration order for events that have should_reverse_callbacks=True.
+        This enables proper cleanup ordering for teardown events.
+
+        Args:
+            event: The event to get callbacks for.
+
+        Yields:
+            Callback functions registered for this event type, in the appropriate order.
+
+        Example:
+            ```python
+            event = EndRequestEvent(agent=my_agent)
+            for callback in registry.get_callbacks_for(event):
+                callback(event)
+            ```
+        """
+        event_type = type(event)
+
+        callbacks = self._registered_callbacks.get(event_type, [])
+        if event.should_reverse_callbacks:
+            yield from reversed(callbacks)
+        else:
+            yield from callbacks

agentrun_sdk/models/__init__.py

@@ -0,0 +1,10 @@
+"""SDK model providers.
+
+This package includes an abstract base Model class along with concrete implementations for specific providers.
+"""
+
+from . import bedrock, model
+from .bedrock import BedrockModel
+from .model import Model
+
+__all__ = ["bedrock", "model", "BedrockModel", "Model"]

agentrun_sdk/models/anthropic.py

@@ -0,0 +1,432 @@
+"""Anthropic Claude model provider.
+
+- Docs: https://docs.anthropic.com/claude/reference/getting-started-with-the-api
+"""
+
+import base64
+import json
+import logging
+import mimetypes
+from typing import Any, AsyncGenerator, Optional, Type, TypedDict, TypeVar, Union, cast
+
+import anthropic
+from pydantic import BaseModel
+from typing_extensions import Required, Unpack, override
+
+from ..event_loop.streaming import process_stream
+from ..tools import convert_pydantic_to_tool_spec
+from ..types.content import ContentBlock, Messages
+from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
+from ..types.streaming import StreamEvent
+from ..types.tools import ToolSpec
+from .model import Model
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class AnthropicModel(Model):
+    """Anthropic model provider implementation."""
+
+    EVENT_TYPES = {
+        "message_start",
+        "content_block_start",
+        "content_block_delta",
+        "content_block_stop",
+        "message_stop",
+    }
+
+    OVERFLOW_MESSAGES = {
+        "input is too long",
+        "input length exceeds context window",
+        "input and output tokens exceed your context limit",
+    }
+
+    class AnthropicConfig(TypedDict, total=False):
+        """Configuration options for Anthropic models.
+
+        Attributes:
+            max_tokens: Maximum number of tokens to generate.
+            model_id: Calude model ID (e.g., "claude-3-7-sonnet-latest").
+                For a complete list of supported models, see
+                https://docs.anthropic.com/en/docs/about-claude/models/all-models.
+            params: Additional model parameters (e.g., temperature).
+                For a complete list of supported parameters, see https://docs.anthropic.com/en/api/messages.
+        """
+
+        max_tokens: Required[int]
+        model_id: Required[str]
+        params: Optional[dict[str, Any]]
+
+    def __init__(self, *, client_args: Optional[dict[str, Any]] = None, **model_config: Unpack[AnthropicConfig]):
+        """Initialize provider instance.
+
+        Args:
+            client_args: Arguments for the underlying Anthropic client (e.g., api_key).
+                For a complete list of supported arguments, see https://docs.anthropic.com/en/api/client-sdks.
+            **model_config: Configuration options for the Anthropic model.
+        """
+        self.config = AnthropicModel.AnthropicConfig(**model_config)
+
+        logger.debug("config=<%s> | initializing", self.config)
+
+        client_args = client_args or {}
+        self.client = anthropic.AsyncAnthropic(**client_args)
+
+    @override
+    def update_config(self, **model_config: Unpack[AnthropicConfig]) -> None:  # type: ignore[override]
+        """Update the Anthropic model configuration with the provided arguments.
+
+        Args:
+            **model_config: Configuration overrides.
+        """
+        self.config.update(model_config)
+
+    @override
+    def get_config(self) -> AnthropicConfig:
+        """Get the Anthropic model configuration.
+
+        Returns:
+            The Anthropic model configuration.
+        """
+        return self.config
+
+    def _format_request_message_content(self, content: ContentBlock) -> dict[str, Any]:
+        """Format an Anthropic content block.
+
+        Args:
+            content: Message content.
+
+        Returns:
+            Anthropic formatted content block.
+
+        Raises:
+            TypeError: If the content block type cannot be converted to an Anthropic-compatible format.
+        """
+        if "document" in content:
+            mime_type = mimetypes.types_map.get(f".{content['document']['format']}", "application/octet-stream")
+            return {
+                "source": {
+                    "data": (
+                        content["document"]["source"]["bytes"].decode("utf-8")
+                        if mime_type == "text/plain"
+                        else base64.b64encode(content["document"]["source"]["bytes"]).decode("utf-8")
+                    ),
+                    "media_type": mime_type,
+                    "type": "text" if mime_type == "text/plain" else "base64",
+                },
+                "title": content["document"]["name"],
+                "type": "document",
+            }
+
+        if "image" in content:
+            return {
+                "source": {
+                    "data": base64.b64encode(content["image"]["source"]["bytes"]).decode("utf-8"),
+                    "media_type": mimetypes.types_map.get(f".{content['image']['format']}", "application/octet-stream"),
+                    "type": "base64",
+                },
+                "type": "image",
+            }
+
+        if "reasoningContent" in content:
+            return {
+                "signature": content["reasoningContent"]["reasoningText"]["signature"],
+                "thinking": content["reasoningContent"]["reasoningText"]["text"],
+                "type": "thinking",
+            }
+
+        if "text" in content:
+            return {"text": content["text"], "type": "text"}
+
+        if "toolUse" in content:
+            return {
+                "id": content["toolUse"]["toolUseId"],
+                "input": content["toolUse"]["input"],
+                "name": content["toolUse"]["name"],
+                "type": "tool_use",
+            }
+
+        if "toolResult" in content:
+            return {
+                "content": [
+                    self._format_request_message_content(
+                        {"text": json.dumps(tool_result_content["json"])}
+                        if "json" in tool_result_content
+                        else cast(ContentBlock, tool_result_content)
+                    )
+                    for tool_result_content in content["toolResult"]["content"]
+                ],
+                "is_error": content["toolResult"]["status"] == "error",
+                "tool_use_id": content["toolResult"]["toolUseId"],
+                "type": "tool_result",
+            }
+
+        raise TypeError(f"content_type=<{next(iter(content))}> | unsupported type")
+
+    def _format_request_messages(self, messages: Messages) -> list[dict[str, Any]]:
+        """Format an Anthropic messages array.
+
+        Args:
+            messages: List of message objects to be processed by the model.
+
+        Returns:
+            An Anthropic messages array.
+        """
+        formatted_messages = []
+
+        for message in messages:
+            formatted_contents: list[dict[str, Any]] = []
+
+            for content in message["content"]:
+                if "cachePoint" in content:
+                    formatted_contents[-1]["cache_control"] = {"type": "ephemeral"}
+                    continue
+
+                formatted_contents.append(self._format_request_message_content(content))
+
+            if formatted_contents:
+                formatted_messages.append({"content": formatted_contents, "role": message["role"]})
+
+        return formatted_messages
+
+    def format_request(
+        self, messages: Messages, tool_specs: Optional[list[ToolSpec]] = None, system_prompt: Optional[str] = None
+    ) -> dict[str, Any]:
+        """Format an Anthropic streaming request.
+
+        Args:
+            messages: List of message objects to be processed by the model.
+            tool_specs: List of tool specifications to make available to the model.
+            system_prompt: System prompt to provide context to the model.
+
+        Returns:
+            An Anthropic streaming request.
+
+        Raises:
+            TypeError: If a message contains a content block type that cannot be converted to an Anthropic-compatible
+                format.
+        """
+        return {
+            "max_tokens": self.config["max_tokens"],
+            "messages": self._format_request_messages(messages),
+            "model": self.config["model_id"],
+            "tools": [
+                {
+                    "name": tool_spec["name"],
+                    "description": tool_spec["description"],
+                    "input_schema": tool_spec["inputSchema"]["json"],
+                }
+                for tool_spec in tool_specs or []
+            ],
+            **({"system": system_prompt} if system_prompt else {}),
+            **(self.config.get("params") or {}),
+        }
+
+    def format_chunk(self, event: dict[str, Any]) -> StreamEvent:
+        """Format the Anthropic response events into standardized message chunks.
+
+        Args:
+            event: A response event from the Anthropic model.
+
+        Returns:
+            The formatted chunk.
+
+        Raises:
+            RuntimeError: If chunk_type is not recognized.
+                This error should never be encountered as we control chunk_type in the stream method.
+        """
+        match event["type"]:
+            case "message_start":
+                return {"messageStart": {"role": "assistant"}}
+
+            case "content_block_start":
+                content = event["content_block"]
+
+                if content["type"] == "tool_use":
+                    return {
+                        "contentBlockStart": {
+                            "contentBlockIndex": event["index"],
+                            "start": {
+                                "toolUse": {
+                                    "name": content["name"],
+                                    "toolUseId": content["id"],
+                                }
+                            },
+                        }
+                    }
+
+                return {"contentBlockStart": {"contentBlockIndex": event["index"], "start": {}}}
+
+            case "content_block_delta":
+                delta = event["delta"]
+
+                match delta["type"]:
+                    case "signature_delta":
+                        return {
+                            "contentBlockDelta": {
+                                "contentBlockIndex": event["index"],
+                                "delta": {
+                                    "reasoningContent": {
+                                        "signature": delta["signature"],
+                                    },
+                                },
+                            },
+                        }
+
+                    case "thinking_delta":
+                        return {
+                            "contentBlockDelta": {
+                                "contentBlockIndex": event["index"],
+                                "delta": {
+                                    "reasoningContent": {
+                                        "text": delta["thinking"],
+                                    },
+                                },
+                            },
+                        }
+
+                    case "input_json_delta":
+                        return {
+                            "contentBlockDelta": {
+                                "contentBlockIndex": event["index"],
+                                "delta": {
+                                    "toolUse": {
+                                        "input": delta["partial_json"],
+                                    },
+                                },
+                            },
+                        }
+
+                    case "text_delta":
+                        return {
+                            "contentBlockDelta": {
+                                "contentBlockIndex": event["index"],
+                                "delta": {
+                                    "text": delta["text"],
+                                },
+                            },
+                        }
+
+                    case _:
+                        raise RuntimeError(
+                            f"event_type=<content_block_delta>, delta_type=<{delta['type']}> | unknown type"
+                        )
+
+            case "content_block_stop":
+                return {"contentBlockStop": {"contentBlockIndex": event["index"]}}
+
+            case "message_stop":
+                message = event["message"]
+
+                return {"messageStop": {"stopReason": message["stop_reason"]}}
+
+            case "metadata":
+                usage = event["usage"]
+
+                return {
+                    "metadata": {
+                        "usage": {
+                            "inputTokens": usage["input_tokens"],
+                            "outputTokens": usage["output_tokens"],
+                            "totalTokens": usage["input_tokens"] + usage["output_tokens"],
+                        },
+                        "metrics": {
+                            "latencyMs": 0,  # TODO
+                        },
+                    }
+                }
+
+            case _:
+                raise RuntimeError(f"event_type=<{event['type']} | unknown type")
+
+    @override
+    async def stream(
+        self,
+        messages: Messages,
+        tool_specs: Optional[list[ToolSpec]] = None,
+        system_prompt: Optional[str] = None,
+        **kwargs: Any,
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """Stream conversation with the Anthropic model.
+
+        Args:
+            messages: List of message objects to be processed by the model.
+            tool_specs: List of tool specifications to make available to the model.
+            system_prompt: System prompt to provide context to the model.
+            **kwargs: Additional keyword arguments for future extensibility.
+
+        Yields:
+            Formatted message chunks from the model.
+
+        Raises:
+            ContextWindowOverflowException: If the input exceeds the model's context window.
+            ModelThrottledException: If the request is throttled by Anthropic.
+        """
+        logger.debug("formatting request")
+        request = self.format_request(messages, tool_specs, system_prompt)
+        logger.debug("request=<%s>", request)
+
+        logger.debug("invoking model")
+        try:
+            async with self.client.messages.stream(**request) as stream:
+                logger.debug("got response from model")
+                async for event in stream:
+                    if event.type in AnthropicModel.EVENT_TYPES:
+                        yield self.format_chunk(event.model_dump())
+
+                usage = event.message.usage  # type: ignore
+                yield self.format_chunk({"type": "metadata", "usage": usage.model_dump()})
+
+        except anthropic.RateLimitError as error:
+            raise ModelThrottledException(str(error)) from error
+
+        except anthropic.BadRequestError as error:
+            if any(overflow_message in str(error).lower() for overflow_message in AnthropicModel.OVERFLOW_MESSAGES):
+                raise ContextWindowOverflowException(str(error)) from error
+
+            raise error
+
+        logger.debug("finished streaming response from model")
+
+    @override
+    async def structured_output(
+        self, output_model: Type[T], prompt: Messages, system_prompt: Optional[str] = None, **kwargs: Any
+    ) -> AsyncGenerator[dict[str, Union[T, Any]], None]:
+        """Get structured output from the model.
+
+        Args:
+            output_model: The output model to use for the agent.
+            prompt: The prompt messages to use for the agent.
+            system_prompt: System prompt to provide context to the model.
+            **kwargs: Additional keyword arguments for future extensibility.
+
+        Yields:
+            Model events with the last being the structured output.
+        """
+        tool_spec = convert_pydantic_to_tool_spec(output_model)
+
+        response = self.stream(messages=prompt, tool_specs=[tool_spec], system_prompt=system_prompt, **kwargs)
+        async for event in process_stream(response):
+            yield event
+
+        stop_reason, messages, _, _ = event["stop"]
+
+        if stop_reason != "tool_use":
+            raise ValueError(f'Model returned stop_reason: {stop_reason} instead of "tool_use".')
+
+        content = messages["content"]
+        output_response: dict[str, Any] | None = None
+        for block in content:
+            # if the tool use name doesn't match the tool spec name, skip, and if the block is not a tool use, skip.
+            # if the tool use name never matches, raise an error.
+            if block.get("toolUse") and block["toolUse"]["name"] == tool_spec["name"]:
+                output_response = block["toolUse"]["input"]
+            else:
+                continue
+
+        if output_response is None:
+            raise ValueError("No valid tool use or tool use input was found in the Anthropic response.")
+
+        yield {"output": output_model(**output_response)}