mirror of
https://github.com/andrewyng/openworker.git
synced 2026-09-13 07:40:18 +00:00
Support Bedrock API keys (bearer auth)
New optional field: paste the console-generated key, no CLI/IAM setup needed. Takes precedence over SigV4 credentials, matching boto3; live-tested on Converse.
This commit is contained in:
@@ -12,8 +12,12 @@ An id with no family segment falls back to Converse as-is — Converse serves ev
|
|||||||
model (including Claude, minus the native extras), so a raw model id pasted without the
|
model (including Claude, minus the native extras), so a raw model id pasted without the
|
||||||
add-model dropdown still works.
|
add-model dropdown still works.
|
||||||
|
|
||||||
Credentials resolve explicit → named profile → ambient (env / `~/.aws` default / role),
|
Credentials, in order:
|
||||||
matching what AWS CLI users expect; `aws sso login` sessions arrive via the named profile.
|
1. A **Bedrock API key** (bearer token from the Bedrock console — the no-CLI path). When
|
||||||
|
present (field or `AWS_BEARER_TOKEN_BEDROCK` env) it WINS over SigV4 credentials, the
|
||||||
|
same precedence boto3 applies; mixing both makes `AnthropicBedrock` raise outright.
|
||||||
|
2. Explicit IAM keys → named profile (covers `aws sso login`) → ambient chain.
|
||||||
|
|
||||||
boto3 is a lazy import (packaged via the `bedrock` extra) and returns PLAIN DICTS — every
|
boto3 is a lazy import (packaged via the `bedrock` extra) and returns PLAIN DICTS — every
|
||||||
response/stream mapping here is dict-shaped, unlike the attribute objects other SDKs return.
|
response/stream mapping here is dict-shaped, unlike the attribute objects other SDKs return.
|
||||||
"""
|
"""
|
||||||
@@ -22,6 +26,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import base64
|
import base64
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
import re
|
import re
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional
|
||||||
|
|
||||||
@@ -260,6 +265,7 @@ class _BedrockConverseClient(ProviderClient):
|
|||||||
client: Any = None,
|
client: Any = None,
|
||||||
*,
|
*,
|
||||||
region: Optional[str] = None,
|
region: Optional[str] = None,
|
||||||
|
bedrock_api_key: Optional[str] = None,
|
||||||
profile_name: Optional[str] = None,
|
profile_name: Optional[str] = None,
|
||||||
access_key_id: Optional[str] = None,
|
access_key_id: Optional[str] = None,
|
||||||
secret_access_key: Optional[str] = None,
|
secret_access_key: Optional[str] = None,
|
||||||
@@ -267,6 +273,7 @@ class _BedrockConverseClient(ProviderClient):
|
|||||||
):
|
):
|
||||||
self._client = client # tests inject a dict-returning fake
|
self._client = client # tests inject a dict-returning fake
|
||||||
self._region = region
|
self._region = region
|
||||||
|
self._bedrock_api_key = bedrock_api_key
|
||||||
self._session_kwargs = _session_kwargs(
|
self._session_kwargs = _session_kwargs(
|
||||||
profile_name, access_key_id, secret_access_key, session_token
|
profile_name, access_key_id, secret_access_key, session_token
|
||||||
)
|
)
|
||||||
@@ -280,6 +287,11 @@ class _BedrockConverseClient(ProviderClient):
|
|||||||
"AWS Bedrock support needs the boto3 package — "
|
"AWS Bedrock support needs the boto3 package — "
|
||||||
"install with `pip install 'openworker[bedrock]'`."
|
"install with `pip install 'openworker[bedrock]'`."
|
||||||
) from exc
|
) from exc
|
||||||
|
# boto3 has no per-client bearer parameter — it only reads the env var, and
|
||||||
|
# prefers bearer auth for Bedrock whenever it's set. The sidecar process is
|
||||||
|
# ours, so publishing the configured key there is the supported path.
|
||||||
|
if self._bedrock_api_key:
|
||||||
|
os.environ["AWS_BEARER_TOKEN_BEDROCK"] = self._bedrock_api_key
|
||||||
session = boto3.session.Session(**self._session_kwargs)
|
session = boto3.session.Session(**self._session_kwargs)
|
||||||
self._client = session.client("bedrock-runtime", region_name=self._region)
|
self._client = session.client("bedrock-runtime", region_name=self._region)
|
||||||
return self._client
|
return self._client
|
||||||
@@ -437,6 +449,7 @@ class BedrockProvider(ProviderClient):
|
|||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
region: Optional[str] = None,
|
region: Optional[str] = None,
|
||||||
|
bedrock_api_key: Optional[str] = None,
|
||||||
profile_name: Optional[str] = None,
|
profile_name: Optional[str] = None,
|
||||||
access_key_id: Optional[str] = None,
|
access_key_id: Optional[str] = None,
|
||||||
secret_access_key: Optional[str] = None,
|
secret_access_key: Optional[str] = None,
|
||||||
@@ -445,6 +458,7 @@ class BedrockProvider(ProviderClient):
|
|||||||
converse_client: Optional[ProviderClient] = None,
|
converse_client: Optional[ProviderClient] = None,
|
||||||
):
|
):
|
||||||
self._region = region
|
self._region = region
|
||||||
|
self._bedrock_api_key = bedrock_api_key
|
||||||
self._profile_name = profile_name
|
self._profile_name = profile_name
|
||||||
self._access_key_id = access_key_id
|
self._access_key_id = access_key_id
|
||||||
self._secret_access_key = secret_access_key
|
self._secret_access_key = secret_access_key
|
||||||
@@ -472,18 +486,26 @@ class BedrockProvider(ProviderClient):
|
|||||||
if family == "claude":
|
if family == "claude":
|
||||||
from anthropic import AnthropicBedrock
|
from anthropic import AnthropicBedrock
|
||||||
|
|
||||||
client = AnthropicProvider(
|
# A Bedrock API key (field or ambient env) takes the bearer path and
|
||||||
client=AnthropicBedrock(
|
# EXCLUDES the SigV4 params — AnthropicBedrock raises on a mix.
|
||||||
|
bearer = self._bedrock_api_key or os.environ.get(
|
||||||
|
"AWS_BEARER_TOKEN_BEDROCK"
|
||||||
|
)
|
||||||
|
if bearer:
|
||||||
|
sdk = AnthropicBedrock(api_key=bearer, aws_region=self._region)
|
||||||
|
else:
|
||||||
|
sdk = AnthropicBedrock(
|
||||||
aws_region=self._region,
|
aws_region=self._region,
|
||||||
aws_profile=self._profile_name,
|
aws_profile=self._profile_name,
|
||||||
aws_access_key=self._access_key_id,
|
aws_access_key=self._access_key_id,
|
||||||
aws_secret_key=self._secret_access_key,
|
aws_secret_key=self._secret_access_key,
|
||||||
aws_session_token=self._session_token,
|
aws_session_token=self._session_token,
|
||||||
)
|
)
|
||||||
)
|
client = AnthropicProvider(client=sdk)
|
||||||
else:
|
else:
|
||||||
client = _BedrockConverseClient(
|
client = _BedrockConverseClient(
|
||||||
region=self._region,
|
region=self._region,
|
||||||
|
bedrock_api_key=self._bedrock_api_key,
|
||||||
profile_name=self._profile_name,
|
profile_name=self._profile_name,
|
||||||
access_key_id=self._access_key_id,
|
access_key_id=self._access_key_id,
|
||||||
secret_access_key=self._secret_access_key,
|
secret_access_key=self._secret_access_key,
|
||||||
|
|||||||
@@ -140,6 +140,7 @@ def _build_bedrock(profile: dict[str, Any], secrets: Any) -> ProviderClient:
|
|||||||
|
|
||||||
return BedrockProvider(
|
return BedrockProvider(
|
||||||
region=get("region"),
|
region=get("region"),
|
||||||
|
bedrock_api_key=get("bedrock_api_key"),
|
||||||
profile_name=get("aws_profile"),
|
profile_name=get("aws_profile"),
|
||||||
access_key_id=get("aws_access_key_id"),
|
access_key_id=get("aws_access_key_id"),
|
||||||
secret_access_key=get("aws_secret_access_key"),
|
secret_access_key=get("aws_secret_access_key"),
|
||||||
@@ -299,6 +300,15 @@ DESCRIPTORS: list[ProviderDescriptor] = [
|
|||||||
placeholder="us-east-1",
|
placeholder="us-east-1",
|
||||||
help="The region your Bedrock model access is enabled in.",
|
help="The region your Bedrock model access is enabled in.",
|
||||||
),
|
),
|
||||||
|
ProviderField(
|
||||||
|
"bedrock_api_key",
|
||||||
|
"Bedrock API key (optional)",
|
||||||
|
secret=True,
|
||||||
|
required=False,
|
||||||
|
placeholder="ABSK…",
|
||||||
|
help="The easiest way in: generate one on the Bedrock console — no AWS "
|
||||||
|
"CLI or IAM setup needed. Takes precedence over the fields below.",
|
||||||
|
),
|
||||||
ProviderField(
|
ProviderField(
|
||||||
"aws_profile",
|
"aws_profile",
|
||||||
"AWS profile (optional)",
|
"AWS profile (optional)",
|
||||||
@@ -551,6 +561,10 @@ def _verify_bedrock(fields: dict[str, Any], timeout: float) -> dict[str, Any]:
|
|||||||
"error": "boto3 is not installed — `pip install 'openworker[bedrock]'`.",
|
"error": "boto3 is not installed — `pip install 'openworker[bedrock]'`.",
|
||||||
}
|
}
|
||||||
try:
|
try:
|
||||||
|
# A Bedrock API key rides the env var (boto3's only bearer channel) and then
|
||||||
|
# wins over any SigV4 credentials, matching the provider's own precedence.
|
||||||
|
if get("bedrock_api_key"):
|
||||||
|
os.environ["AWS_BEARER_TOKEN_BEDROCK"] = get("bedrock_api_key")
|
||||||
session = boto3.session.Session(
|
session = boto3.session.Session(
|
||||||
**_session_kwargs(
|
**_session_kwargs(
|
||||||
get("aws_profile"),
|
get("aws_profile"),
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ export const KEY_HELP: Record<string, { url: string; label: string }> = {
|
|||||||
openai: { url: "https://platform.openai.com/api-keys", label: "platform.openai.com" },
|
openai: { url: "https://platform.openai.com/api-keys", label: "platform.openai.com" },
|
||||||
gemini: { url: "https://aistudio.google.com/apikey", label: "aistudio.google.com" },
|
gemini: { url: "https://aistudio.google.com/apikey", label: "aistudio.google.com" },
|
||||||
openrouter: { url: "https://openrouter.ai/keys", label: "openrouter.ai" },
|
openrouter: { url: "https://openrouter.ai/keys", label: "openrouter.ai" },
|
||||||
|
bedrock: { url: "https://console.aws.amazon.com/bedrock/home#/api-keys", label: "the AWS Bedrock console" },
|
||||||
fireworks: { url: "https://fireworks.ai/account/api-keys", label: "fireworks.ai" },
|
fireworks: { url: "https://fireworks.ai/account/api-keys", label: "fireworks.ai" },
|
||||||
together: { url: "https://api.together.xyz/settings/api-keys", label: "together.xyz" },
|
together: { url: "https://api.together.xyz/settings/api-keys", label: "together.xyz" },
|
||||||
zai: { url: "https://z.ai/manage-apikey/apikey-list", label: "z.ai" },
|
zai: { url: "https://z.ai/manage-apikey/apikey-list", label: "z.ai" },
|
||||||
|
|||||||
@@ -288,17 +288,52 @@ def test_family_dispatch():
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
def test_claude_family_builds_native_anthropic_over_bedrock():
|
def test_claude_family_builds_native_anthropic_over_bedrock(monkeypatch):
|
||||||
from anthropic import AnthropicBedrock
|
from anthropic import AnthropicBedrock
|
||||||
|
|
||||||
from coworker.providers import AnthropicProvider
|
from coworker.providers import AnthropicProvider
|
||||||
|
|
||||||
|
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||||
p = BedrockProvider(region="us-east-1", profile_name="work")
|
p = BedrockProvider(region="us-east-1", profile_name="work")
|
||||||
sub = p._family_client("claude")
|
sub = p._family_client("claude")
|
||||||
assert isinstance(sub, AnthropicProvider)
|
assert isinstance(sub, AnthropicProvider)
|
||||||
assert isinstance(sub._client, AnthropicBedrock)
|
assert isinstance(sub._client, AnthropicBedrock)
|
||||||
|
|
||||||
|
|
||||||
|
def test_claude_family_prefers_bedrock_api_key_over_sigv4(monkeypatch):
|
||||||
|
"""A Bedrock API key must take the bearer path WITHOUT the SigV4 params —
|
||||||
|
AnthropicBedrock raises outright when both are passed."""
|
||||||
|
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||||
|
p = BedrockProvider(
|
||||||
|
region="us-east-1", bedrock_api_key="ABSKtest", profile_name="work"
|
||||||
|
)
|
||||||
|
sub = p._family_client("claude")
|
||||||
|
assert sub._client.api_key == "ABSKtest"
|
||||||
|
assert sub._client.aws_profile is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_converse_client_publishes_api_key_as_bearer_env(monkeypatch):
|
||||||
|
import os
|
||||||
|
|
||||||
|
import boto3
|
||||||
|
|
||||||
|
from coworker.providers.bedrock_provider import _BedrockConverseClient
|
||||||
|
|
||||||
|
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||||
|
|
||||||
|
class _FakeSession:
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def client(self, service, **kwargs):
|
||||||
|
return object()
|
||||||
|
|
||||||
|
monkeypatch.setattr(boto3.session, "Session", _FakeSession)
|
||||||
|
client = _BedrockConverseClient(region="us-east-1", bedrock_api_key="ABSKtest")
|
||||||
|
client._ensure_client()
|
||||||
|
assert os.environ["AWS_BEARER_TOKEN_BEDROCK"] == "ABSKtest"
|
||||||
|
|
||||||
|
|
||||||
# -- capabilities / matrix ------------------------------------------------------------
|
# -- capabilities / matrix ------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
@@ -334,6 +369,7 @@ def test_bedrock_descriptor_and_builder():
|
|||||||
keys = [f.key for f in d.fields]
|
keys = [f.key for f in d.fields]
|
||||||
assert keys == [
|
assert keys == [
|
||||||
"region",
|
"region",
|
||||||
|
"bedrock_api_key",
|
||||||
"aws_profile",
|
"aws_profile",
|
||||||
"aws_access_key_id",
|
"aws_access_key_id",
|
||||||
"aws_secret_access_key",
|
"aws_secret_access_key",
|
||||||
@@ -341,7 +377,7 @@ def test_bedrock_descriptor_and_builder():
|
|||||||
]
|
]
|
||||||
assert [f.key for f in d.fields if f.required] == ["region"]
|
assert [f.key for f in d.fields if f.required] == ["region"]
|
||||||
secret = {f.key for f in d.fields if f.secret}
|
secret = {f.key for f in d.fields if f.secret}
|
||||||
assert secret == {"aws_secret_access_key", "aws_session_token"}
|
assert secret == {"bedrock_api_key", "aws_secret_access_key", "aws_session_token"}
|
||||||
# Recommended model is curated in the matrix (set_provider's auto-add depends on it).
|
# Recommended model is curated in the matrix (set_provider's auto-add depends on it).
|
||||||
from coworker.providers.matrix import models_for_provider
|
from coworker.providers.matrix import models_for_provider
|
||||||
|
|
||||||
@@ -416,6 +452,22 @@ def test_verify_bedrock_ok(monkeypatch):
|
|||||||
assert captured["client"]["region_name"] == "us-east-1"
|
assert captured["client"]["region_name"] == "us-east-1"
|
||||||
|
|
||||||
|
|
||||||
|
def test_verify_bedrock_api_key_rides_the_bearer_env(monkeypatch):
|
||||||
|
import os
|
||||||
|
|
||||||
|
from coworker.providers.registry import verify_provider_key
|
||||||
|
|
||||||
|
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||||
|
captured: dict = {}
|
||||||
|
_patch_session(monkeypatch, _FakeBedrockControl(), captured)
|
||||||
|
out = verify_provider_key(
|
||||||
|
"bedrock", fields={"region": "us-east-1", "bedrock_api_key": "ABSKtest"}
|
||||||
|
)
|
||||||
|
assert out == {"ok": True}
|
||||||
|
assert captured["session"] == {} # bearer only — no SigV4 session kwargs
|
||||||
|
assert os.environ["AWS_BEARER_TOKEN_BEDROCK"] == "ABSKtest"
|
||||||
|
|
||||||
|
|
||||||
def test_verify_bedrock_maps_client_errors(monkeypatch):
|
def test_verify_bedrock_maps_client_errors(monkeypatch):
|
||||||
from botocore.exceptions import ClientError
|
from botocore.exceptions import ClientError
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user