From e678074f824a1a73a1ea0d5792244303a8264ad1 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Mon, 18 May 2026 14:48:48 -0700 Subject: [PATCH] feat: Add bedrock provider package MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce launchdarkly-server-sdk-ai-bedrock, an AWS Bedrock provider for the LaunchDarkly AI SDK. The model runner wraps boto3 bedrock-runtime.converse(...) via asyncio.to_thread. The agent runner integrates the Strands Agents SDK, available through the optional [agents] extra. Agent-graph support is intentionally deferred — create_agent_graph raises NotImplementedError. The package mirrors the OpenAI provider layout: BedrockRunnerFactory, BedrockModelRunner, BedrockAgentRunner, helper module, and tests. Register the new package in the workspace pyproject.toml, root Makefile, release-please configuration, and CI / release workflows. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/ci.yml | 42 +++ .github/workflows/release-please.yml | 39 +++ .release-please-manifest.json | 1 + Makefile | 20 ++ .../ai-providers/server-ai-bedrock/Makefile | 29 ++ .../ai-providers/server-ai-bedrock/README.md | 185 +++++++++++++ .../server-ai-bedrock/pyproject.toml | 66 +++++ .../ai-providers/server-ai-bedrock/setup.cfg | 2 + .../src/ldai_bedrock/__init__.py | 21 ++ .../src/ldai_bedrock/bedrock_agent_runner.py | 202 ++++++++++++++ .../src/ldai_bedrock/bedrock_helper.py | 201 ++++++++++++++ .../src/ldai_bedrock/bedrock_model_runner.py | 205 ++++++++++++++ .../ldai_bedrock/bedrock_runner_factory.py | 125 +++++++++ .../server-ai-bedrock/tests/__init__.py | 1 + .../tests/test_agent_runner.py | 144 ++++++++++ .../server-ai-bedrock/tests/test_helper.py | 183 +++++++++++++ .../tests/test_model_runner.py | 254 ++++++++++++++++++ .../tests/test_runner_factory.py | 129 +++++++++ pyproject.toml | 2 + release-please-config.json | 13 +- 20 files changed, 1863 insertions(+), 1 deletion(-) create mode 100644 packages/ai-providers/server-ai-bedrock/Makefile create mode 100644 packages/ai-providers/server-ai-bedrock/README.md create mode 100644 packages/ai-providers/server-ai-bedrock/pyproject.toml create mode 100644 packages/ai-providers/server-ai-bedrock/setup.cfg create mode 100644 packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/__init__.py create mode 100644 packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_agent_runner.py create mode 100644 packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_helper.py create mode 100644 packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_model_runner.py create mode 100644 packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_runner_factory.py create mode 100644 packages/ai-providers/server-ai-bedrock/tests/__init__.py create mode 100644 packages/ai-providers/server-ai-bedrock/tests/test_agent_runner.py create mode 100644 packages/ai-providers/server-ai-bedrock/tests/test_helper.py create mode 100644 packages/ai-providers/server-ai-bedrock/tests/test_model_runner.py create mode 100644 packages/ai-providers/server-ai-bedrock/tests/test_runner_factory.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e98ef9f6..ab68aff9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -139,6 +139,48 @@ jobs: - name: Run tests run: make -C packages/ai-providers/server-ai-openai test + server-ai-bedrock-linux: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + + steps: + - uses: actions/checkout@v4 + + - uses: ./.github/actions/ci + with: + workspace_path: packages/ai-providers/server-ai-bedrock + python_version: ${{ matrix.python-version }} + + - uses: ./.github/actions/build + with: + workspace_path: packages/ai-providers/server-ai-bedrock + + server-ai-bedrock-windows: + runs-on: windows-latest + defaults: + run: + shell: powershell + + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up uv + uses: astral-sh/setup-uv@6ee6290f1cbc4156c0bdd66691b2c144ef8df19a # v7.4 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: make -C packages/ai-providers/server-ai-bedrock install + + - name: Run tests + run: make -C packages/ai-providers/server-ai-bedrock test + server-ai-optimization-linux: runs-on: ubuntu-latest strategy: diff --git a/.github/workflows/release-please.yml b/.github/workflows/release-please.yml index ecddce95..7c387cbd 100644 --- a/.github/workflows/release-please.yml +++ b/.github/workflows/release-please.yml @@ -27,6 +27,7 @@ on: - packages/sdk/server-ai - packages/ai-providers/server-ai-langchain - packages/ai-providers/server-ai-openai + - packages/ai-providers/server-ai-bedrock - packages/optimization dry_run: description: 'Is this a dry run. If so no package will be published.' @@ -47,6 +48,8 @@ jobs: package-server-ai-langchain-tag-name: ${{ steps.release.outputs['packages/ai-providers/server-ai-langchain--tag_name'] }} package-server-ai-openai-released: ${{ steps.release.outputs['packages/ai-providers/server-ai-openai--release_created'] }} package-server-ai-openai-tag-name: ${{ steps.release.outputs['packages/ai-providers/server-ai-openai--tag_name'] }} + package-server-ai-bedrock-released: ${{ steps.release.outputs['packages/ai-providers/server-ai-bedrock--release_created'] }} + package-server-ai-bedrock-tag-name: ${{ steps.release.outputs['packages/ai-providers/server-ai-bedrock--tag_name'] }} package-server-ai-optimization-released: ${{ steps.release.outputs['packages/optimization--release_created'] }} package-server-ai-optimization-tag-name: ${{ steps.release.outputs['packages/optimization--tag_name'] }} steps: @@ -193,6 +196,42 @@ jobs: password: ${{ env.PYPI_AUTH_TOKEN }} packages-dir: packages/ai-providers/server-ai-openai/dist/ + release-server-ai-bedrock: + runs-on: ubuntu-latest + needs: ['release-please'] + permissions: + id-token: write # Needed for OIDC to get release secrets from AWS. + attestations: write # Needed for actions/attest. + if: ${{ needs.release-please.outputs.package-server-ai-bedrock-released == 'true' }} + steps: + - uses: actions/checkout@v4 + + - uses: ./.github/actions/ci + with: + workspace_path: packages/ai-providers/server-ai-bedrock + + - uses: ./.github/actions/build + id: build + with: + workspace_path: packages/ai-providers/server-ai-bedrock + + - name: Attest build provenance + uses: actions/attest@v4 + with: + subject-path: 'packages/ai-providers/server-ai-bedrock/dist/*' + + - uses: launchdarkly/gh-actions/actions/release-secrets@release-secrets-v1.2.0 + name: 'Get PyPI token' + with: + aws_assume_role: ${{ vars.AWS_ROLE_ARN }} + ssm_parameter_pairs: '/production/common/releasing/pypi/token = PYPI_AUTH_TOKEN' + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 + with: + password: ${{ env.PYPI_AUTH_TOKEN }} + packages-dir: packages/ai-providers/server-ai-bedrock/dist/ + release-server-ai-optimization: runs-on: ubuntu-latest needs: ['release-please'] diff --git a/.release-please-manifest.json b/.release-please-manifest.json index c6e93245..63b319c9 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -2,5 +2,6 @@ "packages/sdk/server-ai": "1.0.1", "packages/ai-providers/server-ai-langchain": "0.8.0", "packages/ai-providers/server-ai-openai": "0.7.0", + "packages/ai-providers/server-ai-bedrock": "0.1.0", "packages/optimization": "0.1.0" } diff --git a/Makefile b/Makefile index 0d3e79fb..b549a8da 100644 --- a/Makefile +++ b/Makefile @@ -10,6 +10,7 @@ BUILDDIR = $(SOURCEDIR)/build SERVER_AI_PKG = packages/sdk/server-ai LANGCHAIN_PKG = packages/ai-providers/server-ai-langchain OPENAI_PKG = packages/ai-providers/server-ai-openai +BEDROCK_PKG = packages/ai-providers/server-ai-bedrock .PHONY: help help: #! Show this help message @@ -38,6 +39,10 @@ install-langchain: #! Install langchain provider package install-openai: #! Install openai provider package $(MAKE) -C $(OPENAI_PKG) install +.PHONY: install-bedrock +install-bedrock: #! Install bedrock provider package + $(MAKE) -C $(BEDROCK_PKG) install + # # Quality control checks # @@ -47,6 +52,7 @@ test: #! Run unit tests for all packages $(MAKE) test-server-ai $(MAKE) test-langchain $(MAKE) test-openai + $(MAKE) test-bedrock .PHONY: test-server-ai test-server-ai: #! Run unit tests for server-ai package @@ -60,11 +66,16 @@ test-langchain: #! Run unit tests for langchain provider package test-openai: #! Run unit tests for openai provider package $(MAKE) -C $(OPENAI_PKG) test +.PHONY: test-bedrock +test-bedrock: #! Run unit tests for bedrock provider package + $(MAKE) -C $(BEDROCK_PKG) test + .PHONY: lint lint: #! Run type analysis and linting checks for all packages $(MAKE) lint-server-ai $(MAKE) lint-langchain $(MAKE) lint-openai + $(MAKE) lint-bedrock .PHONY: lint-server-ai lint-server-ai: #! Run type analysis and linting checks for server-ai package @@ -78,6 +89,10 @@ lint-langchain: #! Run type analysis and linting checks for langchain provider p lint-openai: #! Run type analysis and linting checks for openai provider package $(MAKE) -C $(OPENAI_PKG) lint +.PHONY: lint-bedrock +lint-bedrock: #! Run type analysis and linting checks for bedrock provider package + $(MAKE) -C $(BEDROCK_PKG) lint + # # Build targets # @@ -87,6 +102,7 @@ build: #! Build all packages $(MAKE) build-server-ai $(MAKE) build-langchain $(MAKE) build-openai + $(MAKE) build-bedrock .PHONY: build-server-ai build-server-ai: #! Build server-ai package @@ -100,6 +116,10 @@ build-langchain: #! Build langchain provider package build-openai: #! Build openai provider package $(MAKE) -C $(OPENAI_PKG) build +.PHONY: build-bedrock +build-bedrock: #! Build bedrock provider package + $(MAKE) -C $(BEDROCK_PKG) build + # # Documentation generation # diff --git a/packages/ai-providers/server-ai-bedrock/Makefile b/packages/ai-providers/server-ai-bedrock/Makefile new file mode 100644 index 00000000..53b39a61 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/Makefile @@ -0,0 +1,29 @@ +PYTEST_FLAGS=-W error::SyntaxWarning + +.PHONY: help +help: #! Show this help message + @echo 'Usage: make [target] ... ' + @echo '' + @echo 'Targets:' + @grep -h -F '#!' $(MAKEFILE_LIST) | grep -v grep | sed 's/:.*#!/:/' | column -t -s":" + +.PHONY: install +install: #! Install package dependencies + uv sync --all-groups + +.PHONY: test +test: #! Run unit tests +test: install + uv run pytest $(PYTEST_FLAGS) + +.PHONY: lint +lint: #! Run type analysis and linting checks +lint: install + uv run mypy src/ldai_bedrock + uv run isort --check --atomic src/ldai_bedrock + uv run pycodestyle src/ldai_bedrock + +.PHONY: build +build: #! Build distribution files +build: install + uv build --out-dir dist diff --git a/packages/ai-providers/server-ai-bedrock/README.md b/packages/ai-providers/server-ai-bedrock/README.md new file mode 100644 index 00000000..dc798e1c --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/README.md @@ -0,0 +1,185 @@ +# LaunchDarkly AI SDK Amazon Bedrock Provider + +[![Actions Status](https://github.com/launchdarkly/python-server-sdk-ai/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/launchdarkly/python-server-sdk-ai/actions/workflows/ci.yml) + +[![PyPI](https://img.shields.io/pypi/v/launchdarkly-server-sdk-ai-bedrock.svg?maxAge=2592000)](https://pypi.org/project/launchdarkly-server-sdk-ai-bedrock/) +[![PyPI](https://img.shields.io/pypi/pyversions/launchdarkly-server-sdk-ai-bedrock.svg)](https://pypi.org/project/launchdarkly-server-sdk-ai-bedrock/) + +> [!CAUTION] +> This package is in pre-release and not subject to backwards compatibility +> guarantees. The API may change based on feedback. +> +> Pin to a specific minor version and review the [changelog](CHANGELOG.md) before upgrading. + +This package provides an Amazon Bedrock integration for the LaunchDarkly AI SDK. +Model completions use the Bedrock [Converse API](https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html) +via `boto3`. Single-turn agents are backed by the +[Strands Agents SDK](https://github.com/strands-agents/sdk-python), which is an +optional dependency. + +## Installation + +Install just the model runner: + +```bash +pip install launchdarkly-server-sdk-ai-bedrock +``` + +Add the agent runner by installing the `agents` extra: + +```bash +pip install "launchdarkly-server-sdk-ai-bedrock[agents]" +``` + +AWS credentials are read from the standard `boto3` credential chain +(environment variables, shared config, instance profile, etc.). + +## Quick Start + +```python +import asyncio +from ldclient import LDClient, Config, Context +from ldai import LDAIClient +from ldai.models import AICompletionConfigDefault, ModelConfig, ProviderConfig + +# Initialize LaunchDarkly client +ld_client = LDClient(Config("your-sdk-key")) +ai_client = LDAIClient(ld_client) + +context = Context.builder("user-123").build() + +async def main(): + # Create a ManagedModel backed by the Bedrock provider + model = await ai_client.create_model( + "ai-config-key", + context, + AICompletionConfigDefault( + enabled=True, + model=ModelConfig("anthropic.claude-3-5-sonnet-20240620-v1:0"), + provider=ProviderConfig("bedrock"), + ), + ) + + if model: + result = await model.run("Hello, how are you?") + print(result.content) + +asyncio.run(main()) +``` + +## Usage + +### Using `create_model` (recommended) + +The recommended entry point is `LDAIClient.create_model`, which evaluates a +LaunchDarkly AI config flag, selects the Bedrock runner automatically, and +returns a `ManagedModel` that wraps the runner: + +```python +model = await ai_client.create_model("ai-config-key", context) + +if model: + result = await model.run("What is feature flagging?") + print(result.content) +``` + +### Using the runner directly + +If you need to construct a runner manually (e.g. for testing), you can use +`BedrockRunnerFactory` from the `ldai_bedrock` package: + +```python +from ldai_bedrock import BedrockRunnerFactory + +# Uses the default boto3 credential chain. Pass ``region_name=`` to pin a region +# or ``client=`` to supply a pre-built bedrock-runtime client. +factory = BedrockRunnerFactory() +runner = factory.create_model(ai_config) + +result = await runner.run("Hello!") +print(result.content) +``` + +### Structured Output + +Pass a JSON schema dict as `output_type` to request structured output. The +Bedrock Converse API does not have a first-class JSON-schema mode, so the +schema is included as a system instruction and the response text is parsed as +JSON: + +```python +response_structure = { + "type": "object", + "properties": { + "sentiment": {"type": "string", "enum": ["positive", "negative", "neutral"]}, + "confidence": {"type": "number"}, + }, + "required": ["sentiment", "confidence"], +} + +result = await runner.run("How do customers feel about flags?", output_type=response_structure) +print(result.parsed) # {"sentiment": "positive", "confidence": 0.95} +``` + +### Tracking Metrics + +`ManagedModel.run()` automatically tracks metrics via the associated +`LDAIConfigTracker`. For manual tracking, use the tracker directly: + +```python +model = await ai_client.create_model("ai-config-key", context) + +if model: + result = await model.run("Explain feature flags.") + # Metrics are tracked automatically; access them via result.metrics + print(result.metrics.tokens) +``` + +### Async note + +`boto3` is synchronous. The Bedrock model runner therefore dispatches each +`converse` call to a worker thread (`asyncio.to_thread`) so the event loop is +not blocked while AWS responds. + +## Static Utility Methods + +The `ldai_bedrock` helper module provides several utility functions: + +### Converting Messages + +```python +from ldai.models import LDMessage +from ldai_bedrock import convert_messages_to_bedrock + +messages = [ + LDMessage(role="system", content="You are helpful."), + LDMessage(role="user", content="Hello!"), +] + +request_fragment = convert_messages_to_bedrock(messages) +# {"messages": [{"role": "user", "content": [{"text": "Hello!"}]}], +# "system": [{"text": "You are helpful."}]} +``` + +### Extracting Metrics + +```python +from ldai_bedrock import get_ai_metrics_from_response + +# After getting a response from bedrock-runtime.converse(...) +metrics = get_ai_metrics_from_response(response) +print(f"Success: {metrics.success}") +print(f"Tokens used: {metrics.tokens.total if metrics.tokens else 'N/A'}") +``` + +## Documentation + +For full documentation, please refer to the [LaunchDarkly AI SDK documentation](https://docs.launchdarkly.com/sdk/ai/python). + +## Contributing + +See [CONTRIBUTING.md](../../../CONTRIBUTING.md) in the repository root. + +## License + +Apache-2.0 diff --git a/packages/ai-providers/server-ai-bedrock/pyproject.toml b/packages/ai-providers/server-ai-bedrock/pyproject.toml new file mode 100644 index 00000000..096c2779 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/pyproject.toml @@ -0,0 +1,66 @@ +[project] +name = "launchdarkly-server-sdk-ai-bedrock" +version = "0.1.0" +description = "LaunchDarkly AI SDK Amazon Bedrock Provider" +authors = [{name = "LaunchDarkly", email = "dev@launchdarkly.com"}] +license = {text = "Apache-2.0"} +readme = "README.md" +requires-python = ">=3.10,<4" +classifiers = [ + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Software Development", + "Topic :: Software Development :: Libraries", +] +dependencies = [ + "launchdarkly-server-sdk-ai>=1.0.1", # x-release-please-version + "boto3>=1.34.0", +] + +[project.optional-dependencies] +agents = ["strands-agents>=0.1.0"] + +[project.urls] +Homepage = "https://docs.launchdarkly.com/sdk/ai/python" +Repository = "https://github.com/launchdarkly/python-server-sdk-ai" + +[dependency-groups] +dev = [ + "pytest>=2.8", + "pytest-cov>=2.4.0", + "pytest-asyncio>=0.21.0,<1.0.0", + "mypy==1.18.2", + "pycodestyle>=2.11.0", + "isort>=5.12.0", + "strands-agents>=0.1.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/ldai_bedrock"] + +[tool.mypy] +python_version = "3.10" +ignore_missing_imports = true +install_types = true +non_interactive = true + +[tool.isort] +profile = "black" +known_third_party = ["boto3", "ldai"] +sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"] + +[tool.pytest.ini_options] +addopts = ["-ra"] +testpaths = ["tests"] +asyncio_mode = "auto" diff --git a/packages/ai-providers/server-ai-bedrock/setup.cfg b/packages/ai-providers/server-ai-bedrock/setup.cfg new file mode 100644 index 00000000..6224f314 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/setup.cfg @@ -0,0 +1,2 @@ +[pycodestyle] +max-line-length = 120 diff --git a/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/__init__.py b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/__init__.py new file mode 100644 index 00000000..50fd062a --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/__init__.py @@ -0,0 +1,21 @@ +from ldai_bedrock.bedrock_agent_runner import BedrockAgentRunner +from ldai_bedrock.bedrock_helper import ( + convert_messages_to_bedrock, + convert_tools_to_bedrock, + get_ai_metrics_from_response, + get_ai_usage_from_response, + map_provider, +) +from ldai_bedrock.bedrock_model_runner import BedrockModelRunner +from ldai_bedrock.bedrock_runner_factory import BedrockRunnerFactory + +__all__ = [ + 'BedrockRunnerFactory', + 'BedrockModelRunner', + 'BedrockAgentRunner', + 'convert_messages_to_bedrock', + 'convert_tools_to_bedrock', + 'get_ai_metrics_from_response', + 'get_ai_usage_from_response', + 'map_provider', +] diff --git a/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_agent_runner.py b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_agent_runner.py new file mode 100644 index 00000000..b9c815c0 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_agent_runner.py @@ -0,0 +1,202 @@ +from typing import Any, Dict, List, Optional + +from ldai import log +from ldai.providers import RunnerResult, ToolRegistry +from ldai.providers.runner import Runner +from ldai.providers.types import LDAIMetrics +from ldai.tracker import TokenUsage + + +class BedrockAgentRunner(Runner): + """ + CAUTION: + This feature is experimental and should NOT be considered ready for production use. + It may change or be removed without notice and is not subject to backwards + compatibility guarantees. + + Runner implementation for a single agent running on Amazon Bedrock via + the Strands Agents SDK (``strands-agents``). + + Tool calling and the agentic loop are handled internally by the Strands + ``Agent`` class. Returned by ``BedrockRunnerFactory.create_agent(config, tools)``. + + Implements the unified :class:`~ldai.providers.runner.Runner` protocol. + Requires ``strands-agents`` to be installed: + + ``pip install launchdarkly-server-sdk-ai-bedrock[agents]`` + """ + + def __init__( + self, + model_id: str, + parameters: Dict[str, Any], + instructions: str, + tool_definitions: List[Dict[str, Any]], + tools: ToolRegistry, + ): + self._model_id = model_id + self._parameters = parameters + self._instructions = instructions + self._tool_definitions = tool_definitions + self._tools = tools + + async def run( + self, + input: str, + output_type: Optional[Dict[str, Any]] = None, + ) -> RunnerResult: + """ + Run the agent with the given input. + + Delegates to Strands' ``Agent.invoke_async``, which handles the + tool-calling loop internally. + + :param input: The user prompt string to the agent + :param output_type: Reserved for future structured output support; + currently ignored. + :return: :class:`RunnerResult` with ``content``, ``raw`` response, and + metrics including aggregated token usage and observed ``tool_calls``. + """ + try: + from strands import Agent + except ImportError: + log.warning( + "strands-agents is required for BedrockAgentRunner. " + "Install it with: pip install launchdarkly-server-sdk-ai-bedrock[agents]" + ) + return RunnerResult( + content='', + metrics=LDAIMetrics(success=False, tokens=None), + ) + + try: + agent_tools = self._build_agent_tools() + agent_kwargs: Dict[str, Any] = { + 'model': self._model_id, + 'tools': agent_tools, + } + if self._instructions: + agent_kwargs['system_prompt'] = self._instructions + + agent = Agent(**agent_kwargs) + result = await agent.invoke_async(str(input)) + + return RunnerResult( + content=_extract_strands_content(result), + metrics=LDAIMetrics( + success=True, + tokens=_extract_strands_usage(result), + tool_calls=_extract_strands_tool_calls(result) or None, + ), + raw=result, + ) + except Exception as error: + log.warning(f'Bedrock agent run failed: {error}') + return RunnerResult( + content='', + metrics=LDAIMetrics(success=False, tokens=None), + ) + + def _build_agent_tools(self) -> List[Any]: + """Build the tool list passed to ``Agent(tools=...)``. + + Plain callables are wrapped with Strands' ``@tool`` decorator; values + that are already Strands tool objects are returned unchanged. + """ + try: + from strands import tool as strands_tool + except ImportError as exc: + raise ImportError( + "strands-agents is required for agent tools. " + "Install it with: pip install launchdarkly-server-sdk-ai-bedrock[agents]" + ) from exc + + tools: List[Any] = [] + for td in self._tool_definitions: + if not isinstance(td, dict): + continue + name = td.get('name', '') + if not name: + continue + tool_fn = self._tools.get(name) + if tool_fn is None: + log.warning( + f"Tool '{name}' is defined in the AI config but was not found in " + "the tool registry; skipping." + ) + continue + if callable(tool_fn) and not _is_strands_tool(tool_fn): + tools.append(strands_tool(tool_fn)) + else: + tools.append(tool_fn) + return tools + + +def _is_strands_tool(value: Any) -> bool: + """Heuristically detect a value already decorated as a Strands tool.""" + return hasattr(value, 'tool_name') or hasattr(value, 'tool_spec') + + +def _extract_strands_content(result: Any) -> str: + """Best-effort extraction of the final text output from a Strands AgentResult.""" + if result is None: + return '' + # Strands' AgentResult is convertible to str and exposes ``message`` with + # Bedrock-style content blocks. Prefer the structured form so we never + # accidentally surface a debug repr. + message = getattr(result, 'message', None) + if isinstance(message, dict): + content = message.get('content') + if isinstance(content, list): + parts: List[str] = [] + for block in content: + if isinstance(block, dict): + text = block.get('text') + if isinstance(text, str) and text: + parts.append(text) + if parts: + return ''.join(parts) + return str(result) + + +def _extract_strands_usage(result: Any) -> Optional[TokenUsage]: + """Extract aggregate token usage from a Strands AgentResult, if available.""" + if result is None: + return None + metrics = getattr(result, 'metrics', None) + if metrics is None: + return None + usage = getattr(metrics, 'accumulated_usage', None) or getattr(metrics, 'usage', None) + if usage is None: + return None + if isinstance(usage, dict): + total = usage.get('totalTokens') or usage.get('total_tokens') or 0 + inp = usage.get('inputTokens') or usage.get('input_tokens') or 0 + out = usage.get('outputTokens') or usage.get('output_tokens') or 0 + else: + total = getattr(usage, 'totalTokens', None) or getattr(usage, 'total_tokens', 0) or 0 + inp = getattr(usage, 'inputTokens', None) or getattr(usage, 'input_tokens', 0) or 0 + out = getattr(usage, 'outputTokens', None) or getattr(usage, 'output_tokens', 0) or 0 + if not (total or inp or out): + return None + return TokenUsage(total=total, input=inp, output=out) + + +def _extract_strands_tool_calls(result: Any) -> List[str]: + """Return the names of any tool calls recorded on a Strands AgentResult.""" + if result is None: + return [] + metrics = getattr(result, 'metrics', None) + if metrics is None: + return [] + tool_metrics = getattr(metrics, 'tool_metrics', None) + if isinstance(tool_metrics, dict): + names: List[str] = [] + for name, entry in tool_metrics.items(): + call_count = getattr(entry, 'call_count', None) + if call_count is None and isinstance(entry, dict): + call_count = entry.get('call_count') + count = int(call_count or 0) if call_count is not None else 1 + names.extend([str(name)] * max(count, 1)) + return names + return [] diff --git a/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_helper.py b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_helper.py new file mode 100644 index 00000000..71437be4 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_helper.py @@ -0,0 +1,201 @@ +from typing import Any, Dict, List, Optional + +from ldai import LDMessage +from ldai.providers.types import LDAIMetrics +from ldai.tracker import TokenUsage + + +def map_provider(ld_provider_name: str) -> str: + """ + Map a LaunchDarkly provider name to its Bedrock-equivalent identifier. + + Amazon Bedrock is routed by ``modelId`` rather than a provider string, + so this is effectively an identity for ``bedrock``. The helper exists + to keep parity with the LangChain and OpenAI providers, which both + expose a similar mapping function. + + :param ld_provider_name: LaunchDarkly provider name + :return: The Bedrock provider identifier (always ``"bedrock"``) + """ + lowercased_name = ld_provider_name.lower() + if lowercased_name == 'bedrock' or lowercased_name.startswith('bedrock:'): + return 'bedrock' + return lowercased_name + + +def convert_messages_to_bedrock(messages: List[LDMessage]) -> Dict[str, Any]: + """ + Convert LaunchDarkly messages into the Bedrock Converse API shape. + + The Converse API splits the conversation into two top-level fields: + a ``system`` list aggregating all system prompts, and a ``messages`` + list containing user/assistant turns. Each non-system message has + a ``role`` plus a ``content`` array of typed content blocks + (``{"text": "..."}``). + + :param messages: List of LDMessage objects + :return: Dict with ``system`` and ``messages`` keys ready to pass to + ``bedrock-runtime.converse(...)``. The ``system`` key is absent + when no system prompts are present. + :raises ValueError: When a message has an unsupported role. + """ + system_blocks: List[Dict[str, Any]] = [] + bedrock_messages: List[Dict[str, Any]] = [] + for msg in messages: + if msg.role == 'system': + system_blocks.append({'text': msg.content}) + elif msg.role in ('user', 'assistant'): + bedrock_messages.append({ + 'role': msg.role, + 'content': [{'text': msg.content}], + }) + else: + raise ValueError(f'Unsupported message role: {msg.role}') + + result: Dict[str, Any] = {'messages': bedrock_messages} + if system_blocks: + result['system'] = system_blocks + return result + + +def convert_tools_to_bedrock(tool_definitions: List[Any]) -> Optional[Dict[str, Any]]: + """ + Convert LaunchDarkly tool definitions into Bedrock Converse ``toolConfig``. + + Bedrock expects ``toolConfig`` to be a dict containing a ``tools`` list + where each entry is wrapped in a ``toolSpec`` envelope. Each ``toolSpec`` + has ``name``, ``description``, and ``inputSchema.json`` keys. The LD + config layout matches this closely so most fields map through directly. + + :param tool_definitions: Tool definitions from the LD AI config + :return: A ``toolConfig`` dict ready to pass to ``converse(...)``, or + ``None`` when no usable tools are provided. + """ + tool_specs: List[Dict[str, Any]] = [] + for td in tool_definitions: + if not isinstance(td, dict): + continue + name = td.get('name') + if not name: + continue + spec: Dict[str, Any] = {'name': name} + description = td.get('description') + if description is not None: + spec['description'] = description + parameters = td.get('parameters') + if parameters is not None: + spec['inputSchema'] = {'json': parameters} + tool_specs.append({'toolSpec': spec}) + + if not tool_specs: + return None + return {'tools': tool_specs} + + +def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: + """ + Extract token usage from a Bedrock Converse response. + + :param response: Response dict returned by ``bedrock-runtime.converse(...)`` + :return: TokenUsage or None if usage data is unavailable + """ + if not isinstance(response, dict): + return None + usage = response.get('usage') + if not isinstance(usage, dict): + return None + total = usage.get('totalTokens') or 0 + inp = usage.get('inputTokens') or 0 + out = usage.get('outputTokens') or 0 + if not (total or inp or out): + return None + return TokenUsage(total=total, input=inp, output=out) + + +def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: + """ + Extract LaunchDarkly AI metrics from a Bedrock Converse response. + + ``success`` is derived from the HTTP status in ``ResponseMetadata``: a + 200 response is treated as success, anything else as failure. Token + usage, request duration, and observed tool-call names are populated + when present in the response. + + :param response: Response dict returned by ``bedrock-runtime.converse(...)`` + :return: LDAIMetrics with success, tokens, duration_ms, and tool_calls + """ + if not isinstance(response, dict): + return LDAIMetrics(success=False) + + status_code = (response.get('ResponseMetadata') or {}).get('HTTPStatusCode') + success = status_code == 200 + + tokens = get_ai_usage_from_response(response) + + duration_ms: Optional[int] = None + metrics_block = response.get('metrics') + if isinstance(metrics_block, dict): + latency = metrics_block.get('latencyMs') + if isinstance(latency, (int, float)): + duration_ms = int(latency) + + tool_calls = _extract_tool_calls(response) + + return LDAIMetrics( + success=success, + tokens=tokens, + tool_calls=tool_calls if tool_calls else None, + duration_ms=duration_ms, + ) + + +def _extract_tool_calls(response: Dict[str, Any]) -> List[str]: + """Return the names of any ``toolUse`` content blocks in the response.""" + output = response.get('output') + if not isinstance(output, dict): + return [] + message = output.get('message') + if not isinstance(message, dict): + return [] + content = message.get('content') + if not isinstance(content, list): + return [] + names: List[str] = [] + for block in content: + if not isinstance(block, dict): + continue + tool_use = block.get('toolUse') + if isinstance(tool_use, dict): + name = tool_use.get('name') + if isinstance(name, str): + names.append(name) + return names + + +def extract_content_from_response(response: Any) -> str: + """ + Pull the first textual content block out of a Bedrock Converse response. + + :param response: Response dict returned by ``bedrock-runtime.converse(...)`` + :return: The concatenated text from the first message's text blocks, or + an empty string when no text content is available. + """ + if not isinstance(response, dict): + return '' + output = response.get('output') + if not isinstance(output, dict): + return '' + message = output.get('message') + if not isinstance(message, dict): + return '' + content = message.get('content') + if not isinstance(content, list): + return '' + parts: List[str] = [] + for block in content: + if not isinstance(block, dict): + continue + text = block.get('text') + if isinstance(text, str) and text: + parts.append(text) + return ''.join(parts) diff --git a/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_model_runner.py b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_model_runner.py new file mode 100644 index 00000000..de9aa30b --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_model_runner.py @@ -0,0 +1,205 @@ +import asyncio +import json +from typing import Any, Dict, List, Optional + +from ldai import LDMessage, log +from ldai.providers.runner import Runner +from ldai.providers.types import LDAIMetrics, RunnerResult + +from ldai_bedrock.bedrock_helper import ( + convert_messages_to_bedrock, + convert_tools_to_bedrock, + extract_content_from_response, + get_ai_metrics_from_response, +) + + +class BedrockModelRunner(Runner): + """ + Runner implementation for Amazon Bedrock chat completions via the + Converse API. + + Holds a fully-configured ``bedrock-runtime`` client, model id, and + parameters. Returned by ``BedrockRunnerFactory.create_model(config)``. + + The Converse API is synchronous (boto3 has no async transport), so + each invocation dispatches to a worker thread via ``asyncio.to_thread`` + to avoid blocking the event loop. + + Implements the unified :class:`~ldai.providers.runner.Runner` protocol + via :meth:`run`. + """ + + # Inference parameters recognised by the Converse API's ``inferenceConfig``. + # All other parameters are forwarded as ``additionalModelRequestFields``. + _INFERENCE_KEYS = frozenset({ + 'maxTokens', + 'temperature', + 'topP', + 'stopSequences', + }) + + def __init__( + self, + client: Any, + model_id: str, + parameters: Dict[str, Any], + config_messages: Optional[List[LDMessage]] = None, + multi_turn: bool = True, + ): + self._client = client + self._model_id = model_id + self._parameters = parameters + self._history: List[LDMessage] = list(config_messages or []) + self._multi_turn = multi_turn + + async def run( + self, + input: str, + output_type: Optional[Dict[str, Any]] = None, + ) -> RunnerResult: + """ + Run the Bedrock model with the given input. + + :param input: A string prompt + :param output_type: Optional JSON schema dict requesting structured output. + When provided, ``parsed`` on the returned :class:`RunnerResult` is + populated with the parsed JSON document. Bedrock does not have a + native structured-output mode comparable to OpenAI's ``response_format``, + so the schema is appended as a system instruction and the model's + response is parsed as JSON. + :return: :class:`RunnerResult` containing ``content``, ``metrics``, + ``raw`` and (when ``output_type`` is set) ``parsed``. + """ + user_message = LDMessage(role='user', content=input) + messages = self._history + [user_message] + + if output_type is not None: + result = await self._run_structured(messages, output_type) + else: + result = await self._run_completion(messages) + + if result.metrics.success and result.content and self._multi_turn: + self._history.append(user_message) + self._history.append(LDMessage(role='assistant', content=result.content)) + + return result + + async def _run_completion(self, messages: List[LDMessage]) -> RunnerResult: + try: + response = await self._invoke_converse(messages) + + metrics = get_ai_metrics_from_response(response) + content = extract_content_from_response(response) + + if not content: + log.warning('Bedrock response has no content available') + return RunnerResult( + content='', + metrics=LDAIMetrics( + success=False, + tokens=metrics.tokens, + duration_ms=metrics.duration_ms, + ), + raw=response, + ) + + return RunnerResult(content=content, metrics=metrics, raw=response) + except Exception as error: + log.warning(f'Bedrock model invocation failed: {error}') + return RunnerResult( + content='', + metrics=LDAIMetrics(success=False, tokens=None), + ) + + async def _run_structured( + self, + messages: List[LDMessage], + output_type: Dict[str, Any], + ) -> RunnerResult: + # Bedrock has no first-class JSON-schema response mode comparable to + # OpenAI's ``response_format``. Inject the schema as a system + # instruction and parse the response text as JSON. + schema_instruction = LDMessage( + role='system', + content=( + 'Respond with a JSON document that conforms exactly to the ' + 'following JSON schema. Output JSON only, with no surrounding ' + f'prose or code fences.\nSchema: {json.dumps(output_type)}' + ), + ) + augmented = [schema_instruction, *messages] + + try: + response = await self._invoke_converse(augmented) + + metrics = get_ai_metrics_from_response(response) + content = extract_content_from_response(response) + + if not content: + log.warning('Bedrock structured response has no content available') + return RunnerResult( + content='', + metrics=LDAIMetrics( + success=False, + tokens=metrics.tokens, + duration_ms=metrics.duration_ms, + ), + raw=response, + ) + + try: + parsed = json.loads(content) + return RunnerResult( + content=content, + metrics=metrics, + raw=response, + parsed=parsed, + ) + except json.JSONDecodeError as parse_error: + log.warning(f'Bedrock structured response contains invalid JSON: {parse_error}') + return RunnerResult( + content=content, + metrics=LDAIMetrics( + success=False, + tokens=metrics.tokens, + duration_ms=metrics.duration_ms, + ), + raw=response, + ) + except Exception as error: + log.warning(f'Bedrock structured model invocation failed: {error}') + return RunnerResult( + content='', + metrics=LDAIMetrics(success=False, tokens=None), + ) + + async def _invoke_converse(self, messages: List[LDMessage]) -> Dict[str, Any]: + request = self._build_request(messages) + return await asyncio.to_thread(self._client.converse, **request) + + def _build_request(self, messages: List[LDMessage]) -> Dict[str, Any]: + """Assemble the kwargs passed to ``bedrock-runtime.converse(...)``.""" + request: Dict[str, Any] = {'modelId': self._model_id} + request.update(convert_messages_to_bedrock(messages)) + + inference_config: Dict[str, Any] = {} + additional_fields: Dict[str, Any] = {} + for key, value in self._parameters.items(): + if key == 'tools': + continue + if key in self._INFERENCE_KEYS: + inference_config[key] = value + else: + additional_fields[key] = value + + if inference_config: + request['inferenceConfig'] = inference_config + if additional_fields: + request['additionalModelRequestFields'] = additional_fields + + tool_config = convert_tools_to_bedrock(self._parameters.get('tools') or []) + if tool_config: + request['toolConfig'] = tool_config + + return request diff --git a/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_runner_factory.py b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_runner_factory.py new file mode 100644 index 00000000..698e2f0a --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/src/ldai_bedrock/bedrock_runner_factory.py @@ -0,0 +1,125 @@ +from typing import TYPE_CHECKING, Any, Optional + +from ldai.models import AIConfigKind +from ldai.providers import AIProvider, ToolRegistry + +from ldai_bedrock.bedrock_model_runner import BedrockModelRunner + +if TYPE_CHECKING: + from ldai_bedrock.bedrock_agent_runner import BedrockAgentRunner + + +class BedrockRunnerFactory(AIProvider): + """Amazon Bedrock ``AIProvider`` implementation for the LaunchDarkly AI SDK.""" + + def __init__(self, client: Any = None, region_name: Optional[str] = None): + """ + Initialize the Bedrock connector. + + :param client: A pre-built ``bedrock-runtime`` boto3 client. When + omitted, one is constructed lazily on first use using the standard + AWS credential chain. + :param region_name: Optional region name passed through to ``boto3.client`` + when the connector creates the client itself. Ignored when + ``client`` is supplied. + """ + self._client = client + self._region_name = region_name + + def _get_client(self) -> Any: + if self._client is None: + import boto3 + kwargs: dict = {} + if self._region_name: + kwargs['region_name'] = self._region_name + self._client = boto3.client('bedrock-runtime', **kwargs) + return self._client + + def get_client(self) -> Any: + """Return the underlying ``bedrock-runtime`` boto3 client.""" + return self._get_client() + + def _extract_model_config(self, config: AIConfigKind) -> tuple: + """ + Extract model id and parameters from an AI config. + + :param config: The LaunchDarkly AI configuration + :return: Tuple of (model_id, parameters) + """ + config_dict = config.to_dict() + model_dict = config_dict.get('model') or {} + return model_dict.get('name', ''), model_dict.get('parameters') or {} + + def create_model(self, config: AIConfigKind, multi_turn: bool = True) -> BedrockModelRunner: + """ + Create a configured ``BedrockModelRunner`` for the given AI config. + + Reuses the underlying ``bedrock-runtime`` client so connection pooling + is preserved across model invocations. + + :param config: The LaunchDarkly AI configuration. ``model.name`` is + interpreted as a Bedrock ``modelId`` (e.g. + ``anthropic.claude-3-5-sonnet-20240620-v1:0``). + :param multi_turn: When ``True`` (the default) the runner accumulates + successful exchanges into its conversation history. Pass ``False`` + to keep history fixed at the configured baseline across ``run()`` + calls. + :return: BedrockModelRunner ready to invoke the model + """ + model_id, parameters = self._extract_model_config(config) + parameters = dict(parameters) + config_messages = list(getattr(config, 'messages', None) or []) + return BedrockModelRunner( + self._get_client(), + model_id, + parameters, + config_messages, + multi_turn=multi_turn, + ) + + def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'BedrockAgentRunner': + """ + CAUTION: + This feature is experimental and should NOT be considered ready for production use. + It may change or be removed without notice and is not subject to backwards + compatibility guarantees. + + Create a configured ``BedrockAgentRunner`` for the given AI agent config. + + The agent runner is backed by the Strands Agents SDK (``strands-agents``), + which is an optional dependency of this package. Install it with:: + + pip install launchdarkly-server-sdk-ai-bedrock[agents] + + :param config: The LaunchDarkly AI agent configuration + :param tools: ToolRegistry mapping tool names to callables + :return: BedrockAgentRunner ready to run the agent + """ + from ldai_bedrock.bedrock_agent_runner import BedrockAgentRunner + + model_id, base_parameters = self._extract_model_config(config) + parameters = dict(base_parameters) + tool_definitions = parameters.pop('tools', []) or [] + instructions = (config.instructions or '') if hasattr(config, 'instructions') else '' + + return BedrockAgentRunner( + model_id, + parameters, + instructions, + tool_definitions, + tools or {}, + ) + + def create_agent_graph( + self, + graph_def: Any, + tools: ToolRegistry, + ) -> Any: + """ + Agent graph execution is not yet supported by the Bedrock provider. + + :raises NotImplementedError: Always. + """ + raise NotImplementedError( + 'Agent graph is not yet supported by the bedrock provider' + ) diff --git a/packages/ai-providers/server-ai-bedrock/tests/__init__.py b/packages/ai-providers/server-ai-bedrock/tests/__init__.py new file mode 100644 index 00000000..9f1d5176 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for LaunchDarkly AI SDK Amazon Bedrock Provider.""" diff --git a/packages/ai-providers/server-ai-bedrock/tests/test_agent_runner.py b/packages/ai-providers/server-ai-bedrock/tests/test_agent_runner.py new file mode 100644 index 00000000..5e808f82 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/tests/test_agent_runner.py @@ -0,0 +1,144 @@ +"""Tests for BedrockAgentRunner.""" + +import sys +import types + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from ldai_bedrock import BedrockAgentRunner + + +def _make_strands_mock(invoke_async_mock): + """ + Construct a stand-in ``strands`` module with the surface BedrockAgentRunner + uses (``Agent``, ``tool`` decorator). + """ + strands_mock = types.ModuleType('strands') + + agent_instances = [] + + class FakeAgent: + def __init__(self, **kwargs): + self.kwargs = kwargs + agent_instances.append(self) + + async def invoke_async(self, prompt): + return await invoke_async_mock(prompt, self.kwargs) + + strands_mock.Agent = FakeAgent # type: ignore[attr-defined] + strands_mock.tool = lambda fn: fn # type: ignore[attr-defined] + strands_mock._agent_instances = agent_instances # type: ignore[attr-defined] + return strands_mock + + +def _make_strands_result(content: str, *, total: int = 15, inp: int = 10, out: int = 5, + tool_metrics: dict = None): + """Mock the public surface of a strands AgentResult we read in the runner.""" + metrics_obj = MagicMock() + metrics_obj.accumulated_usage = { + 'totalTokens': total, + 'inputTokens': inp, + 'outputTokens': out, + } + metrics_obj.tool_metrics = tool_metrics or {} + + result = MagicMock() + result.metrics = metrics_obj + result.message = { + 'role': 'assistant', + 'content': [{'text': content}], + } + return result + + +class TestBedrockAgentRunner: + """Tests for BedrockAgentRunner.run.""" + + @pytest.mark.asyncio + async def test_runs_agent_and_returns_result_with_no_tool_calls(self): + invoke_mock = AsyncMock(return_value=_make_strands_result('The answer is 42.')) + strands_mock = _make_strands_mock(invoke_mock) + + runner = BedrockAgentRunner( + 'anthropic.claude-3-5-sonnet-20240620-v1:0', + {}, + 'You are helpful.', + [], + {}, + ) + with patch.dict(sys.modules, {'strands': strands_mock}): + result = await runner.run('What is the answer?') + + assert result.content == 'The answer is 42.' + assert result.metrics.success is True + assert result.metrics.tokens is not None + assert result.metrics.tokens.total == 15 + assert result.metrics.tool_calls is None + + @pytest.mark.asyncio + async def test_records_tool_calls_from_metrics(self): + tool_metrics = { + 'get_weather': MagicMock(call_count=2), + } + invoke_mock = AsyncMock(return_value=_make_strands_result( + 'It is sunny.', + tool_metrics=tool_metrics, + )) + strands_mock = _make_strands_mock(invoke_mock) + + runner = BedrockAgentRunner( + 'model-id', + {}, + 'You are helpful.', + [{'name': 'get_weather', 'description': 'Get weather', 'parameters': {}}], + {'get_weather': lambda loc: 'sunny'}, + ) + with patch.dict(sys.modules, {'strands': strands_mock}): + result = await runner.run('Weather in Paris?') + + assert result.content == 'It is sunny.' + assert result.metrics.success is True + assert result.metrics.tool_calls == ['get_weather', 'get_weather'] + + @pytest.mark.asyncio + async def test_passes_instructions_and_tools_to_agent(self): + invoke_mock = AsyncMock(return_value=_make_strands_result('done')) + strands_mock = _make_strands_mock(invoke_mock) + + runner = BedrockAgentRunner( + 'model-id', + {}, + 'Be terse.', + [{'name': 'lookup', 'description': 'Look up', 'parameters': {}}], + {'lookup': lambda q: 'answer'}, + ) + with patch.dict(sys.modules, {'strands': strands_mock}): + await runner.run('go') + + assert len(strands_mock._agent_instances) == 1 + kwargs = strands_mock._agent_instances[0].kwargs + assert kwargs['model'] == 'model-id' + assert kwargs['system_prompt'] == 'Be terse.' + assert len(kwargs['tools']) == 1 + + @pytest.mark.asyncio + async def test_returns_failure_when_exception_thrown(self): + invoke_mock = AsyncMock(side_effect=Exception('boom')) + strands_mock = _make_strands_mock(invoke_mock) + + runner = BedrockAgentRunner('model-id', {}, '', [], {}) + with patch.dict(sys.modules, {'strands': strands_mock}): + result = await runner.run('Hello') + + assert result.content == '' + assert result.metrics.success is False + + @pytest.mark.asyncio + async def test_returns_failure_when_strands_not_installed(self): + runner = BedrockAgentRunner('model-id', {}, '', [], {}) + with patch.dict(sys.modules, {'strands': None}): + result = await runner.run('Hello') + + assert result.content == '' + assert result.metrics.success is False diff --git a/packages/ai-providers/server-ai-bedrock/tests/test_helper.py b/packages/ai-providers/server-ai-bedrock/tests/test_helper.py new file mode 100644 index 00000000..48d2d68d --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/tests/test_helper.py @@ -0,0 +1,183 @@ +"""Tests for ldai_bedrock.bedrock_helper.""" + +import pytest +from ldai import LDMessage + +from ldai_bedrock import ( + convert_messages_to_bedrock, + convert_tools_to_bedrock, + get_ai_metrics_from_response, + get_ai_usage_from_response, + map_provider, +) + + +class TestMapProvider: + """Tests for map_provider.""" + + def test_maps_bedrock_to_bedrock(self): + assert map_provider('bedrock') == 'bedrock' + + def test_maps_bedrock_with_family_prefix(self): + assert map_provider('Bedrock:Anthropic') == 'bedrock' + assert map_provider('bedrock:amazon') == 'bedrock' + + def test_passes_through_unrelated_provider(self): + assert map_provider('openai') == 'openai' + + +class TestConvertMessagesToBedrock: + """Tests for convert_messages_to_bedrock.""" + + def test_splits_system_and_conversation_messages(self): + result = convert_messages_to_bedrock([ + LDMessage(role='system', content='You are helpful.'), + LDMessage(role='user', content='Hello!'), + LDMessage(role='assistant', content='Hi there.'), + ]) + assert result == { + 'system': [{'text': 'You are helpful.'}], + 'messages': [ + {'role': 'user', 'content': [{'text': 'Hello!'}]}, + {'role': 'assistant', 'content': [{'text': 'Hi there.'}]}, + ], + } + + def test_omits_system_key_when_no_system_messages(self): + result = convert_messages_to_bedrock([ + LDMessage(role='user', content='Hello!'), + ]) + assert 'system' not in result + assert result['messages'] == [ + {'role': 'user', 'content': [{'text': 'Hello!'}]}, + ] + + def test_aggregates_multiple_system_messages(self): + result = convert_messages_to_bedrock([ + LDMessage(role='system', content='First.'), + LDMessage(role='system', content='Second.'), + LDMessage(role='user', content='Go.'), + ]) + assert result['system'] == [ + {'text': 'First.'}, + {'text': 'Second.'}, + ] + + def test_raises_on_unsupported_role(self): + with pytest.raises(ValueError): + convert_messages_to_bedrock([ + LDMessage(role='function', content='Nope.'), + ]) + + +class TestConvertToolsToBedrock: + """Tests for convert_tools_to_bedrock.""" + + def test_wraps_tools_in_tool_spec_envelopes(self): + result = convert_tools_to_bedrock([ + { + 'name': 'get_weather', + 'description': 'Look up the weather for a location.', + 'parameters': { + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + 'required': ['location'], + }, + }, + ]) + assert result == { + 'tools': [ + { + 'toolSpec': { + 'name': 'get_weather', + 'description': 'Look up the weather for a location.', + 'inputSchema': { + 'json': { + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + 'required': ['location'], + }, + }, + }, + }, + ], + } + + def test_returns_none_for_empty_tools(self): + assert convert_tools_to_bedrock([]) is None + + def test_skips_non_dict_and_unnamed_entries(self): + assert convert_tools_to_bedrock(['not-a-dict', {'description': 'no name'}]) is None + + +class TestGetAIUsageFromResponse: + """Tests for get_ai_usage_from_response.""" + + def test_extracts_token_counts(self): + usage = get_ai_usage_from_response({ + 'usage': {'totalTokens': 100, 'inputTokens': 60, 'outputTokens': 40}, + }) + assert usage is not None + assert usage.total == 100 + assert usage.input == 60 + assert usage.output == 40 + + def test_returns_none_when_usage_absent(self): + assert get_ai_usage_from_response({}) is None + + def test_returns_none_when_all_counts_zero(self): + assert get_ai_usage_from_response({ + 'usage': {'totalTokens': 0, 'inputTokens': 0, 'outputTokens': 0}, + }) is None + + +class TestGetAIMetricsFromResponse: + """Tests for get_ai_metrics_from_response.""" + + def test_success_when_status_200(self): + result = get_ai_metrics_from_response({ + 'ResponseMetadata': {'HTTPStatusCode': 200}, + 'usage': {'totalTokens': 5, 'inputTokens': 3, 'outputTokens': 2}, + 'metrics': {'latencyMs': 42}, + }) + assert result.success is True + assert result.tokens is not None + assert result.tokens.total == 5 + assert result.duration_ms == 42 + assert result.tool_calls is None + + def test_failure_when_status_not_200(self): + result = get_ai_metrics_from_response({ + 'ResponseMetadata': {'HTTPStatusCode': 500}, + 'usage': {'totalTokens': 5, 'inputTokens': 3, 'outputTokens': 2}, + 'metrics': {'latencyMs': 42}, + }) + assert result.success is False + assert result.tokens is not None + assert result.duration_ms == 42 + + def test_records_observed_tool_call_names(self): + result = get_ai_metrics_from_response({ + 'ResponseMetadata': {'HTTPStatusCode': 200}, + 'output': { + 'message': { + 'role': 'assistant', + 'content': [ + {'toolUse': {'name': 'get_weather', 'input': {}}}, + {'text': 'Looking that up.'}, + ], + }, + }, + }) + assert result.success is True + assert result.tool_calls == ['get_weather'] + + def test_handles_missing_metadata(self): + result = get_ai_metrics_from_response({}) + assert result.success is False + assert result.tokens is None + assert result.duration_ms is None + + def test_returns_failure_for_non_dict(self): + result = get_ai_metrics_from_response(None) + assert result.success is False diff --git a/packages/ai-providers/server-ai-bedrock/tests/test_model_runner.py b/packages/ai-providers/server-ai-bedrock/tests/test_model_runner.py new file mode 100644 index 00000000..67084d38 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/tests/test_model_runner.py @@ -0,0 +1,254 @@ +"""Tests for BedrockModelRunner.""" + +import pytest +from unittest.mock import MagicMock + +from ldai_bedrock import BedrockModelRunner + + +def _make_response(text: str = 'Hello!', *, status: int = 200, + tokens: tuple = (25, 10, 15), latency_ms: int = 50, + tool_use: dict = None) -> dict: + """Build a synthetic Bedrock Converse response dict.""" + total, inp, out = tokens + content_blocks = [] + if text is not None: + content_blocks.append({'text': text}) + if tool_use is not None: + content_blocks.append({'toolUse': tool_use}) + return { + 'ResponseMetadata': {'HTTPStatusCode': status}, + 'output': { + 'message': { + 'role': 'assistant', + 'content': content_blocks, + }, + }, + 'usage': { + 'totalTokens': total, + 'inputTokens': inp, + 'outputTokens': out, + }, + 'metrics': {'latencyMs': latency_ms}, + 'stopReason': 'end_turn', + } + + +class TestRunCompletion: + """Tests for the run() method (chat-completion path).""" + + @pytest.fixture + def mock_client(self): + return MagicMock() + + @pytest.mark.asyncio + async def test_invokes_converse_and_returns_response(self, mock_client): + mock_client.converse = MagicMock(return_value=_make_response('Hello! How can I help you today?')) + + runner = BedrockModelRunner( + mock_client, + 'anthropic.claude-3-5-sonnet-20240620-v1:0', + {'temperature': 0.5, 'maxTokens': 200}, + ) + result = await runner.run('Hello!') + + # The model invocation should request inferenceConfig with the + # Converse-recognised keys, the user prompt as a Bedrock message, and + # the expected modelId. + mock_client.converse.assert_called_once() + kwargs = mock_client.converse.call_args.kwargs + assert kwargs['modelId'] == 'anthropic.claude-3-5-sonnet-20240620-v1:0' + assert kwargs['messages'] == [ + {'role': 'user', 'content': [{'text': 'Hello!'}]}, + ] + assert kwargs['inferenceConfig'] == {'temperature': 0.5, 'maxTokens': 200} + assert 'toolConfig' not in kwargs + + assert result.content == 'Hello! How can I help you today?' + assert result.metrics.success is True + assert result.metrics.tokens is not None + assert result.metrics.tokens.total == 25 + assert result.metrics.duration_ms == 50 + + @pytest.mark.asyncio + async def test_returns_unsuccessful_response_when_no_content(self, mock_client): + mock_client.converse = MagicMock(return_value={ + 'ResponseMetadata': {'HTTPStatusCode': 200}, + 'output': {'message': {'role': 'assistant', 'content': []}}, + 'usage': {'totalTokens': 0, 'inputTokens': 0, 'outputTokens': 0}, + }) + + runner = BedrockModelRunner(mock_client, 'model-id', {}) + result = await runner.run('Hello!') + + assert result.content == '' + assert result.metrics.success is False + + @pytest.mark.asyncio + async def test_returns_unsuccessful_response_when_exception_thrown(self, mock_client): + mock_client.converse = MagicMock(side_effect=Exception('AWS error')) + + runner = BedrockModelRunner(mock_client, 'model-id', {}) + result = await runner.run('Hello!') + + assert result.content == '' + assert result.metrics.success is False + + @pytest.mark.asyncio + async def test_passes_through_tool_config_when_tools_present(self, mock_client): + mock_client.converse = MagicMock(return_value=_make_response()) + + runner = BedrockModelRunner( + mock_client, + 'model-id', + { + 'temperature': 0.2, + 'tools': [ + {'name': 'get_time', 'description': 'Return the current time.', 'parameters': {'type': 'object'}}, + ], + }, + ) + await runner.run('What time is it?') + + kwargs = mock_client.converse.call_args.kwargs + assert kwargs['toolConfig'] == { + 'tools': [ + { + 'toolSpec': { + 'name': 'get_time', + 'description': 'Return the current time.', + 'inputSchema': {'json': {'type': 'object'}}, + }, + }, + ], + } + # ``tools`` should not leak into inferenceConfig or additional fields. + assert kwargs['inferenceConfig'] == {'temperature': 0.2} + assert 'additionalModelRequestFields' not in kwargs + + @pytest.mark.asyncio + async def test_routes_unknown_parameters_to_additional_fields(self, mock_client): + mock_client.converse = MagicMock(return_value=_make_response()) + + runner = BedrockModelRunner( + mock_client, + 'model-id', + {'temperature': 0.1, 'top_k': 50, 'reasoning_budget': 1024}, + ) + await runner.run('Question') + + kwargs = mock_client.converse.call_args.kwargs + assert kwargs['inferenceConfig'] == {'temperature': 0.1} + assert kwargs['additionalModelRequestFields'] == { + 'top_k': 50, + 'reasoning_budget': 1024, + } + + @pytest.mark.asyncio + async def test_accumulates_history_across_successful_calls(self, mock_client): + mock_client.converse = MagicMock(side_effect=[ + _make_response('First response'), + _make_response('Second response'), + ]) + + runner = BedrockModelRunner(mock_client, 'model-id', {}) + await runner.run('First question') + await runner.run('Second question') + + second_call_messages = mock_client.converse.call_args_list[1].kwargs['messages'] + assert second_call_messages == [ + {'role': 'user', 'content': [{'text': 'First question'}]}, + {'role': 'assistant', 'content': [{'text': 'First response'}]}, + {'role': 'user', 'content': [{'text': 'Second question'}]}, + ] + + @pytest.mark.asyncio + async def test_multi_turn_false_does_not_accumulate_history(self, mock_client): + mock_client.converse = MagicMock(side_effect=[ + _make_response('First response'), + _make_response('Second response'), + ]) + + runner = BedrockModelRunner(mock_client, 'model-id', {}, multi_turn=False) + baseline_len = len(runner._history) + + await runner.run('First question') + assert len(runner._history) == baseline_len + + await runner.run('Second question') + assert len(runner._history) == baseline_len + + second_call_messages = mock_client.converse.call_args_list[1].kwargs['messages'] + assert second_call_messages == [ + {'role': 'user', 'content': [{'text': 'Second question'}]}, + ] + + @pytest.mark.asyncio + async def test_does_not_accumulate_history_on_failed_call(self, mock_client): + mock_client.converse = MagicMock(side_effect=[Exception('boom'), _make_response('Recovery')]) + + runner = BedrockModelRunner(mock_client, 'model-id', {}) + await runner.run('Hello!') + await runner.run('Try again') + + second_call_messages = mock_client.converse.call_args_list[1].kwargs['messages'] + assert second_call_messages == [ + {'role': 'user', 'content': [{'text': 'Try again'}]}, + ] + + +class TestRunStructured: + """Tests for the structured-output path of run().""" + + @pytest.fixture + def mock_client(self): + return MagicMock() + + @pytest.mark.asyncio + async def test_returns_parsed_json_when_schema_supplied(self, mock_client): + mock_client.converse = MagicMock( + return_value=_make_response('{"name": "John", "age": 30}'), + ) + + runner = BedrockModelRunner(mock_client, 'model-id', {}) + schema = { + 'type': 'object', + 'properties': { + 'name': {'type': 'string'}, + 'age': {'type': 'number'}, + }, + 'required': ['name', 'age'], + } + + result = await runner.run('Tell me about a person', output_type=schema) + + assert result.parsed == {'name': 'John', 'age': 30} + assert result.content == '{"name": "John", "age": 30}' + assert result.metrics.success is True + + # The schema should be injected as a system-prompt block to steer Bedrock. + kwargs = mock_client.converse.call_args.kwargs + assert 'system' in kwargs + assert any('Schema:' in block['text'] for block in kwargs['system']) + + @pytest.mark.asyncio + async def test_returns_unsuccessful_when_response_is_not_valid_json(self, mock_client): + mock_client.converse = MagicMock(return_value=_make_response('not json')) + + runner = BedrockModelRunner(mock_client, 'model-id', {}) + result = await runner.run('Question', output_type={'type': 'object'}) + + assert result.parsed is None + assert result.content == 'not json' + assert result.metrics.success is False + + @pytest.mark.asyncio + async def test_returns_unsuccessful_when_exception_thrown(self, mock_client): + mock_client.converse = MagicMock(side_effect=Exception('AWS error')) + + runner = BedrockModelRunner(mock_client, 'model-id', {}) + result = await runner.run('Question', output_type={'type': 'object'}) + + assert result.parsed is None + assert result.content == '' + assert result.metrics.success is False diff --git a/packages/ai-providers/server-ai-bedrock/tests/test_runner_factory.py b/packages/ai-providers/server-ai-bedrock/tests/test_runner_factory.py new file mode 100644 index 00000000..21a85993 --- /dev/null +++ b/packages/ai-providers/server-ai-bedrock/tests/test_runner_factory.py @@ -0,0 +1,129 @@ +"""Tests for BedrockRunnerFactory.""" + +import pytest +from unittest.mock import MagicMock, patch + +from ldai_bedrock import BedrockAgentRunner, BedrockModelRunner, BedrockRunnerFactory + + +class TestGetClient: + """Tests for get_client / client lazy-construction.""" + + def test_returns_supplied_client(self): + mock_client = MagicMock() + factory = BedrockRunnerFactory(client=mock_client) + assert factory.get_client() is mock_client + + def test_constructs_client_lazily_when_omitted(self): + with patch('boto3.client') as mock_boto: + built_client = MagicMock() + mock_boto.return_value = built_client + + factory = BedrockRunnerFactory(region_name='us-east-1') + client = factory.get_client() + + mock_boto.assert_called_once_with( + 'bedrock-runtime', + region_name='us-east-1', + ) + assert client is built_client + + def test_reuses_client_on_second_call(self): + with patch('boto3.client') as mock_boto: + mock_boto.return_value = MagicMock() + factory = BedrockRunnerFactory() + first = factory.get_client() + second = factory.get_client() + assert first is second + assert mock_boto.call_count == 1 + + +class TestCreateModel: + """Tests for BedrockRunnerFactory.create_model.""" + + def test_creates_runner_with_correct_model_and_parameters(self): + mock_ai_config = MagicMock() + mock_ai_config.messages = None + mock_ai_config.to_dict.return_value = { + 'model': { + 'name': 'anthropic.claude-3-5-sonnet-20240620-v1:0', + 'parameters': { + 'temperature': 0.7, + 'maxTokens': 1024, + }, + }, + 'provider': {'name': 'bedrock'}, + } + + mock_client = MagicMock() + result = BedrockRunnerFactory(client=mock_client).create_model(mock_ai_config) + + assert isinstance(result, BedrockModelRunner) + assert result._model_id == 'anthropic.claude-3-5-sonnet-20240620-v1:0' + assert result._parameters == {'temperature': 0.7, 'maxTokens': 1024} + + def test_handles_missing_model_config(self): + mock_ai_config = MagicMock() + mock_ai_config.messages = None + mock_ai_config.to_dict.return_value = {} + + mock_client = MagicMock() + result = BedrockRunnerFactory(client=mock_client).create_model(mock_ai_config) + + assert isinstance(result, BedrockModelRunner) + assert result._model_id == '' + assert result._parameters == {} + + +class TestCreateAgent: + """Tests for BedrockRunnerFactory.create_agent.""" + + def test_creates_agent_runner_with_instructions_and_tool_definitions(self): + mock_ai_config = MagicMock() + mock_ai_config.instructions = 'You are a helpful assistant.' + mock_ai_config.to_dict.return_value = { + 'model': { + 'name': 'anthropic.claude-3-5-sonnet-20240620-v1:0', + 'parameters': { + 'temperature': 0.7, + 'tools': [ + {'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}, + ], + }, + }, + } + + mock_client = MagicMock() + factory = BedrockRunnerFactory(client=mock_client) + result = factory.create_agent(mock_ai_config, {'get-weather': lambda loc: 'sunny'}) + + assert isinstance(result, BedrockAgentRunner) + assert result._model_id == 'anthropic.claude-3-5-sonnet-20240620-v1:0' + assert result._instructions == 'You are a helpful assistant.' + # The ``tools`` parameter is consumed and surfaced via ``_tool_definitions``. + assert result._parameters == {'temperature': 0.7} + assert len(result._tool_definitions) == 1 + assert result._tool_definitions[0]['name'] == 'get-weather' + + def test_creates_agent_runner_with_no_tools(self): + mock_ai_config = MagicMock() + mock_ai_config.instructions = 'You are a helpful assistant.' + mock_ai_config.to_dict.return_value = { + 'model': {'name': 'anthropic.claude-3-5-sonnet-20240620-v1:0', 'parameters': {}}, + } + + mock_client = MagicMock() + factory = BedrockRunnerFactory(client=mock_client) + result = factory.create_agent(mock_ai_config, {}) + + assert isinstance(result, BedrockAgentRunner) + assert result._tool_definitions == [] + + +class TestCreateAgentGraph: + """Tests for BedrockRunnerFactory.create_agent_graph.""" + + def test_raises_not_implemented_error(self): + factory = BedrockRunnerFactory(client=MagicMock()) + with pytest.raises(NotImplementedError): + factory.create_agent_graph(MagicMock(), {}) diff --git a/pyproject.toml b/pyproject.toml index e74c8f42..f3c88ba1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ # - packages/sdk/server-ai (launchdarkly-server-sdk-ai) # - packages/ai-providers/server-ai-openai (launchdarkly-server-sdk-ai-openai) # - packages/ai-providers/server-ai-langchain (launchdarkly-server-sdk-ai-langchain) +# - packages/ai-providers/server-ai-bedrock (launchdarkly-server-sdk-ai-bedrock) # - packages/optimization (launchdarkly-server-sdk-ai-optimization) # # For development, use the package-specific pyproject.toml files in each member. @@ -16,6 +17,7 @@ members = [ "packages/sdk/server-ai", "packages/ai-providers/server-ai-openai", "packages/ai-providers/server-ai-langchain", + "packages/ai-providers/server-ai-bedrock", "packages/optimization", ] # Resolve launchdarkly-server-sdk-ai from the local workspace member rather diff --git a/release-please-config.json b/release-please-config.json index 44fb6bef..cd6991fa 100644 --- a/release-please-config.json +++ b/release-please-config.json @@ -10,7 +10,8 @@ "src/ldai/__init__.py", "PROVENANCE.md", "/packages/ai-providers/server-ai-langchain/pyproject.toml", - "/packages/ai-providers/server-ai-openai/pyproject.toml" + "/packages/ai-providers/server-ai-openai/pyproject.toml", + "/packages/ai-providers/server-ai-bedrock/pyproject.toml" ], "component": "launchdarkly-server-sdk-ai" }, @@ -34,6 +35,16 @@ ], "component": "launchdarkly-server-sdk-ai-openai" }, + "packages/ai-providers/server-ai-bedrock": { + "release-type": "python", + "versioning": "default", + "bump-minor-pre-major": true, + "include-v-in-tag": false, + "extra-files": [ + "src/ldai_bedrock/__init__.py" + ], + "component": "launchdarkly-server-sdk-ai-bedrock" + }, "packages/optimization": { "release-type": "python", "versioning": "default",