agentrun-sdk 0.1.2__py3-none-any.whl

agentrun_sdk/models/writer.py

@@ -0,0 +1,449 @@
+"""Writer model provider.
+
+- Docs: https://dev.writer.com/home/introduction
+"""
+
+import base64
+import json
+import logging
+import mimetypes
+from typing import Any, AsyncGenerator, Dict, List, Optional, Type, TypedDict, TypeVar, Union, cast
+
+import writerai
+from pydantic import BaseModel
+from typing_extensions import Unpack, override
+
+from ..types.content import ContentBlock, Messages
+from ..types.exceptions import ModelThrottledException
+from ..types.streaming import StreamEvent
+from ..types.tools import ToolResult, ToolSpec, ToolUse
+from .model import Model
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class WriterModel(Model):
+    """Writer API model provider implementation."""
+
+    class WriterConfig(TypedDict, total=False):
+        """Configuration options for Writer API.
+
+        Attributes:
+            model_id: Model name to use (e.g. palmyra-x5, palmyra-x4, etc.).
+            max_tokens: Maximum number of tokens to generate.
+            stop: Default stop sequences.
+            stream_options: Additional options for streaming.
+            temperature: What sampling temperature to use.
+            top_p: Threshold for 'nucleus sampling'
+        """
+
+        model_id: str
+        max_tokens: Optional[int]
+        stop: Optional[Union[str, List[str]]]
+        stream_options: Dict[str, Any]
+        temperature: Optional[float]
+        top_p: Optional[float]
+
+    def __init__(self, client_args: Optional[dict[str, Any]] = None, **model_config: Unpack[WriterConfig]):
+        """Initialize provider instance.
+
+        Args:
+            client_args: Arguments for the Writer client (e.g., api_key, base_url, timeout, etc.).
+            **model_config: Configuration options for the Writer model.
+        """
+        self.config = WriterModel.WriterConfig(**model_config)
+
+        logger.debug("config=<%s> | initializing", self.config)
+
+        client_args = client_args or {}
+        self.client = writerai.AsyncClient(**client_args)
+
+    @override
+    def update_config(self, **model_config: Unpack[WriterConfig]) -> None:  # type: ignore[override]
+        """Update the Writer Model configuration with the provided arguments.
+
+        Args:
+            **model_config: Configuration overrides.
+        """
+        self.config.update(model_config)
+
+    @override
+    def get_config(self) -> WriterConfig:
+        """Get the Writer model configuration.
+
+        Returns:
+            The Writer model configuration.
+        """
+        return self.config
+
+    def _format_request_message_contents_vision(self, contents: list[ContentBlock]) -> list[dict[str, Any]]:
+        def _format_content_vision(content: ContentBlock) -> dict[str, Any]:
+            """Format a Writer content block for Palmyra V5 request.
+
+            - NOTE: "reasoningContent", "document" and "video" are not supported currently.
+
+            Args:
+                content: Message content.
+
+            Returns:
+                Writer formatted content block for models, which support vision content format.
+
+            Raises:
+                TypeError: If the content block type cannot be converted to a Writer-compatible format.
+            """
+            if "text" in content:
+                return {"text": content["text"], "type": "text"}
+
+            if "image" in content:
+                mime_type = mimetypes.types_map.get(f".{content['image']['format']}", "application/octet-stream")
+                image_data = base64.b64encode(content["image"]["source"]["bytes"]).decode("utf-8")
+
+                return {
+                    "image_url": {
+                        "url": f"data:{mime_type};base64,{image_data}",
+                    },
+                    "type": "image_url",
+                }
+
+            raise TypeError(f"content_type=<{next(iter(content))}> | unsupported type")
+
+        return [
+            _format_content_vision(content)
+            for content in contents
+            if not any(block_type in content for block_type in ["toolResult", "toolUse"])
+        ]
+
+    def _format_request_message_contents(self, contents: list[ContentBlock]) -> str:
+        def _format_content(content: ContentBlock) -> str:
+            """Format a Writer content block for Palmyra models (except V5) request.
+
+            - NOTE: "reasoningContent", "document", "video" and "image" are not supported currently.
+
+            Args:
+                content: Message content.
+
+            Returns:
+                Writer formatted content block.
+
+            Raises:
+                TypeError: If the content block type cannot be converted to a Writer-compatible format.
+            """
+            if "text" in content:
+                return content["text"]
+
+            raise TypeError(f"content_type=<{next(iter(content))}> | unsupported type")
+
+        content_blocks = list(
+            filter(
+                lambda content: content.get("text")
+                and not any(block_type in content for block_type in ["toolResult", "toolUse"]),
+                contents,
+            )
+        )
+
+        if len(content_blocks) > 1:
+            raise ValueError(
+                f"Model with name {self.get_config().get('model_id', 'N/A')} doesn't support multiple contents"
+            )
+        elif len(content_blocks) == 1:
+            return _format_content(content_blocks[0])
+        else:
+            return ""
+
+    def _format_request_message_tool_call(self, tool_use: ToolUse) -> dict[str, Any]:
+        """Format a Writer tool call.
+
+        Args:
+            tool_use: Tool use requested by the model.
+
+        Returns:
+            Writer formatted tool call.
+        """
+        return {
+            "function": {
+                "arguments": json.dumps(tool_use["input"]),
+                "name": tool_use["name"],
+            },
+            "id": tool_use["toolUseId"],
+            "type": "function",
+        }
+
+    def _format_request_tool_message(self, tool_result: ToolResult) -> dict[str, Any]:
+        """Format a Writer tool message.
+
+        Args:
+            tool_result: Tool result collected from a tool execution.
+
+        Returns:
+            Writer formatted tool message.
+        """
+        contents = cast(
+            list[ContentBlock],
+            [
+                {"text": json.dumps(content["json"])} if "json" in content else content
+                for content in tool_result["content"]
+            ],
+        )
+
+        if self.get_config().get("model_id", "") == "palmyra-x5":
+            formatted_contents = self._format_request_message_contents_vision(contents)
+        else:
+            formatted_contents = self._format_request_message_contents(contents)  # type: ignore [assignment]
+
+        return {
+            "role": "tool",
+            "tool_call_id": tool_result["toolUseId"],
+            "content": formatted_contents,
+        }
+
+    def _format_request_messages(self, messages: Messages, system_prompt: Optional[str] = None) -> list[dict[str, Any]]:
+        """Format a Writer compatible messages array.
+
+        Args:
+            messages: List of message objects to be processed by the model.
+            system_prompt: System prompt to provide context to the model.
+
+        Returns:
+            Writer compatible messages array.
+        """
+        formatted_messages: list[dict[str, Any]]
+        formatted_messages = [{"role": "system", "content": system_prompt}] if system_prompt else []
+
+        for message in messages:
+            contents = message["content"]
+
+            # Only palmyra V5 support multiple content. Other models support only '{"content": "text_content"}'
+            if self.get_config().get("model_id", "") == "palmyra-x5":
+                formatted_contents: str | list[dict[str, Any]] = self._format_request_message_contents_vision(contents)
+            else:
+                formatted_contents = self._format_request_message_contents(contents)
+
+            formatted_tool_calls = [
+                self._format_request_message_tool_call(content["toolUse"])
+                for content in contents
+                if "toolUse" in content
+            ]
+            formatted_tool_messages = [
+                self._format_request_tool_message(content["toolResult"])
+                for content in contents
+                if "toolResult" in content
+            ]
+
+            formatted_message = {
+                "role": message["role"],
+                "content": formatted_contents if len(formatted_contents) > 0 else "",
+                **({"tool_calls": formatted_tool_calls} if formatted_tool_calls else {}),
+            }
+            formatted_messages.append(formatted_message)
+            formatted_messages.extend(formatted_tool_messages)
+
+        return [message for message in formatted_messages if message["content"] or "tool_calls" in message]
+
+    def format_request(
+        self, messages: Messages, tool_specs: Optional[list[ToolSpec]] = None, system_prompt: Optional[str] = None
+    ) -> Any:
+        """Format a streaming request to the underlying 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.
+
+        Returns:
+            The formatted request.
+        """
+        request = {
+            **{k: v for k, v in self.config.items()},
+            "messages": self._format_request_messages(messages, system_prompt),
+            "stream": True,
+        }
+        try:
+            request["model"] = request.pop(
+                "model_id"
+            )  # To be consisted with other models WriterConfig use 'model_id' arg, but Writer API wait for 'model' arg
+        except KeyError as e:
+            raise KeyError("Please specify a model ID. Use 'model_id' keyword argument.") from e
+
+        # Writer don't support empty tools attribute
+        if tool_specs:
+            request["tools"] = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": tool_spec["name"],
+                        "description": tool_spec["description"],
+                        "parameters": tool_spec["inputSchema"]["json"],
+                    },
+                }
+                for tool_spec in tool_specs
+            ]
+
+        return request
+
+    def format_chunk(self, event: Any) -> StreamEvent:
+        """Format the model response events into standardized message chunks.
+
+        Args:
+            event: A response event from the model.
+
+        Returns:
+            The formatted chunk.
+        """
+        match event.get("chunk_type", ""):
+            case "message_start":
+                return {"messageStart": {"role": "assistant"}}
+
+            case "content_block_start":
+                if event["data_type"] == "text":
+                    return {"contentBlockStart": {"start": {}}}
+
+                return {
+                    "contentBlockStart": {
+                        "start": {
+                            "toolUse": {
+                                "name": event["data"].function.name,
+                                "toolUseId": event["data"].id,
+                            }
+                        }
+                    }
+                }
+
+            case "content_block_delta":
+                if event["data_type"] == "text":
+                    return {"contentBlockDelta": {"delta": {"text": event["data"]}}}
+
+                return {"contentBlockDelta": {"delta": {"toolUse": {"input": event["data"].function.arguments}}}}
+
+            case "content_block_stop":
+                return {"contentBlockStop": {}}
+
+            case "message_stop":
+                match event["data"]:
+                    case "tool_calls":
+                        return {"messageStop": {"stopReason": "tool_use"}}
+                    case "length":
+                        return {"messageStop": {"stopReason": "max_tokens"}}
+                    case _:
+                        return {"messageStop": {"stopReason": "end_turn"}}
+
+            case "metadata":
+                return {
+                    "metadata": {
+                        "usage": {
+                            "inputTokens": event["data"].prompt_tokens if event["data"] else 0,
+                            "outputTokens": event["data"].completion_tokens if event["data"] else 0,
+                            "totalTokens": event["data"].total_tokens if event["data"] else 0,
+                        },  # If 'stream_options' param is unset, empty metadata will be provided.
+                        # To avoid errors replacing expected fields with default zero value
+                        "metrics": {
+                            "latencyMs": 0,  # All palmyra models don't provide 'latency' metadata
+                        },
+                    },
+                }
+
+            case _:
+                raise RuntimeError(f"chunk_type=<{event['chunk_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 Writer 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:
+            ModelThrottledException: When the model service is throttling requests from the client.
+        """
+        logger.debug("formatting request")
+        request = self.format_request(messages, tool_specs, system_prompt)
+        logger.debug("request=<%s>", request)
+
+        logger.debug("invoking model")
+        try:
+            response = await self.client.chat.chat(**request)
+        except writerai.RateLimitError as e:
+            raise ModelThrottledException(str(e)) from e
+
+        yield self.format_chunk({"chunk_type": "message_start"})
+        yield self.format_chunk({"chunk_type": "content_block_start", "data_type": "text"})
+
+        tool_calls: dict[int, list[Any]] = {}
+
+        async for chunk in response:
+            if not getattr(chunk, "choices", None):
+                continue
+            choice = chunk.choices[0]
+
+            if choice.delta.content:
+                yield self.format_chunk(
+                    {"chunk_type": "content_block_delta", "data_type": "text", "data": choice.delta.content}
+                )
+
+            for tool_call in choice.delta.tool_calls or []:
+                tool_calls.setdefault(tool_call.index, []).append(tool_call)
+
+            if choice.finish_reason:
+                break
+
+        yield self.format_chunk({"chunk_type": "content_block_stop", "data_type": "text"})
+
+        for tool_deltas in tool_calls.values():
+            tool_start, tool_deltas = tool_deltas[0], tool_deltas[1:]
+            yield self.format_chunk({"chunk_type": "content_block_start", "data_type": "tool", "data": tool_start})
+
+            for tool_delta in tool_deltas:
+                yield self.format_chunk({"chunk_type": "content_block_delta", "data_type": "tool", "data": tool_delta})
+
+            yield self.format_chunk({"chunk_type": "content_block_stop", "data_type": "tool"})
+
+        yield self.format_chunk({"chunk_type": "message_stop", "data": choice.finish_reason})
+
+        # Iterating until the end to fetch metadata chunk
+        async for chunk in response:
+            _ = chunk
+
+        yield self.format_chunk({"chunk_type": "metadata", "data": chunk.usage})
+
+        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.
+        """
+        formatted_request = self.format_request(messages=prompt, tool_specs=None, system_prompt=system_prompt)
+        formatted_request["response_format"] = {
+            "type": "json_schema",
+            "json_schema": {"schema": output_model.model_json_schema()},
+        }
+        formatted_request["stream"] = False
+        formatted_request.pop("stream_options", None)
+
+        response = await self.client.chat.chat(**formatted_request)
+
+        try:
+            content = response.choices[0].message.content.strip()
+            yield {"output": output_model.model_validate_json(content)}
+        except Exception as e:
+            raise ValueError(f"Failed to parse or load content into model: {e}") from e

agentrun_sdk/multiagent/__init__.py

@@ -0,0 +1,22 @@
+"""Multiagent capabilities for Strands Agents.
+
+This module provides support for multiagent systems, including agent-to-agent (A2A)
+communication protocols and coordination mechanisms.
+
+Submodules:
+    a2a: Implementation of the Agent-to-Agent (A2A) protocol, which enables
+         standardized communication between agents.
+"""
+
+from .base import MultiAgentBase, MultiAgentResult
+from .graph import GraphBuilder, GraphResult
+from .swarm import Swarm, SwarmResult
+
+__all__ = [
+    "GraphBuilder",
+    "GraphResult",
+    "MultiAgentBase",
+    "MultiAgentResult",
+    "Swarm",
+    "SwarmResult",
+]

agentrun_sdk/multiagent/a2a/__init__.py

@@ -0,0 +1,15 @@
+"""Agent-to-Agent (A2A) communication protocol implementation for Strands Agents.
+
+This module provides classes and utilities for enabling Strands Agents to communicate
+with other agents using the Agent-to-Agent (A2A) protocol.
+
+Docs: https://google-a2a.github.io/A2A/latest/
+
+Classes:
+    A2AAgent: A wrapper that adapts a Strands Agent to be A2A-compatible.
+"""
+
+from .executor import StrandsA2AExecutor
+from .server import A2AServer
+
+__all__ = ["A2AServer", "StrandsA2AExecutor"]

agentrun_sdk/multiagent/a2a/executor.py

@@ -0,0 +1,148 @@
+"""Strands Agent executor for the A2A protocol.
+
+This module provides the StrandsA2AExecutor class, which adapts a Strands Agent
+to be used as an executor in the A2A protocol. It handles the execution of agent
+requests and the conversion of Strands Agent streamed responses to A2A events.
+
+The A2A AgentExecutor ensures clients receive responses for synchronous and
+streamed requests to the A2AServer.
+"""
+
+import logging
+from typing import Any
+
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events import EventQueue
+from a2a.server.tasks import TaskUpdater
+from a2a.types import InternalError, Part, TaskState, TextPart, UnsupportedOperationError
+from a2a.utils import new_agent_text_message, new_task
+from a2a.utils.errors import ServerError
+
+from ...agent.agent import Agent as SAAgent
+from ...agent.agent import AgentResult as SAAgentResult
+
+logger = logging.getLogger(__name__)
+
+
+class StrandsA2AExecutor(AgentExecutor):
+    """Executor that adapts a Strands Agent to the A2A protocol.
+
+    This executor uses streaming mode to handle the execution of agent requests
+    and converts Strands Agent responses to A2A protocol events.
+    """
+
+    def __init__(self, agent: SAAgent):
+        """Initialize a StrandsA2AExecutor.
+
+        Args:
+            agent: The Strands Agent instance to adapt to the A2A protocol.
+        """
+        self.agent = agent
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        """Execute a request using the Strands Agent and send the response as A2A events.
+
+        This method executes the user's input using the Strands Agent in streaming mode
+        and converts the agent's response to A2A events.
+
+        Args:
+            context: The A2A request context, containing the user's input and task metadata.
+            event_queue: The A2A event queue used to send response events back to the client.
+
+        Raises:
+            ServerError: If an error occurs during agent execution
+        """
+        task = context.current_task
+        if not task:
+            task = new_task(context.message)  # type: ignore
+            await event_queue.enqueue_event(task)
+
+        updater = TaskUpdater(event_queue, task.id, task.context_id)
+
+        try:
+            await self._execute_streaming(context, updater)
+        except Exception as e:
+            raise ServerError(error=InternalError()) from e
+
+    async def _execute_streaming(self, context: RequestContext, updater: TaskUpdater) -> None:
+        """Execute request in streaming mode.
+
+        Streams the agent's response in real-time, sending incremental updates
+        as they become available from the agent.
+
+        Args:
+            context: The A2A request context, containing the user's input and other metadata.
+            updater: The task updater for managing task state and sending updates.
+        """
+        logger.info("Executing request in streaming mode")
+        user_input = context.get_user_input()
+        try:
+            async for event in self.agent.stream_async(user_input):
+                await self._handle_streaming_event(event, updater)
+        except Exception:
+            logger.exception("Error in streaming execution")
+            raise
+
+    async def _handle_streaming_event(self, event: dict[str, Any], updater: TaskUpdater) -> None:
+        """Handle a single streaming event from the Strands Agent.
+
+        Processes streaming events from the agent, converting data chunks to A2A
+        task updates and handling the final result when streaming is complete.
+
+        Args:
+            event: The streaming event from the agent, containing either 'data' for
+                incremental content or 'result' for the final response.
+            updater: The task updater for managing task state and sending updates.
+        """
+        logger.debug("Streaming event: %s", event)
+        if "data" in event:
+            if text_content := event["data"]:
+                await updater.update_status(
+                    TaskState.working,
+                    new_agent_text_message(
+                        text_content,
+                        updater.context_id,
+                        updater.task_id,
+                    ),
+                )
+        elif "result" in event:
+            await self._handle_agent_result(event["result"], updater)
+
+    async def _handle_agent_result(self, result: SAAgentResult | None, updater: TaskUpdater) -> None:
+        """Handle the final result from the Strands Agent.
+
+        Processes the agent's final result, extracts text content from the response,
+        and adds it as an artifact to the task before marking the task as complete.
+
+        Args:
+            result: The agent result object containing the final response, or None if no result.
+            updater: The task updater for managing task state and adding the final artifact.
+        """
+        if final_content := str(result):
+            await updater.add_artifact(
+                [Part(root=TextPart(text=final_content))],
+                name="agent_response",
+            )
+        await updater.complete()
+
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Cancel an ongoing execution.
+
+        This method is called when a request cancellation is requested. Currently,
+        cancellation is not supported by the Strands Agent executor, so this method
+        always raises an UnsupportedOperationError.
+
+        Args:
+            context: The A2A request context.
+            event_queue: The A2A event queue.
+
+        Raises:
+            ServerError: Always raised with an UnsupportedOperationError, as cancellation
+                is not currently supported.
+        """
+        logger.warning("Cancellation requested but not supported")
+        raise ServerError(error=UnsupportedOperationError())