From 0bf59dc8470a1104b5c45f797d3e017e08f58726 Mon Sep 17 00:00:00 2001 From: Vaquar Khan Date: Mon, 20 Apr 2026 23:41:46 -0500 Subject: [PATCH] Bedrock integration backup (#677) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Bedrock integration for Burr Adds first-class support for Amazon Bedrock's Converse API to Burr, so users can build state machines and agents on top of Bedrock models (Claude, Titan, Llama, etc.) without writing boto3 plumbing themselves. ### Architecture The integration is structured in three layers: a pure "core" that knows how to talk to Bedrock, a thin mixin that adapts it to Burr's action interface, and two concrete action classes that plug into Burr's state machine. Layer 1 _BedrockCore (transport + request building). This is a plain Python class with no knowledge of Burr actions. It holds the configuration (model id, region, guardrail, inference config, retry count, optional injected client) and exposes two things: get_client(), which returns the bedrock-runtime boto3 client (either the one injected at construction time, or a new one built lazily with a Config(retries=...) on first use), and build_converse_request(state), which calls the user's input_mapper(state) and assembles the kwargs dict passed to converse() / converse_stream() — modelId, messages, inferenceConfig, and the optional system and guardrailConfig blocks. Guardrail validation lives here: if guardrail_id is set without an explicit guardrail_version, the constructor raises ValueError so nobody accidentally ships against an unpublished DRAFT. Layer 2 - _BedrockBase (Burr adapter mixin). A small mixin that holds a _BedrockCore instance and exposes the three properties Burr actions need: reads, writes, and name. It exists purely to avoid duplicating constructor wiring and property boilerplate between the sync and streaming classes. Layer 3 - concrete actions. BedrockAction inherits from _BedrockBase and Burr's SingleStepAction. Its run_and_update() builds the Converse request, calls converse(), runs the response through _text_from_content_blocks() (which joins every text block in the response, so multi-block replies and mixed text/tool-use messages aren't silently truncated), and then uses _model_result_for_writes() to map the model output to whichever keys the user declared in writes — "response" gets the joined text, "usage" gets the token counts, "stop_reason" gets the Bedrock stop reason, and any other custom key also gets the joined text. BedrockStreamingAction inherits from _BedrockBase and Burr's StreamingAction. Its stream_run() is a generator that iterates the Bedrock event stream, collecting text chunks into a list[str] (joined once at the end — not concatenated in a loop), and yields one dict per contentBlockDelta event plus a final complete: True payload. Its update() runs only on the final chunk and reuses the same _model_result_for_writes() helper, so streaming and non-streaming actions write state using identical logic. Helper functions. Two pure functions sit at module scope and are shared by both actions: _text_from_content_blocks(blocks) tolerates non-text blocks (e.g. toolUse) and returns the joined text, and _model_result_for_writes(text, usage, stop_reason, writes) centralises the "which result goes into which state key" logic so the sync and streaming paths can't drift out of sync. #### Cross-cutting concerns. boto3 is imported inside a try/except that calls require_plugin(...) on failure, so installing Burr without the [bedrock] extra doesn't break anything until a user actually touches the module. The BedrockAction / BedrockStreamingAction symbols are exposed via a __getattr__ in __init__.py , giving callers the ergonomic from burr.integrations import BedrockAction without importing boto3 for every Burr user. Errors from boto3 (ClientError) are logged at ERROR level and re-raised unchanged, so Burr's own retry/lifecycle hooks can handle them at the application level. State machine shape. Both actions are drop-in nodes in a Burr graph. A typical chatbot wiring is human_input → bedrock_call → save_response → human_input, where bedrock_call is one of these two actions reading chat_history and writing response. For agents, bedrock_call is followed by a conditional transition into a tool-executor action and back again — the integration itself stays stateless across calls; all conversation and tool state lives in Burr's State. ## Changes - `burr/integrations/bedrock.py` - `BedrockAction` — single-step action wrapping Bedrock `converse()`. Handles lazy client creation, retries, guardrails, inference config, and maps model output into state. - `BedrockStreamingAction` — streaming variant using `converse_stream()`, yields chunks incrementally and merges final text into state on completion. - Shared `_BedrockBase` + `_BedrockCore` to avoid duplication across the two classes. - Helpers `_text_from_content_blocks` (joins all text blocks, tolerates tool-use blocks) and `_model_result_for_writes` (dynamically maps user's `writes` keys to response fields). - `burr/integrations/__init__.py` — lazy exports for `BedrockAction` / `BedrockStreamingAction` so `boto3` is only imported when the classes are actually used. - `pyproject.toml` — adds `[bedrock]` optional extra (`boto3`). - `docs/reference/integrations/bedrock.rst` — Sphinx reference docs with install and IAM setup notes. - `docs/getting_started/install.rst` — adds Bedrock to the install matrix. - `examples/integrations/bedrock/` — runnable example with `application()` (sync) and `streaming_application()` (streaming) functions, plus README and requirements. - `.github/workflows/python-package.yml` — dedicated `test-bedrock` job so AWS deps stay out of the default test matrix. - `tests/integrations/test_bip0042_bedrock.py` — 23 unit tests covering imports, guardrail validation, sync/streaming interfaces, multi-block responses, custom `writes` keys, and error propagation. ## Design notes - **No `boto3` import at module load** — guarded by `require_plugin`, so installs without the `[bedrock]` extra still work. - **Client injection** — both actions accept a pre-built `bedrock-runtime` client for tests and distributed execution; otherwise the client is created lazily on first use. - **Guardrail safety** — setting `guardrail_id` without an explicit `guardrail_version` raises `ValueError` rather than silently defaulting to `DRAFT`. - **Empty vs. default inference config** — `inference_config={}` is respected; `None` falls back to `{"maxTokens": 4096}`. - **Custom `writes` keys** — both sync and streaming actions map the model's text output to whichever `writes` key the user declares (e.g. `writes=["answer"]`), with `"response"` kept for backwards compatibility. ## How I tested this - Ran the full unit suite: 23/23 bedrock tests pass, broader Burr suite unchanged. - Ran the example against real Bedrock (Claude 3 Haiku, `us-east-1`) — both sync and streaming return valid output. - Exercised end-to-end scenarios: multi-turn chat with state-managed history, custom `writes` keys, multi-content-block responses, guardrail validation paths, and `ClientError` propagation. ## Notes - Follow-ups I'd suggest as separate PRs: - Async variants (`AsyncBedrockAction` / `AsyncBedrockStreamingAction`) on top of `AsyncStreamingAction`. - Native `toolConfig` support in `BedrockAction` for building agents without going around the abstraction. - `stream_result` with typed state via `pydantic` integration. ## Checklist - [x] PR has an informative and human-readable title - [x] Changes are limited to a single goal (Bedrock integration only — tracking/SQS work split out per earlier review) - [x] Code passed pre-commit - [x] New functionality is tested (23 unit tests + real-API smoke) - [x] New functions are documented (docstrings + Sphinx reference) - [x] Project documentation updated (`docs/reference/integrations/bedrock.rst`, `docs/getting_started/install.rst`, example README) --- Commits * Add Bedrock integration, tests, docs, and CI * fix(bedrock): streaming writes parity, multi-block text, list join for stream * style: black-format bedrock integration * Format Bedrock integration; add runnable Bedrock example Apply Black formatting to burr/integrations/bedrock.py for CI. Add examples/integrations/bedrock with BedrockAction and BedrockStreamingAction, requirements and README, and list it in examples/README. * chore: fix burr/examples newline for pre-commit end-of-file-fixer expects a trailing newline on burr/examples so pre-commit run --all-files passes. * fix bedrock converse content blocks and set_user_input return Made-with: Cursor * fixed review comments Made-with: Cursor * Fix Bedrock test imports for isort compatibility --------- Co-authored-by: vaquarkhan --- .github/workflows/python-package.yml | 25 +- burr/examples | 2 +- burr/integrations/bedrock.py | 367 +++++++++++++++++ docs/getting_started/install.rst | 6 + docs/reference/integrations/bedrock.rst | 44 +++ docs/reference/integrations/index.rst | 1 + examples/README.md | 1 + examples/integrations/bedrock/README.md | 45 +++ examples/integrations/bedrock/__init__.py | 16 + examples/integrations/bedrock/application.py | 140 +++++++ .../integrations/bedrock/requirements.txt | 18 + pyproject.toml | 5 + tests/integrations/test_bip0042_bedrock.py | 371 ++++++++++++++++++ 13 files changed, 1039 insertions(+), 2 deletions(-) create mode 100644 burr/integrations/bedrock.py create mode 100644 docs/reference/integrations/bedrock.rst create mode 100644 examples/integrations/bedrock/README.md create mode 100644 examples/integrations/bedrock/__init__.py create mode 100644 examples/integrations/bedrock/application.py create mode 100644 examples/integrations/bedrock/requirements.txt create mode 100644 tests/integrations/test_bip0042_bedrock.py diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 794677227..407db459c 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -73,7 +73,7 @@ jobs: - name: Run tests run: | - python -m pytest tests --ignore=tests/integrations/persisters + python -m pytest tests --ignore=tests/integrations/persisters --ignore=tests/integrations/test_bip0042_bedrock.py test-tracking-server-s3: runs-on: ubuntu-latest @@ -98,6 +98,29 @@ jobs: run: | python -m pytest tests/tracking/test_bip0042_s3_buffering.py -v + test-bedrock: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.9', '3.10', '3.11', '3.12'] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install -e ".[tests,bedrock]" + + - name: Run Bedrock integration tests + run: | + python -m pytest tests/integrations/test_bip0042_bedrock.py -v + validate-examples: runs-on: ubuntu-latest steps: diff --git a/burr/examples b/burr/examples index a6573af9c..beeced1f9 120000 --- a/burr/examples +++ b/burr/examples @@ -1 +1 @@ -../examples \ No newline at end of file +../examples diff --git a/burr/integrations/bedrock.py b/burr/integrations/bedrock.py new file mode 100644 index 000000000..780494e5a --- /dev/null +++ b/burr/integrations/bedrock.py @@ -0,0 +1,367 @@ +# 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": [{"text": 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) + +If ``guardrail_id`` is set, you must pass an explicit ``guardrail_version`` +(including the string ``DRAFT`` if you intend to use the unpublished draft). +""" + +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]: + ... # noqa: E704 + + +def _text_from_content_blocks(content_blocks: list[Any]) -> str: + """Join text from every content block (multi-block replies, tool use + text, etc.).""" + parts: list[str] = [] + for block in content_blocks: + if isinstance(block, dict) and "text" in block: + parts.append(block["text"]) + return "\n".join(parts) + + +def _model_result_for_writes( + text: str, + usage: dict[str, Any], + stop_reason: Any, + writes: list[str], +) -> dict[str, Any]: + """Build the result dict and ensure each ``writes`` key maps to the right value.""" + result: dict[str, Any] = { + "response": text, + "usage": usage, + "stop_reason": stop_reason, + } + for w in writes: + if w == "usage": + result[w] = usage + elif w == "stop_reason": + result[w] = stop_reason + else: + result[w] = text + return result + + +class _BedrockCore: + """Shared Bedrock client, inference config, and Converse request shape.""" + + def __init__( + self, + model_id: str, + input_mapper: StateToPromptMapper, + reads: list[str], + writes: list[str], + name: str, + region: Optional[str], + guardrail_id: Optional[str], + guardrail_version: Optional[str], + inference_config: Optional[dict[str, Any]], + max_retries: int, + client: Optional[BedrockClient], + ): + if guardrail_id is not None and guardrail_version is None: + raise ValueError( + "guardrail_version is required when guardrail_id is set " + '(pass an explicit published version, or "DRAFT" for the draft).' + ) + 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 + self._inference_config = ( + {"maxTokens": 4096} if inference_config is None else inference_config + ) + self._max_retries = max_retries + self._client = 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 get_client(self) -> BedrockClient: + """Return the Bedrock runtime client, creating it 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 + + def build_converse_request(self, state: State) -> dict[str, Any]: + """Build the kwargs dict for ``converse`` / ``converse_stream`` from current 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, + } + return request + + +class _BedrockBase: + """Shared Bedrock wiring: core state and reads/writes/name for action subclasses.""" + + _bedrock: _BedrockCore + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + def _init_bedrock_core( + self, + model_id: str, + input_mapper: StateToPromptMapper, + reads: list[str], + writes: list[str], + name: str, + region: Optional[str], + guardrail_id: Optional[str], + guardrail_version: Optional[str], + inference_config: Optional[dict[str, Any]], + max_retries: int, + client: Optional[BedrockClient], + ) -> None: + self._bedrock = _BedrockCore( + model_id=model_id, + input_mapper=input_mapper, + reads=reads, + writes=writes, + name=name, + region=region, + guardrail_id=guardrail_id, + guardrail_version=guardrail_version, + inference_config=inference_config, + max_retries=max_retries, + client=client, + ) + + @property + def reads(self) -> list[str]: + return self._bedrock.reads + + @property + def writes(self) -> list[str]: + return self._bedrock.writes + + @property + def name(self) -> str: + return self._bedrock.name + + +class BedrockAction(_BedrockBase, SingleStepAction): + """Action that invokes Amazon Bedrock models using the Converse API. + + :param model_id: Bedrock model identifier (e.g. Anthropic Claude on Bedrock). + :param input_mapper: Callable mapping :class:`~burr.core.state.State` to Bedrock + ``messages`` / optional ``system`` keys. + :param reads: State keys this action reads. + :param writes: State keys to update (typically include ``response``). + :param name: Action name for the graph. + :param region: AWS region for the Bedrock runtime client (optional). + :param guardrail_id: If set, ``guardrail_version`` must also be set explicitly. + :param guardrail_version: Guardrail version string (use ``DRAFT`` only when intended). + :param inference_config: Passed as ``inferenceConfig``; if omitted, defaults to + a ``maxTokens`` limit. Pass an empty dict explicitly to send an empty config. + :param max_retries: Botocore retry configuration for the runtime client. + :param client: Optional pre-built ``bedrock-runtime`` client (for tests or injection). + + Use :meth:`run_and_update` to run the model and merge outputs into state. + """ + + 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._init_bedrock_core( + model_id=model_id, + input_mapper=input_mapper, + reads=reads, + writes=writes, + name=name, + region=region, + guardrail_id=guardrail_id, + guardrail_version=guardrail_version, + inference_config=inference_config, + max_retries=max_retries, + client=client, + ) + + def run_and_update(self, state: State, **run_kwargs) -> tuple[dict, State]: + request = self._bedrock.build_converse_request(state) + + try: + response = self._bedrock.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 = _text_from_content_blocks(content_blocks) + + result = _model_result_for_writes( + text, + response.get("usage", {}), + response.get("stopReason"), + self._bedrock.writes, + ) + + updates = {key: result[key] for key in self._bedrock.writes if key in result} + new_state = state.update(**updates) + + return result, new_state + + +class BedrockStreamingAction(_BedrockBase, StreamingAction): + """Streaming Bedrock action using the Converse Stream API. + + Parameters match :class:`BedrockAction` except the default ``name`` is + ``bedrock_stream``. Yields chunk dicts from :meth:`stream_run` and merges the + final response in :meth:`update`. + """ + + 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._init_bedrock_core( + model_id=model_id, + input_mapper=input_mapper, + reads=reads, + writes=writes, + name=name, + region=region, + guardrail_id=guardrail_id, + guardrail_version=guardrail_version, + inference_config=inference_config, + max_retries=max_retries, + client=client, + ) + + def stream_run(self, state: State, **run_kwargs) -> Generator[dict, None, None]: + request = self._bedrock.build_converse_request(state) + + try: + response = self._bedrock.get_client().converse_stream(**request) + except ClientError as e: + logger.error("Bedrock streaming API error: %s", e) + raise + + text_parts: list[str] = [] + stream = response.get("stream", []) + for event in stream: + if "contentBlockDelta" in event: + chunk = event["contentBlockDelta"]["delta"].get("text", "") + text_parts.append(chunk) + full_response = "".join(text_parts) + yield {"chunk": chunk, "response": full_response} + + full_text = "".join(text_parts) + payload = _model_result_for_writes(full_text, {}, None, self._bedrock.writes) + yield {"chunk": "", "complete": True, **payload} + + def update(self, result: dict, state: State) -> State: + if result.get("complete"): + updates = {key: result[key] for key in self._bedrock.writes if key in result} + return state.update(**updates) + return state diff --git a/docs/getting_started/install.rst b/docs/getting_started/install.rst index 1afa3fbd2..335401758 100644 --- a/docs/getting_started/install.rst +++ b/docs/getting_started/install.rst @@ -174,3 +174,9 @@ This installs the server dependencies to run the UI and load tracking that was s pip install "burr[tracking-server]" This installs the server dependencies for running the UI off a filesystem. + +.. code-block:: bash + + pip install "burr[bedrock]" + +This installs ``boto3`` for the :ref:`Amazon Bedrock integration ` (``BedrockAction`` / ``BedrockStreamingAction``). diff --git a/docs/reference/integrations/bedrock.rst b/docs/reference/integrations/bedrock.rst new file mode 100644 index 000000000..934b62694 --- /dev/null +++ b/docs/reference/integrations/bedrock.rst @@ -0,0 +1,44 @@ +.. + 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. + + +.. _bedrock-integration: + +============== +Amazon Bedrock +============== + +`Amazon Bedrock `_ models can be called through the +Bedrock Runtime **Converse** and **Converse Stream** APIs using the actions below. + +Install the optional extra (pulls ``boto3``): + +.. code-block:: bash + + pip install "burr[bedrock]" + +IAM permissions must allow ``bedrock:InvokeModel`` / streaming equivalents for your +chosen models; see AWS documentation for your account and model IDs. + +.. autoclass:: burr.integrations.bedrock.BedrockAction + :members: run_and_update + :show-inheritance: + +.. autoclass:: burr.integrations.bedrock.BedrockStreamingAction + :members: stream_run, update + :show-inheritance: diff --git a/docs/reference/integrations/index.rst b/docs/reference/integrations/index.rst index ee489c0cc..562a8cd1b 100644 --- a/docs/reference/integrations/index.rst +++ b/docs/reference/integrations/index.rst @@ -34,3 +34,4 @@ Integrations -- we will be adding more pydantic haystack ray + bedrock diff --git a/examples/README.md b/examples/README.md index 740880768..9d4447cc2 100644 --- a/examples/README.md +++ b/examples/README.md @@ -48,5 +48,6 @@ Note we have a few more in [other-examples](other-examples/), but those do not y - [multi-agent-collaboration](multi-agent-collaboration/) - This example shows how to use Burr to create a multi-agent collaboration. This is a clone of the following [LangGraph example](https://github.com/langchain-ai/langgraph/blob/main/examples/multi_agent/multi-agent-collaboration.ipynb). - [multi-modal-chatbot](multi-modal-chatbot/) - This example shows how to use Burr to create a multi-modal chatbot. This demonstrates how to use a model to delegate to other models conditionally. - [streaming-overview](streaming-overview/) - This example shows how we can use the streaming API to respond to return quicker results to the user and build a seamless experience +- [integrations/bedrock](integrations/bedrock/) - Minimal graphs using Amazon Bedrock (`BedrockAction` and `BedrockStreamingAction`). - [tracing-and-spans](tracing-and-spans/) - This example shows how to use Burr to create a simple chatbot with additional visibility. This is a good starting point for understanding how to use Burr's tracing functionality. - [web-server](web-server/) - This example shows how to use Burr in a web server. This is a good starting point for understanding how to use Burr for interaction. diff --git a/examples/integrations/bedrock/README.md b/examples/integrations/bedrock/README.md new file mode 100644 index 000000000..0515a6d0a --- /dev/null +++ b/examples/integrations/bedrock/README.md @@ -0,0 +1,45 @@ + + +# Amazon Bedrock integration + +This example shows how to use Burr’s Bedrock helpers: + +- `BedrockAction` — single-step [Converse](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html) call. +- `BedrockStreamingAction` — streaming [ConverseStream](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStream.html) with `Application.stream_result`. + +## Setup + +1. Install dependencies (from the repo root): + + ```bash + pip install -r examples/integrations/bedrock/requirements.txt + ``` + +2. Configure AWS credentials and a region where Bedrock is available (for example `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_DEFAULT_REGION`). + +3. Ensure your account can invoke the model you choose. Override the default model with `BEDROCK_MODEL_ID` if needed (default in code is Claude 3 Haiku on Bedrock). + +## Run + +```bash +python examples/integrations/bedrock/application.py +``` + +The script runs a non-streaming call, then a streaming call, using two small Burr graphs defined in `application()` and `streaming_application()`. diff --git a/examples/integrations/bedrock/__init__.py b/examples/integrations/bedrock/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/examples/integrations/bedrock/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/examples/integrations/bedrock/application.py b/examples/integrations/bedrock/application.py new file mode 100644 index 000000000..be291acea --- /dev/null +++ b/examples/integrations/bedrock/application.py @@ -0,0 +1,140 @@ +# 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. + +"""Minimal Burr apps using :class:`~burr.integrations.bedrock.BedrockAction` +and :class:`~burr.integrations.bedrock.BedrockStreamingAction`. + +Configure AWS credentials (for example via ``aws configure`` or environment +variables) and optionally set ``BEDROCK_MODEL_ID`` and ``AWS_REGION`` / +``AWS_DEFAULT_REGION`` before running. +""" + +from __future__ import annotations + +import os + +from burr.core import Application, ApplicationBuilder, State, default +from burr.core.action import action +from burr.integrations import BedrockAction, BedrockStreamingAction + + +def _default_model_id() -> str: + return os.environ.get( + "BEDROCK_MODEL_ID", + "anthropic.claude-3-haiku-20240307-v1:0", + ) + + +def _aws_region() -> str | None: + return os.environ.get("AWS_DEFAULT_REGION") or os.environ.get("AWS_REGION") + + +def prompt_mapper(state: State) -> dict: + """Map Burr state to Bedrock Converse ``messages`` / ``system``.""" + return { + "messages": [{"role": "user", "content": [{"text": state["user_input"]}]}], + "system": [{"text": "You are a concise assistant."}], + } + + +@action(reads=[], writes=["user_input"]) +def set_user_input(state: State, user_input: str) -> State: + return state.update(user_input=user_input) + + +def application( + model_id: str | None = None, + region: str | None = None, +) -> Application: + """Builds a graph with :class:`~burr.integrations.bedrock.BedrockAction` (non-streaming).""" + invoke = BedrockAction( + model_id=model_id or _default_model_id(), + input_mapper=prompt_mapper, + reads=["user_input"], + writes=["response"], + name="invoke_bedrock", + region=region if region is not None else _aws_region(), + inference_config={"maxTokens": 512}, + ) + return ( + ApplicationBuilder() + .with_actions(set_prompt=set_user_input, invoke_bedrock=invoke) + .with_transitions( + ("set_prompt", "invoke_bedrock", default), + ) + .with_state(user_input="", response="") + .with_entrypoint("set_prompt") + .build() + ) + + +def streaming_application( + model_id: str | None = None, + region: str | None = None, +) -> Application: + """Builds a graph with :class:`~burr.integrations.bedrock.BedrockStreamingAction`.""" + stream = BedrockStreamingAction( + model_id=model_id or _default_model_id(), + input_mapper=prompt_mapper, + reads=["user_input"], + writes=["response"], + name="stream_bedrock", + region=region if region is not None else _aws_region(), + inference_config={"maxTokens": 512}, + ) + return ( + ApplicationBuilder() + .with_actions(set_prompt=set_user_input, stream_bedrock=stream) + .with_transitions( + ("set_prompt", "stream_bedrock", default), + ) + .with_state(user_input="", response="") + .with_entrypoint("set_prompt") + .build() + ) + + +def _demo_invoke() -> None: + app = application() + _, _, state = app.run( + halt_after=["invoke_bedrock"], + inputs={"user_input": "Explain what Burr is in one short sentence."}, + ) + print(state["response"]) + + +def _demo_stream() -> None: + app = streaming_application() + _, streaming_result = app.stream_result( + halt_after=["stream_bedrock"], + inputs={"user_input": "Count from 1 to 3, separated by commas."}, + ) + for item in streaming_result: + chunk = item.get("chunk") or "" + if chunk: + print(chunk, end="", flush=True) + print() + _, state = streaming_result.get() + print("Final response:", state["response"]) + + +if __name__ == "__main__": + print("--- BedrockAction (converse) ---") + _demo_invoke() + print() + print("--- BedrockStreamingAction (converse_stream) ---") + _demo_stream() diff --git a/examples/integrations/bedrock/requirements.txt b/examples/integrations/bedrock/requirements.txt new file mode 100644 index 000000000..d2d75e5a2 --- /dev/null +++ b/examples/integrations/bedrock/requirements.txt @@ -0,0 +1,18 @@ +# 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. + +burr[bedrock] diff --git a/pyproject.toml b/pyproject.toml index 70fb47e51..71b74dfa3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -103,6 +103,7 @@ tests = [ documentation = [ "apache-burr[tests]", + "apache-burr[bedrock]", "sphinx<9", "sphinx-autobuild", "myst-nb", @@ -128,6 +129,10 @@ tracking-client-s3 = [ "boto3" ] +bedrock = [ + "boto3" +] + tracking-server-s3 = [ "aerich", "aiobotocore", diff --git a/tests/integrations/test_bip0042_bedrock.py b/tests/integrations/test_bip0042_bedrock.py new file mode 100644 index 000000000..e2509edef --- /dev/null +++ b/tests/integrations/test_bip0042_bedrock.py @@ -0,0 +1,371 @@ +# 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. + +"""Tests for Bedrock integration.""" + +import inspect +from unittest.mock import MagicMock + +import pytest +from botocore.exceptions import ClientError + +from burr.core.action import SingleStepAction, StreamingAction +from burr.core.state import State + +boto3 = pytest.importorskip("boto3", reason="boto3 required for Bedrock tests") +import burr.integrations as integrations +from burr.integrations import bedrock +from burr.integrations.bedrock import BedrockAction, BedrockStreamingAction, StateToPromptMapper + + +class TestBedrockImports: + """Test Bedrock import paths.""" + + def test_import_bedrock_module_from_integrations(self): + """Verify `from burr.integrations import bedrock` works.""" + assert bedrock is not None + + def test_bedrock_module_has_expected_classes(self): + assert hasattr(bedrock, "BedrockAction") + assert hasattr(bedrock, "BedrockStreamingAction") + + def test_direct_import_bedrock_module(self): + """Verify bedrock.py module exists and has expected classes.""" + assert BedrockAction is not None + assert BedrockStreamingAction is not None + assert StateToPromptMapper is not None + + +class TestBedrockGuardrailValidation: + """guardrail_id requires an explicit guardrail_version.""" + + def test_raises_when_guardrail_id_without_version(self): + with pytest.raises(ValueError, match="guardrail_version is required"): + BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": []}, + reads=[], + writes=["response"], + guardrail_id="gr-123", + ) + + def test_accepts_explicit_draft_version(self): + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": [{"text": "ok"}]}}, + "usage": {}, + "stopReason": "end_turn", + } + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + guardrail_id="gr-123", + guardrail_version="DRAFT", + client=mock_client, + ) + action.run_and_update({}) + call_kw = mock_client.converse.call_args[1] + assert call_kw["guardrailConfig"]["guardrailVersion"] == "DRAFT" + + +class TestBedrockActionInterface: + """Test BedrockAction class interface with mocked boto3.""" + + def test_bedrock_action_extends_single_step_action(self): + """Verify BedrockAction extends SingleStepAction.""" + assert issubclass(BedrockAction, SingleStepAction) + + def test_bedrock_streaming_action_extends_streaming_action(self): + """Verify BedrockStreamingAction extends StreamingAction.""" + assert issubclass(BedrockStreamingAction, StreamingAction) + + def test_bedrock_action_has_required_properties(self): + """Verify BedrockAction has reads, writes, name properties.""" + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": []}, + reads=["input"], + writes=["output"], + ) + assert action.reads == ["input"] + assert action.writes == ["output"] + assert action.name == "bedrock_invoke" + + def test_bedrock_action_accepts_all_parameters(self): + """Verify BedrockAction accepts all specified parameters.""" + sig = inspect.signature(BedrockAction.__init__) + params = list(sig.parameters.keys()) + assert "model_id" in params + assert "input_mapper" in params + assert "reads" in params + assert "writes" in params + assert "name" in params + assert "region" in params + assert "guardrail_id" in params + assert "guardrail_version" in params + assert "inference_config" in params + assert "max_retries" in params + assert "client" in params + + def test_bedrock_action_uses_injected_client(self): + """Verify BedrockAction uses injected client when provided.""" + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": [{"text": "hi"}]}}, + "usage": {}, + "stopReason": "end_turn", + } + + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + client=mock_client, + ) + + result, _ = action.run_and_update({}) + assert result["response"] == "hi" + mock_client.converse.assert_called_once() + + def test_empty_inference_config_passed_through(self): + """inference_config={} must not be replaced by defaults (falsy dict bug).""" + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": [{"text": "x"}]}}, + "usage": {}, + "stopReason": "end_turn", + } + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + inference_config={}, + client=mock_client, + ) + action.run_and_update({}) + call_kw = mock_client.converse.call_args[1] + assert call_kw["inferenceConfig"] == {} + + def test_default_inference_config_when_omitted(self): + """When inference_config is None, Bedrock receives default maxTokens.""" + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": [{"text": "ok"}]}}, + "usage": {}, + "stopReason": "end_turn", + } + action = BedrockAction( + model_id="mid-1", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + client=mock_client, + ) + action.run_and_update({}) + assert mock_client.converse.call_args[1]["inferenceConfig"] == {"maxTokens": 4096} + + def test_converse_request_includes_model_id_and_system(self): + """Request passes modelId, messages, and optional system from input_mapper.""" + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": [{"text": "x"}]}}, + "usage": {}, + "stopReason": "end_turn", + } + + def mapper(s): + return { + "messages": [{"role": "user", "content": "hello"}], + "system": [{"text": "You are concise."}], + } + + action = BedrockAction( + model_id="anthropic.claude-v2", + input_mapper=mapper, + reads=[], + writes=["response"], + client=mock_client, + ) + action.run_and_update({}) + kw = mock_client.converse.call_args[1] + assert kw["modelId"] == "anthropic.claude-v2" + assert kw["messages"][0]["content"] == "hello" + assert kw["system"] == [{"text": "You are concise."}] + + def test_empty_content_blocks_return_empty_string(self): + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": []}}, + "usage": {}, + "stopReason": "end_turn", + } + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + client=mock_client, + ) + result, _ = action.run_and_update({}) + assert result["response"] == "" + + def test_multiple_text_blocks_joined_with_newline(self): + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": { + "message": { + "content": [ + {"text": "first"}, + {"text": "second"}, + ] + } + }, + "usage": {}, + "stopReason": "end_turn", + } + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + client=mock_client, + ) + result, _ = action.run_and_update({}) + assert result["response"] == "first\nsecond" + + def test_custom_write_key_updates_state(self): + mock_client = MagicMock() + mock_client.converse.return_value = { + "output": {"message": {"content": [{"text": "hello"}]}}, + "usage": {}, + "stopReason": "end_turn", + } + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["answer"], + client=mock_client, + ) + _, new_state = action.run_and_update(State({})) + assert new_state.get("answer") == "hello" + + def test_client_error_from_converse_propagates(self): + err = ClientError({"Error": {"Code": "ValidationException", "Message": "bad"}}, "Converse") + mock_client = MagicMock() + mock_client.converse.side_effect = err + + action = BedrockAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + client=mock_client, + ) + with pytest.raises(ClientError): + action.run_and_update({}) + + +class TestBedrockStreamingActionInterface: + """Test BedrockStreamingAction class interface with mocked boto3.""" + + def test_bedrock_streaming_action_uses_injected_client(self): + """Verify BedrockStreamingAction uses injected client when provided.""" + mock_client = MagicMock() + mock_client.converse_stream.return_value = { + "stream": [ + {"contentBlockDelta": {"delta": {"text": "hello "}}}, + {"contentBlockDelta": {"delta": {"text": "world"}}}, + ] + } + + action = BedrockStreamingAction( + model_id="test-model", + input_mapper=lambda s: {"messages": [{"role": "user", "content": "hi"}]}, + reads=[], + writes=["response"], + client=mock_client, + ) + + chunks = list(action.stream_run({})) + assert len(chunks) == 3 # 2 content chunks + 1 complete + assert chunks[0]["chunk"] == "hello " + assert chunks[1]["chunk"] == "world" + assert chunks[2]["complete"] is True + assert chunks[2]["response"] == "hello world" + mock_client.converse_stream.assert_called_once() + + def test_streaming_guardrail_requires_version(self): + with pytest.raises(ValueError, match="guardrail_version is required"): + BedrockStreamingAction( + model_id="test-model", + input_mapper=lambda s: {"messages": []}, + reads=[], + writes=["response"], + guardrail_id="gr-1", + ) + + def test_streaming_update_only_merges_on_complete(self): + """Non-final chunks should not update state via update().""" + action = BedrockStreamingAction( + model_id="test-model", + input_mapper=lambda s: {"messages": []}, + reads=[], + writes=["response"], + ) + s0 = State({"response": "old"}) + s1 = action.update({"chunk": "a", "response": "partial"}, s0) + assert s1.get("response") == "old" + s2 = action.update({"complete": True, "response": "final"}, s0) + assert s2.get("response") == "final" + + def test_streaming_custom_write_key_updates_state(self): + action = BedrockStreamingAction( + model_id="test-model", + input_mapper=lambda s: {"messages": []}, + reads=[], + writes=["answer"], + ) + s0 = State({}) + final = { + "chunk": "", + "complete": True, + "response": "done", + "usage": {}, + "stop_reason": None, + "answer": "done", + } + s1 = action.update(final, s0) + assert s1.get("answer") == "done" + + +class TestStateToPromptMapperProtocol: + """Test StateToPromptMapper Protocol exists.""" + + def test_protocol_exists(self): + """Verify StateToPromptMapper Protocol is defined.""" + assert StateToPromptMapper is not None + + +class TestIntegrationsLazyExports: + def test_unknown_attribute_raises(self): + with pytest.raises(AttributeError, match="has no attribute"): + _ = integrations.NotARealBedrockClass # noqa: SLF001