diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index d1aa475ec4..1ec66f6b8b 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -483,33 +483,65 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str: "aws_access_key_id", "aws_secret_access_key", ) +# LiteLLM expects the Bedrock bearer token as `api_key`; the UI uses the +# more descriptive `aws_bearer_token` and the resolver translates between them. +_BEDROCK_BEARER_TOKEN_FIELD: str = "aws_bearer_token" +_BEDROCK_LITELLM_BEARER_KWARG: str = "api_key" _BEDROCK_VALID_AUTH_TYPES: frozenset[str | None] = frozenset( - {None, "access_keys", "iam_role"} + {None, "access_keys", "iam_role", "bearer_token"} ) +def _drop_bedrock_access_keys(validated: dict[str, "Any"]) -> None: + for key in _BEDROCK_AWS_KEY_FIELDS: + validated.pop(key, None) + + +def _require_bedrock_access_keys(validated: dict[str, "Any"]) -> None: + for key in _BEDROCK_AWS_KEY_FIELDS: + value = validated.get(key) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{key} is required when auth_type is 'access_keys'.") + + +def _translate_bedrock_bearer_token(validated: dict[str, "Any"]) -> None: + """Move ``aws_bearer_token`` to ``api_key``; raise if missing or blank. + + Stripped before storing so surrounding whitespace doesn't reach the + ``Authorization`` header — AWS rejects mismatched bearer values with + an opaque 401. + """ + token = validated.pop(_BEDROCK_BEARER_TOKEN_FIELD, None) + if not isinstance(token, str) or not token.strip(): + raise ValueError( + f"{_BEDROCK_BEARER_TOKEN_FIELD} is required when " + "auth_type is 'bearer_token'." + ) + validated[_BEDROCK_LITELLM_BEARER_KWARG] = token.strip() + + def _resolve_bedrock_aws_credentials( adapter_metadata: dict[str, "Any"], validated: dict[str, "Any"], ) -> dict[str, "Any"]: """Apply auth_type semantics to the validated LiteLLM kwargs. - Three cases: - - ``auth_type == "iam_role"``: drop access keys unconditionally so a - previously-saved adapter switched into IAM Role mode does not leak - stale long-lived credentials. boto3's default credential chain - (IRSA / instance profile / env vars / AWS Profile) takes over. - - ``auth_type == "access_keys"`` (explicit): require non-blank values. - A blank submission must surface as a clear error rather than - silently fall through to the default chain (which would hide the - mistake and authenticate with whatever ambient creds the host has). - - ``auth_type is None`` (legacy adapters created before this field - existed): lenient strip of empty/missing keys. Preserves backwards - compatibility for stored configurations. + Each branch drops the credentials belonging to the *other* modes so a + re-saved adapter cannot leak stale long-lived secrets. Blank values + in an explicitly-selected mode raise rather than fall through to the + boto3 default chain, which would mask the misconfiguration. + + - ``iam_role``: drop access keys and bearer token; boto3's default + chain (IRSA / instance profile / env vars / AWS Profile) takes over. + - ``access_keys``: require non-blank access keys; drop bearer token. + - ``bearer_token``: require non-blank ``aws_bearer_token``; drop + access keys; translate the token to ``api_key`` for LiteLLM. + - ``None`` (legacy, pre-``auth_type``): lenient strip of empty access + keys; drop any bearer token (must be opted into explicitly). Raises: - ValueError: on unknown ``auth_type`` (typo / non-UI client) or on - blank credentials when ``auth_type == "access_keys"``. + ValueError: on unknown ``auth_type``, blank access keys in + ``access_keys`` mode, or blank token in ``bearer_token`` mode. """ auth_type = adapter_metadata.get("auth_type") if auth_type not in _BEDROCK_VALID_AUTH_TYPES: @@ -519,23 +551,27 @@ def _resolve_bedrock_aws_credentials( ) if auth_type == "iam_role": - for key in _BEDROCK_AWS_KEY_FIELDS: - validated.pop(key, None) + _drop_bedrock_access_keys(validated) + validated.pop(_BEDROCK_BEARER_TOKEN_FIELD, None) return validated if auth_type == "access_keys": - for key in _BEDROCK_AWS_KEY_FIELDS: - value = validated.get(key) - if not isinstance(value, str) or not value.strip(): - raise ValueError(f"{key} is required when auth_type is 'access_keys'.") + _require_bedrock_access_keys(validated) + validated.pop(_BEDROCK_BEARER_TOKEN_FIELD, None) + return validated + + if auth_type == "bearer_token": + _drop_bedrock_access_keys(validated) + _translate_bedrock_bearer_token(validated) return validated - # Legacy adapters with no auth_type: strip blanks silently to - # preserve the pre-PR behaviour where empty key fields fell through - # to boto3's default chain. + # No auth_type: strip blank access keys (boto3 chain takes over) and + # drop any bearer token — bearer auth must be opted into explicitly + # via auth_type='bearer_token' rather than promoted from this branch. for key in _BEDROCK_AWS_KEY_FIELDS: if not validated.get(key): validated.pop(key, None) + validated.pop(_BEDROCK_BEARER_TOKEN_FIELD, None) return validated @@ -544,6 +580,8 @@ class AWSBedrockLLMParameters(BaseChatCompletionParameters): aws_access_key_id: str | None = None aws_secret_access_key: str | None = None + # AWS_BEARER_TOKEN_BEDROCK; resolver translates to LiteLLM's `api_key`. + aws_bearer_token: str | None = None aws_region_name: str | None = None aws_profile_name: str | None = None # For AWS SSO authentication model_id: str | None = None # For Application Inference Profile (cost tracking) @@ -1033,6 +1071,8 @@ class AWSBedrockEmbeddingParameters(BaseEmbeddingParameters): # may be absent (IAM Role / Instance Profile mode). aws_access_key_id: str | None = None aws_secret_access_key: str | None = None + # AWS_BEARER_TOKEN_BEDROCK; resolver translates to LiteLLM's `api_key`. + aws_bearer_token: str | None = None aws_region_name: str | None @staticmethod diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.json b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.json index 200d8f5bfb..c953205617 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.json @@ -45,14 +45,16 @@ "title": "Authentication Type", "enum": [ "access_keys", - "iam_role" + "iam_role", + "bearer_token" ], "enumNames": [ "Access Keys", - "IAM Role / Instance Profile (on-prem AWS only)" + "IAM Role / Instance Profile (on-prem AWS only)", + "Bedrock API Key (Bearer Token)" ], "default": "access_keys", - "description": "Choose **Access Keys** for any deployment (provide your AWS Access Key ID and Secret). Choose **IAM Role / Instance Profile** only when Unstract is hosted on AWS infrastructure that already has ambient credentials — for example EKS pods with IRSA, ECS tasks with a task role, or EC2 instances with an instance profile. The on-prem option requires no further input; boto3 picks up the host's identity automatically." + "description": "Choose **Access Keys** for any deployment (provide your AWS Access Key ID and Secret). Choose **IAM Role / Instance Profile** only when Unstract is hosted on AWS infrastructure that already has ambient credentials — for example EKS pods with IRSA, ECS tasks with a task role, or EC2 instances with an instance profile. Choose **Bedrock API Key (Bearer Token)** to authenticate with an AWS Bedrock API key (`AWS_BEARER_TOKEN_BEDROCK`). The token is region-scoped — ensure the AWS Region below matches the key's region." } }, "dependencies": { @@ -91,6 +93,25 @@ ] } } + }, + { + "properties": { + "auth_type": { + "enum": [ + "bearer_token" + ] + }, + "aws_bearer_token": { + "type": "string", + "title": "Bedrock API Key", + "description": "Paste the Bedrock API key (bearer token) issued for your AWS account. See AWS docs on `AWS_BEARER_TOKEN_BEDROCK`.", + "format": "password", + "minLength": 1 + } + }, + "required": [ + "aws_bearer_token" + ] } ] } diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json index 03f23d1376..fcad4991e7 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json @@ -69,14 +69,16 @@ "title": "Authentication Type", "enum": [ "access_keys", - "iam_role" + "iam_role", + "bearer_token" ], "enumNames": [ "Access Keys", - "IAM Role / Instance Profile (on-prem AWS only)" + "IAM Role / Instance Profile (on-prem AWS only)", + "Bedrock API Key (Bearer Token)" ], "default": "access_keys", - "description": "Choose **Access Keys** for any deployment (provide your AWS Access Key ID and Secret). Choose **IAM Role / Instance Profile** only when Unstract is hosted on AWS infrastructure that already has ambient credentials — for example EKS pods with IRSA, ECS tasks with a task role, or EC2 instances with an instance profile. The on-prem option requires no further input; boto3 picks up the host's identity automatically." + "description": "Choose **Access Keys** for any deployment (provide your AWS Access Key ID and Secret). Choose **IAM Role / Instance Profile** only when Unstract is hosted on AWS infrastructure that already has ambient credentials — for example EKS pods with IRSA, ECS tasks with a task role, or EC2 instances with an instance profile. Choose **Bedrock API Key (Bearer Token)** to authenticate with an AWS Bedrock API key (`AWS_BEARER_TOKEN_BEDROCK`). The token is region-scoped — ensure the AWS Region below matches the key's region." } }, "dependencies": { @@ -115,6 +117,25 @@ ] } } + }, + { + "properties": { + "auth_type": { + "enum": [ + "bearer_token" + ] + }, + "aws_bearer_token": { + "type": "string", + "title": "Bedrock API Key", + "description": "Paste the Bedrock API key (bearer token) issued for your AWS account. See AWS docs on `AWS_BEARER_TOKEN_BEDROCK`.", + "format": "password", + "minLength": 1 + } + }, + "required": [ + "aws_bearer_token" + ] } ] } diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 14bc4b33ca..c64923dea5 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -10,7 +10,6 @@ # from litellm import get_supported_openai_params from litellm import get_max_tokens -from pydantic import ValidationError from unstract.sdk1.adapters.constants import Common from unstract.sdk1.adapters.llm1 import adapters from unstract.sdk1.constants import Common as SdkCommon @@ -190,7 +189,8 @@ def __init__( # noqa: C901 # if s not in self.kwargs: # logger.warning("Missing supported parameter for '%s': %s", # self.adapter.get_provider(), s) - except ValidationError as e: + except ValueError as e: + # `pydantic.ValidationError` subclasses `ValueError` — this catches both. raise SdkError("Invalid LLM adapter metadata: " + str(e)) from e self._system_prompt = system_prompt or self.SYSTEM_PROMPT diff --git a/unstract/sdk1/tests/test_bedrock_adapter.py b/unstract/sdk1/tests/test_bedrock_adapter.py index 4d4b7c1cd7..cfc728a379 100644 --- a/unstract/sdk1/tests/test_bedrock_adapter.py +++ b/unstract/sdk1/tests/test_bedrock_adapter.py @@ -124,6 +124,18 @@ def test_llm_unknown_auth_type_raises() -> None: ) +def test_llm_unknown_bearer_token_typo_raises() -> None: + with pytest.raises(ValueError, match="Unknown auth_type"): + AWSBedrockLLMParameters.validate( + { + "auth_type": "bearer-token", # typo: hyphen instead of underscore + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_bearer_token": "bedrock-key-abc", + } + ) + + def test_llm_other_params_preserved_through_strip() -> None: """Non-credential params survive the auth-type handling. @@ -149,6 +161,143 @@ def test_llm_other_params_preserved_through_strip() -> None: assert out["thinking"] == {"type": "enabled", "budget_tokens": 4096} +# ── LLM: bearer token (AWS_BEARER_TOKEN_BEDROCK) ───────────────────────────── + + +def test_llm_bearer_token_mode_translates_to_api_key() -> None: + """Bearer token is exposed to LiteLLM under its `api_key` kwarg.""" + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "bearer_token", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_bearer_token": "bedrock-key-abc", + } + ) + assert out["api_key"] == "bedrock-key-abc" + assert "aws_bearer_token" not in out + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + assert "auth_type" not in out + + +def test_llm_bearer_token_mode_drops_stale_access_keys() -> None: + """Switching a saved adapter to bearer mode must not leak old access keys.""" + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "bearer_token", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_access_key_id": "STALE_KEY", + "aws_secret_access_key": "STALE_SECRET", + "aws_bearer_token": "bedrock-key-abc", + } + ) + assert out["api_key"] == "bedrock-key-abc" + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + + +def test_llm_bearer_token_mode_blank_token_raises() -> None: + with pytest.raises(ValueError, match="aws_bearer_token is required"): + AWSBedrockLLMParameters.validate( + { + "auth_type": "bearer_token", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_bearer_token": "", + } + ) + + +def test_llm_bearer_token_mode_whitespace_token_raises() -> None: + with pytest.raises(ValueError, match="aws_bearer_token is required"): + AWSBedrockLLMParameters.validate( + { + "auth_type": "bearer_token", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_bearer_token": " ", + } + ) + + +def test_llm_bearer_token_mode_missing_token_raises() -> None: + """Field absent (not just blank) must surface the same clear error.""" + with pytest.raises(ValueError, match="aws_bearer_token is required"): + AWSBedrockLLMParameters.validate( + { + "auth_type": "bearer_token", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + } + ) + + +def test_llm_bearer_token_strips_surrounding_whitespace() -> None: + """Stray whitespace around a pasted key must not reach the header. + + Storing the unstripped value would produce + ``Authorization: Bearer `` which AWS rejects with an opaque 401. + """ + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "bearer_token", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_bearer_token": " bedrock-key-abc ", + } + ) + assert out["api_key"] == "bedrock-key-abc" + + +def test_llm_iam_role_drops_stale_bearer_token() -> None: + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "iam_role", + "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "region_name": "us-east-1", + "aws_bearer_token": "STALE_TOKEN", + } + ) + assert "aws_bearer_token" not in out + assert "api_key" not in out + + +def test_llm_access_keys_drops_stale_bearer_token() -> None: + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "access_keys", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "secret", + "aws_bearer_token": "STALE_TOKEN", + } + ) + assert "aws_bearer_token" not in out + assert "api_key" not in out + assert out["aws_access_key_id"] == "AKIAFAKE" + + +def test_llm_legacy_drops_stray_bearer_token() -> None: + """Legacy mode (no auth_type) must not stealth-promote a bearer token. + + Auto-translating would silently override env-injected + ``AWS_BEARER_TOKEN_BEDROCK`` or boto3 default-chain credentials with + no log line; opting into bearer auth must be explicit. + """ + out = AWSBedrockLLMParameters.validate( + { + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_bearer_token": "STRAY_TOKEN", + } + ) + assert "aws_bearer_token" not in out + assert "api_key" not in out + + # ── Embedding: same auth_type matrix ───────────────────────────────────────── @@ -233,6 +382,18 @@ def test_embedding_unknown_auth_type_raises() -> None: ) +def test_embedding_unknown_bearer_token_typo_raises() -> None: + with pytest.raises(ValueError, match="Unknown auth_type"): + AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "bearer-token", # typo: hyphen instead of underscore + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_bearer_token": "bedrock-key-abc", + } + ) + + def test_embedding_region_required_when_absent() -> None: """aws_region_name is still mandatory even though credentials are not.""" from pydantic import ValidationError @@ -244,3 +405,126 @@ def test_embedding_region_required_when_absent() -> None: "model": "amazon.titan-embed-text-v2:0", } ) + + +# ── Embedding: bearer token (AWS_BEARER_TOKEN_BEDROCK) ─────────────────────── + + +def test_embedding_bearer_token_mode_translates_to_api_key() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "bearer_token", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_bearer_token": "bedrock-key-abc", + } + ) + assert out["api_key"] == "bedrock-key-abc" + assert "aws_bearer_token" not in out + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + assert "auth_type" not in out + + +def test_embedding_bearer_token_mode_drops_stale_access_keys() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "bearer_token", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_access_key_id": "STALE_KEY", + "aws_secret_access_key": "STALE_SECRET", + "aws_bearer_token": "bedrock-key-abc", + } + ) + assert out["api_key"] == "bedrock-key-abc" + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + + +def test_embedding_bearer_token_mode_blank_token_raises() -> None: + with pytest.raises(ValueError, match="aws_bearer_token is required"): + AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "bearer_token", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_bearer_token": "", + } + ) + + +def test_embedding_bearer_token_mode_whitespace_token_raises() -> None: + with pytest.raises(ValueError, match="aws_bearer_token is required"): + AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "bearer_token", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_bearer_token": " ", + } + ) + + +def test_embedding_bearer_token_mode_missing_token_raises() -> None: + with pytest.raises(ValueError, match="aws_bearer_token is required"): + AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "bearer_token", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + } + ) + + +def test_embedding_bearer_token_strips_surrounding_whitespace() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "bearer_token", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_bearer_token": " bedrock-key-abc ", + } + ) + assert out["api_key"] == "bedrock-key-abc" + + +def test_embedding_iam_role_drops_stale_bearer_token() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "iam_role", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_bearer_token": "STALE_TOKEN", + } + ) + assert "aws_bearer_token" not in out + assert "api_key" not in out + + +def test_embedding_access_keys_drops_stale_bearer_token() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "access_keys", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "secret", + "aws_bearer_token": "STALE_TOKEN", + } + ) + assert "aws_bearer_token" not in out + assert "api_key" not in out + assert out["aws_access_key_id"] == "AKIAFAKE" + + +def test_embedding_legacy_drops_stray_bearer_token() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_bearer_token": "STRAY_TOKEN", + } + ) + assert "aws_bearer_token" not in out + assert "api_key" not in out