feat: add Bedrock integration (BIP-0042) as separate PR

Author: vaquarkhan

burr/integrations/__init__.py

@@ -14,3 +14,14 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+
+def __getattr__(name: str):
+    """Lazy load Bedrock integration to avoid requiring boto3 unless used."""
+    if name == "BedrockAction":
+        from burr.integrations.bedrock import BedrockAction
+        return BedrockAction
+    if name == "BedrockStreamingAction":
+        from burr.integrations.bedrock import BedrockStreamingAction
+        return BedrockStreamingAction
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

burr/integrations/bedrock.py

@@ -0,0 +1,261 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Amazon Bedrock integration for Burr.
+
+This module provides Action classes for invoking Amazon Bedrock models
+within Burr applications.
+
+Example usage:
+    from burr.integrations.bedrock import BedrockAction
+
+    def prompt_mapper(state):
+        return {
+            "messages": [{"role": "user", "content": state["user_input"]}],
+            "system": [{"text": "You are a helpful assistant."}],
+        }
+
+    # With default client (created lazily on first use):
+    action = BedrockAction(
+        model_id="anthropic.claude-3-sonnet-20240229-v1:0",
+        input_mapper=prompt_mapper,
+        reads=["user_input"],
+        writes=["response"],
+    )
+
+    # With injected client (for tests or distributed execution):
+    # client = boto3.client("bedrock-runtime", region_name="us-east-1")
+    # action = BedrockAction(..., client=client)
+"""
+
+import logging
+from typing import Any, Generator, Optional, Protocol
+
+from burr.core.action import SingleStepAction, StreamingAction
+from burr.core.state import State
+from burr.integrations.base import require_plugin
+
+logger = logging.getLogger(__name__)
+
+# Type for injected Bedrock client (avoids boto3 import at type-check time)
+BedrockClient = Any
+
+try:
+    import boto3
+    from botocore.config import Config
+    from botocore.exceptions import ClientError
+except ImportError as e:
+    require_plugin(e, "bedrock")
+
+
+class StateToPromptMapper(Protocol):
+    """Protocol for mapping Burr state to Bedrock prompt format."""
+
+    def __call__(self, state: State) -> dict[str, Any]:
+        ...
+
+
+class BedrockAction(SingleStepAction):
+    """Action that invokes Amazon Bedrock models using the Converse API."""
+
+    def __init__(
+        self,
+        model_id: str,
+        input_mapper: StateToPromptMapper,
+        reads: list[str],
+        writes: list[str],
+        name: str = "bedrock_invoke",
+        region: Optional[str] = None,
+        guardrail_id: Optional[str] = None,
+        guardrail_version: Optional[str] = None,
+        inference_config: Optional[dict[str, Any]] = None,
+        max_retries: int = 3,
+        client: Optional[BedrockClient] = None,
+    ):
+        super().__init__()
+        self._model_id = model_id
+        self._input_mapper = input_mapper
+        self._reads = reads
+        self._writes = writes
+        self._name = name
+        self._region = region
+        self._guardrail_id = guardrail_id
+        self._guardrail_version = guardrail_version or "DRAFT"
+        self._inference_config = inference_config or {"maxTokens": 4096}
+        self._max_retries = max_retries
+        self._client = client
+
+    def _get_client(self) -> BedrockClient:
+        """Return the Bedrock client, creating lazily if not injected."""
+        if self._client is not None:
+            return self._client
+        config = Config(
+            retries={"max_attempts": self._max_retries, "mode": "adaptive"}
+        )
+        self._client = boto3.client(
+            "bedrock-runtime", region_name=self._region, config=config
+        )
+        return self._client
+
+    @property
+    def reads(self) -> list[str]:
+        return self._reads
+
+    @property
+    def writes(self) -> list[str]:
+        return self._writes
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def run_and_update(self, state: State, **run_kwargs) -> tuple[dict, State]:
+        prompt = self._input_mapper(state)
+
+        request: dict[str, Any] = {
+            "modelId": self._model_id,
+            "messages": prompt["messages"],
+            "inferenceConfig": self._inference_config,
+        }
+
+        if "system" in prompt:
+            request["system"] = prompt["system"]
+
+        if self._guardrail_id:
+            request["guardrailConfig"] = {
+                "guardrailIdentifier": self._guardrail_id,
+                "guardrailVersion": self._guardrail_version,
+            }
+
+        try:
+            response = self._get_client().converse(**request)
+        except ClientError as e:
+            logger.error("Bedrock API error: %s", e)
+            raise
+
+        output_message = response["output"]["message"]
+        content_blocks = output_message.get("content", [])
+        text = content_blocks[0]["text"] if content_blocks else ""
+
+        result: dict[str, Any] = {
+            "response": text,
+            "usage": response.get("usage", {}),
+            "stop_reason": response.get("stopReason"),
+        }
+
+        updates = {key: result[key] for key in self._writes if key in result}
+        new_state = state.update(**updates)
+
+        return result, new_state
+
+
+class BedrockStreamingAction(StreamingAction):
+    """Streaming variant of BedrockAction using Converse Stream API."""
+
+    def __init__(
+        self,
+        model_id: str,
+        input_mapper: StateToPromptMapper,
+        reads: list[str],
+        writes: list[str],
+        name: str = "bedrock_stream",
+        region: Optional[str] = None,
+        guardrail_id: Optional[str] = None,
+        guardrail_version: Optional[str] = None,
+        inference_config: Optional[dict[str, Any]] = None,
+        max_retries: int = 3,
+        client: Optional[BedrockClient] = None,
+    ):
+        super().__init__()
+        self._model_id = model_id
+        self._input_mapper = input_mapper
+        self._reads = reads
+        self._writes = writes
+        self._name = name
+        self._region = region
+        self._guardrail_id = guardrail_id
+        self._guardrail_version = guardrail_version or "DRAFT"
+        self._inference_config = inference_config or {"maxTokens": 4096}
+        self._max_retries = max_retries
+        self._client = client
+
+    def _get_client(self) -> BedrockClient:
+        """Return the Bedrock client, creating lazily if not injected."""
+        if self._client is not None:
+            return self._client
+        config = Config(
+            retries={"max_attempts": self._max_retries, "mode": "adaptive"}
+        )
+        self._client = boto3.client(
+            "bedrock-runtime", region_name=self._region, config=config
+        )
+        return self._client
+
+    @property
+    def reads(self) -> list[str]:
+        return self._reads
+
+    @property
+    def writes(self) -> list[str]:
+        return self._writes
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def stream_run(
+        self, state: State, **run_kwargs
+    ) -> Generator[dict, None, None]:
+        prompt = self._input_mapper(state)
+
+        request: dict[str, Any] = {
+            "modelId": self._model_id,
+            "messages": prompt["messages"],
+            "inferenceConfig": self._inference_config,
+        }
+
+        if "system" in prompt:
+            request["system"] = prompt["system"]
+
+        if self._guardrail_id:
+            request["guardrailConfig"] = {
+                "guardrailIdentifier": self._guardrail_id,
+                "guardrailVersion": self._guardrail_version,
+            }
+
+        try:
+            response = self._get_client().converse_stream(**request)
+        except ClientError as e:
+            logger.error("Bedrock streaming API error: %s", e)
+            raise
+
+        full_response = ""
+        stream = response.get("stream", [])
+        for event in stream:
+            if "contentBlockDelta" in event:
+                chunk = event["contentBlockDelta"]["delta"].get("text", "")
+                full_response += chunk
+                yield {"chunk": chunk, "response": full_response}
+
+        yield {"chunk": "", "response": full_response, "complete": True}
+
+    def update(self, result: dict, state: State) -> State:
+        if result.get("complete"):
+            updates = {"response": result.get("response", "")}
+            filtered = {k: v for k, v in updates.items() if k in self._writes}
+            return state.update(**filtered)
+        return state