agentrun-sdk 0.1.2__py3-none-any.whl

agentrun_sdk/event_loop/streaming.py

@@ -0,0 +1,319 @@
+"""Utilities for handling streaming responses from language models."""
+
+import json
+import logging
+from typing import Any, AsyncGenerator, AsyncIterable, Optional
+
+from ..models.model import Model
+from ..types.content import ContentBlock, Message, Messages
+from ..types.streaming import (
+    ContentBlockDeltaEvent,
+    ContentBlockStart,
+    ContentBlockStartEvent,
+    MessageStartEvent,
+    MessageStopEvent,
+    MetadataEvent,
+    Metrics,
+    RedactContentEvent,
+    StopReason,
+    StreamEvent,
+    Usage,
+)
+from ..types.tools import ToolSpec, ToolUse
+
+logger = logging.getLogger(__name__)
+
+
+def remove_blank_messages_content_text(messages: Messages) -> Messages:
+    """Remove or replace blank text in message content.
+
+    Args:
+        messages: Conversation messages to update.
+
+    Returns:
+        Updated messages.
+    """
+    removed_blank_message_content_text = False
+    replaced_blank_message_content_text = False
+
+    for message in messages:
+        # only modify assistant messages
+        if "role" in message and message["role"] != "assistant":
+            continue
+
+        if "content" in message:
+            content = message["content"]
+            has_tool_use = any("toolUse" in item for item in content)
+
+            if has_tool_use:
+                # Remove blank 'text' items for assistant messages
+                before_len = len(content)
+                content[:] = [item for item in content if "text" not in item or item["text"].strip()]
+                if not removed_blank_message_content_text and before_len != len(content):
+                    removed_blank_message_content_text = True
+            else:
+                # Replace blank 'text' with '[blank text]' for assistant messages
+                for item in content:
+                    if "text" in item and not item["text"].strip():
+                        replaced_blank_message_content_text = True
+                        item["text"] = "[blank text]"
+
+    if removed_blank_message_content_text:
+        logger.debug("removed blank message context text")
+    if replaced_blank_message_content_text:
+        logger.debug("replaced blank message context text")
+
+    return messages
+
+
+def handle_message_start(event: MessageStartEvent, message: Message) -> Message:
+    """Handles the start of a message by setting the role in the message dictionary.
+
+    Args:
+        event: A message start event.
+        message: The message dictionary being constructed.
+
+    Returns:
+        Updated message dictionary with the role set.
+    """
+    message["role"] = event["role"]
+    return message
+
+
+def handle_content_block_start(event: ContentBlockStartEvent) -> dict[str, Any]:
+    """Handles the start of a content block by extracting tool usage information if any.
+
+    Args:
+        event: Start event.
+
+    Returns:
+        Dictionary with tool use id and name if tool use request, empty dictionary otherwise.
+    """
+    start: ContentBlockStart = event["start"]
+    current_tool_use = {}
+
+    if "toolUse" in start and start["toolUse"]:
+        tool_use_data = start["toolUse"]
+        current_tool_use["toolUseId"] = tool_use_data["toolUseId"]
+        current_tool_use["name"] = tool_use_data["name"]
+        current_tool_use["input"] = ""
+
+    return current_tool_use
+
+
+def handle_content_block_delta(
+    event: ContentBlockDeltaEvent, state: dict[str, Any]
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Handles content block delta updates by appending text, tool input, or reasoning content to the state.
+
+    Args:
+        event: Delta event.
+        state: The current state of message processing.
+
+    Returns:
+        Updated state with appended text or tool input.
+    """
+    delta_content = event["delta"]
+
+    callback_event = {}
+
+    if "toolUse" in delta_content:
+        if "input" not in state["current_tool_use"]:
+            state["current_tool_use"]["input"] = ""
+
+        state["current_tool_use"]["input"] += delta_content["toolUse"]["input"]
+        callback_event["callback"] = {"delta": delta_content, "current_tool_use": state["current_tool_use"]}
+
+    elif "text" in delta_content:
+        state["text"] += delta_content["text"]
+        callback_event["callback"] = {"data": delta_content["text"], "delta": delta_content}
+
+    elif "reasoningContent" in delta_content:
+        if "text" in delta_content["reasoningContent"]:
+            if "reasoningText" not in state:
+                state["reasoningText"] = ""
+
+            state["reasoningText"] += delta_content["reasoningContent"]["text"]
+            callback_event["callback"] = {
+                "reasoningText": delta_content["reasoningContent"]["text"],
+                "delta": delta_content,
+                "reasoning": True,
+            }
+
+        elif "signature" in delta_content["reasoningContent"]:
+            if "signature" not in state:
+                state["signature"] = ""
+
+            state["signature"] += delta_content["reasoningContent"]["signature"]
+            callback_event["callback"] = {
+                "reasoning_signature": delta_content["reasoningContent"]["signature"],
+                "delta": delta_content,
+                "reasoning": True,
+            }
+
+    return state, callback_event
+
+
+def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]:
+    """Handles the end of a content block by finalizing tool usage, text content, or reasoning content.
+
+    Args:
+        state: The current state of message processing.
+
+    Returns:
+        Updated state with finalized content block.
+    """
+    content: list[ContentBlock] = state["content"]
+
+    current_tool_use = state["current_tool_use"]
+    text = state["text"]
+    reasoning_text = state["reasoningText"]
+
+    if current_tool_use:
+        if "input" not in current_tool_use:
+            current_tool_use["input"] = ""
+
+        try:
+            current_tool_use["input"] = json.loads(current_tool_use["input"])
+        except ValueError:
+            current_tool_use["input"] = {}
+
+        tool_use_id = current_tool_use["toolUseId"]
+        tool_use_name = current_tool_use["name"]
+
+        tool_use = ToolUse(
+            toolUseId=tool_use_id,
+            name=tool_use_name,
+            input=current_tool_use["input"],
+        )
+        content.append({"toolUse": tool_use})
+        state["current_tool_use"] = {}
+
+    elif text:
+        content.append({"text": text})
+        state["text"] = ""
+
+    elif reasoning_text:
+        content.append(
+            {
+                "reasoningContent": {
+                    "reasoningText": {
+                        "text": state["reasoningText"],
+                        "signature": state["signature"],
+                    }
+                }
+            }
+        )
+        state["reasoningText"] = ""
+
+    return state
+
+
+def handle_message_stop(event: MessageStopEvent) -> StopReason:
+    """Handles the end of a message by returning the stop reason.
+
+    Args:
+        event: Stop event.
+
+    Returns:
+        The reason for stopping the stream.
+    """
+    return event["stopReason"]
+
+
+def handle_redact_content(event: RedactContentEvent, state: dict[str, Any]) -> None:
+    """Handles redacting content from the input or output.
+
+    Args:
+        event: Redact Content Event.
+        state: The current state of message processing.
+    """
+    if event.get("redactAssistantContentMessage") is not None:
+        state["message"]["content"] = [{"text": event["redactAssistantContentMessage"]}]
+
+
+def extract_usage_metrics(event: MetadataEvent) -> tuple[Usage, Metrics]:
+    """Extracts usage metrics from the metadata chunk.
+
+    Args:
+        event: metadata.
+
+    Returns:
+        The extracted usage metrics and latency.
+    """
+    usage = Usage(**event["usage"])
+    metrics = Metrics(**event["metrics"])
+
+    return usage, metrics
+
+
+async def process_stream(chunks: AsyncIterable[StreamEvent]) -> AsyncGenerator[dict[str, Any], None]:
+    """Processes the response stream from the API, constructing the final message and extracting usage metrics.
+
+    Args:
+        chunks: The chunks of the response stream from the model.
+
+    Yields:
+        The reason for stopping, the constructed message, and the usage metrics.
+    """
+    stop_reason: StopReason = "end_turn"
+
+    state: dict[str, Any] = {
+        "message": {"role": "assistant", "content": []},
+        "text": "",
+        "current_tool_use": {},
+        "reasoningText": "",
+        "signature": "",
+    }
+    state["content"] = state["message"]["content"]
+
+    usage: Usage = Usage(inputTokens=0, outputTokens=0, totalTokens=0)
+    metrics: Metrics = Metrics(latencyMs=0)
+
+    async for chunk in chunks:
+        yield {"callback": {"event": chunk}}
+
+        if "messageStart" in chunk:
+            state["message"] = handle_message_start(chunk["messageStart"], state["message"])
+        elif "contentBlockStart" in chunk:
+            state["current_tool_use"] = handle_content_block_start(chunk["contentBlockStart"])
+        elif "contentBlockDelta" in chunk:
+            state, callback_event = handle_content_block_delta(chunk["contentBlockDelta"], state)
+            yield callback_event
+        elif "contentBlockStop" in chunk:
+            state = handle_content_block_stop(state)
+        elif "messageStop" in chunk:
+            stop_reason = handle_message_stop(chunk["messageStop"])
+        elif "metadata" in chunk:
+            usage, metrics = extract_usage_metrics(chunk["metadata"])
+        elif "redactContent" in chunk:
+            handle_redact_content(chunk["redactContent"], state)
+
+    yield {"stop": (stop_reason, state["message"], usage, metrics)}
+
+
+async def stream_messages(
+    model: Model,
+    system_prompt: Optional[str],
+    messages: Messages,
+    tool_specs: list[ToolSpec],
+) -> AsyncGenerator[dict[str, Any], None]:
+    """Streams messages to the model and processes the response.
+
+    Args:
+        model: Model provider.
+        system_prompt: The system prompt to send.
+        messages: List of messages to send.
+        tool_specs: The list of tool specs.
+
+    Yields:
+        The reason for stopping, the final message, and the usage metrics
+    """
+    logger.debug("model=<%s> | streaming messages", model)
+
+    messages = remove_blank_messages_content_text(messages)
+
+    chunks = model.stream(messages, tool_specs if tool_specs else None, system_prompt)
+
+    async for event in process_stream(chunks):
+        yield event

agentrun_sdk/experimental/__init__.py

@@ -0,0 +1,4 @@
+"""Experimental features.
+
+This module implements experimental features that are subject to change in future revisions without notice.
+"""

agentrun_sdk/experimental/hooks/__init__.py

@@ -0,0 +1,15 @@
+"""Experimental hook functionality that has not yet reached stability."""
+
+from .events import (
+    AfterModelInvocationEvent,
+    AfterToolInvocationEvent,
+    BeforeModelInvocationEvent,
+    BeforeToolInvocationEvent,
+)
+
+__all__ = [
+    "BeforeToolInvocationEvent",
+    "AfterToolInvocationEvent",
+    "BeforeModelInvocationEvent",
+    "AfterModelInvocationEvent",
+]

agentrun_sdk/experimental/hooks/events.py

@@ -0,0 +1,123 @@
+"""Experimental hook events emitted as part of invoking Agents.
+
+This module defines the events that are emitted as Agents run through the lifecycle of a request.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from ...hooks import HookEvent
+from ...types.content import Message
+from ...types.streaming import StopReason
+from ...types.tools import AgentTool, ToolResult, ToolUse
+
+
+@dataclass
+class BeforeToolInvocationEvent(HookEvent):
+    """Event triggered before a tool is invoked.
+
+    This event is fired just before the agent executes a tool, allowing hook
+    providers to inspect, modify, or replace the tool that will be executed.
+    The selected_tool can be modified by hook callbacks to change which tool
+    gets executed.
+
+    Attributes:
+        selected_tool: The tool that will be invoked. Can be modified by hooks
+            to change which tool gets executed. This may be None if tool lookup failed.
+        tool_use: The tool parameters that will be passed to selected_tool.
+        invocation_state: Keyword arguments that will be passed to the tool.
+    """
+
+    selected_tool: Optional[AgentTool]
+    tool_use: ToolUse
+    invocation_state: dict[str, Any]
+
+    def _can_write(self, name: str) -> bool:
+        return name in ["selected_tool", "tool_use"]
+
+
+@dataclass
+class AfterToolInvocationEvent(HookEvent):
+    """Event triggered after a tool invocation completes.
+
+    This event is fired after the agent has finished executing a tool,
+    regardless of whether the execution was successful or resulted in an error.
+    Hook providers can use this event for cleanup, logging, or post-processing.
+
+    Note: This event uses reverse callback ordering, meaning callbacks registered
+    later will be invoked first during cleanup.
+
+    Attributes:
+        selected_tool: The tool that was invoked. It may be None if tool lookup failed.
+        tool_use: The tool parameters that were passed to the tool invoked.
+        invocation_state: Keyword arguments that were passed to the tool
+        result: The result of the tool invocation. Either a ToolResult on success
+            or an Exception if the tool execution failed.
+    """
+
+    selected_tool: Optional[AgentTool]
+    tool_use: ToolUse
+    invocation_state: dict[str, Any]
+    result: ToolResult
+    exception: Optional[Exception] = None
+
+    def _can_write(self, name: str) -> bool:
+        return name == "result"
+
+    @property
+    def should_reverse_callbacks(self) -> bool:
+        """True to invoke callbacks in reverse order."""
+        return True
+
+
+@dataclass
+class BeforeModelInvocationEvent(HookEvent):
+    """Event triggered before the model is invoked.
+
+    This event is fired just before the agent calls the model for inference,
+    allowing hook providers to inspect or modify the messages and configuration
+    that will be sent to the model.
+
+    Note: This event is not fired for invocations to structured_output.
+    """
+
+    pass
+
+
+@dataclass
+class AfterModelInvocationEvent(HookEvent):
+    """Event triggered after the model invocation completes.
+
+    This event is fired after the agent has finished calling the model,
+    regardless of whether the invocation was successful or resulted in an error.
+    Hook providers can use this event for cleanup, logging, or post-processing.
+
+    Note: This event uses reverse callback ordering, meaning callbacks registered
+    later will be invoked first during cleanup.
+
+    Note: This event is not fired for invocations to structured_output.
+
+    Attributes:
+        stop_response: The model response data if invocation was successful, None if failed.
+        exception: Exception if the model invocation failed, None if successful.
+    """
+
+    @dataclass
+    class ModelStopResponse:
+        """Model response data from successful invocation.
+
+        Attributes:
+            stop_reason: The reason the model stopped generating.
+            message: The generated message from the model.
+        """
+
+        message: Message
+        stop_reason: StopReason
+
+    stop_response: Optional[ModelStopResponse] = None
+    exception: Optional[Exception] = None
+
+    @property
+    def should_reverse_callbacks(self) -> bool:
+        """True to invoke callbacks in reverse order."""
+        return True

agentrun_sdk/handlers/__init__.py

@@ -0,0 +1,10 @@
+"""Various handlers for performing custom actions on agent state.
+
+Examples include:
+
+- Displaying events from the event stream
+"""
+
+from .callback_handler import CompositeCallbackHandler, PrintingCallbackHandler, null_callback_handler
+
+__all__ = ["CompositeCallbackHandler", "null_callback_handler", "PrintingCallbackHandler"]

agentrun_sdk/handlers/callback_handler.py

@@ -0,0 +1,70 @@
+"""This module provides handlers for formatting and displaying events from the agent."""
+
+from collections.abc import Callable
+from typing import Any
+
+
+class PrintingCallbackHandler:
+    """Handler for streaming text output and tool invocations to stdout."""
+
+    def __init__(self) -> None:
+        """Initialize handler."""
+        self.tool_count = 0
+        self.previous_tool_use = None
+
+    def __call__(self, **kwargs: Any) -> None:
+        """Stream text output and tool invocations to stdout.
+
+        Args:
+            **kwargs: Callback event data including:
+                - reasoningText (Optional[str]): Reasoning text to print if provided.
+                - data (str): Text content to stream.
+                - complete (bool): Whether this is the final chunk of a response.
+                - current_tool_use (dict): Information about the current tool being used.
+        """
+        reasoningText = kwargs.get("reasoningText", False)
+        data = kwargs.get("data", "")
+        complete = kwargs.get("complete", False)
+        current_tool_use = kwargs.get("current_tool_use", {})
+
+        if reasoningText:
+            print(reasoningText, end="")
+
+        if data:
+            print(data, end="" if not complete else "\n")
+
+        if current_tool_use and current_tool_use.get("name"):
+            tool_name = current_tool_use.get("name", "Unknown tool")
+            if self.previous_tool_use != current_tool_use:
+                self.previous_tool_use = current_tool_use
+                self.tool_count += 1
+                print(f"\nTool #{self.tool_count}: {tool_name}")
+
+        if complete and data:
+            print("\n")
+
+
+class CompositeCallbackHandler:
+    """Class-based callback handler that combines multiple callback handlers.
+
+    This handler allows multiple callback handlers to be invoked for the same events,
+    enabling different processing or output formats for the same stream data.
+    """
+
+    def __init__(self, *handlers: Callable) -> None:
+        """Initialize handler."""
+        self.handlers = handlers
+
+    def __call__(self, **kwargs: Any) -> None:
+        """Invoke all handlers in the chain."""
+        for handler in self.handlers:
+            handler(**kwargs)
+
+
+def null_callback_handler(**_kwargs: Any) -> None:
+    """Callback handler that discards all output.
+
+    Args:
+        **_kwargs: Event data (ignored).
+    """
+    return None

agentrun_sdk/hooks/__init__.py

@@ -0,0 +1,49 @@
+"""Typed hook system for extending agent functionality.
+
+This module provides a composable mechanism for building objects that can hook
+into specific events during the agent lifecycle. The hook system enables both
+built-in SDK components and user code to react to or modify agent behavior
+through strongly-typed event callbacks.
+
+Example Usage:
+    ```python
+    from strands.hooks import HookProvider, HookRegistry
+    from strands.hooks.events import StartRequestEvent, EndRequestEvent
+
+    class LoggingHooks(HookProvider):
+        def register_hooks(self, registry: HookRegistry) -> None:
+            registry.add_callback(StartRequestEvent, self.log_start)
+            registry.add_callback(EndRequestEvent, self.log_end)
+
+        def log_start(self, event: StartRequestEvent) -> None:
+            print(f"Request started for {event.agent.name}")
+
+        def log_end(self, event: EndRequestEvent) -> None:
+            print(f"Request completed for {event.agent.name}")
+
+    # Use with agent
+    agent = Agent(hooks=[LoggingHooks()])
+    ```
+
+This replaces the older callback_handler approach with a more composable,
+type-safe system that supports multiple subscribers per event type.
+"""
+
+from .events import (
+    AfterInvocationEvent,
+    AgentInitializedEvent,
+    BeforeInvocationEvent,
+    MessageAddedEvent,
+)
+from .registry import HookCallback, HookEvent, HookProvider, HookRegistry
+
+__all__ = [
+    "AgentInitializedEvent",
+    "BeforeInvocationEvent",
+    "AfterInvocationEvent",
+    "MessageAddedEvent",
+    "HookEvent",
+    "HookProvider",
+    "HookCallback",
+    "HookRegistry",
+]

agentrun_sdk/hooks/events.py

@@ -0,0 +1,80 @@
+"""Hook events emitted as part of invoking Agents.
+
+This module defines the events that are emitted as Agents run through the lifecycle of a request.
+"""
+
+from dataclasses import dataclass
+
+from ..types.content import Message
+from .registry import HookEvent
+
+
+@dataclass
+class AgentInitializedEvent(HookEvent):
+    """Event triggered when an agent has finished initialization.
+
+    This event is fired after the agent has been fully constructed and all
+    built-in components have been initialized. Hook providers can use this
+    event to perform setup tasks that require a fully initialized agent.
+    """
+
+    pass
+
+
+@dataclass
+class BeforeInvocationEvent(HookEvent):
+    """Event triggered at the beginning of a new agent request.
+
+    This event is fired before the agent begins processing a new user request,
+    before any model inference or tool execution occurs. Hook providers can
+    use this event to perform request-level setup, logging, or validation.
+
+    This event is triggered at the beginning of the following api calls:
+      - Agent.__call__
+      - Agent.stream_async
+      - Agent.structured_output
+    """
+
+    pass
+
+
+@dataclass
+class AfterInvocationEvent(HookEvent):
+    """Event triggered at the end of an agent request.
+
+    This event is fired after the agent has completed processing a request,
+    regardless of whether it completed successfully or encountered an error.
+    Hook providers can use this event for cleanup, logging, or state persistence.
+
+    Note: This event uses reverse callback ordering, meaning callbacks registered
+    later will be invoked first during cleanup.
+
+    This event is triggered at the end of the following api calls:
+      - Agent.__call__
+      - Agent.stream_async
+      - Agent.structured_output
+    """
+
+    @property
+    def should_reverse_callbacks(self) -> bool:
+        """True to invoke callbacks in reverse order."""
+        return True
+
+
+@dataclass
+class MessageAddedEvent(HookEvent):
+    """Event triggered when a message is added to the agent's conversation.
+
+    This event is fired whenever the agent adds a new message to its internal
+    message history, including user messages, assistant responses, and tool
+    results. Hook providers can use this event for logging, monitoring, or
+    implementing custom message processing logic.
+
+    Note: This event is only triggered for messages added by the framework
+    itself, not for messages manually added by tools or external code.
+
+    Attributes:
+        message: The message that was added to the conversation history.
+    """
+
+    message: Message