diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index f7d98f0fa1..bac149b9ff 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -474,6 +474,71 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str: return f"vertex_ai/{model}" +# AWS Bedrock auth helpers: shared by LLM and Embedding param classes. +# `auth_type` is a UI-only selector (Access Keys vs IAM Role / Instance +# Profile) that drives form rendering. The backend translates the user's +# choice into actual credential handling here so that both validate() +# methods stay symmetric and a single bug fix applies to both paths. +_BEDROCK_AWS_KEY_FIELDS: tuple[str, ...] = ( + "aws_access_key_id", + "aws_secret_access_key", +) +_BEDROCK_VALID_AUTH_TYPES: frozenset[str | None] = frozenset( + {None, "access_keys", "iam_role"} +) + + +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. + + Raises: + ValueError: on unknown ``auth_type`` (typo / non-UI client) or on + blank credentials when ``auth_type == "access_keys"``. + """ + auth_type = adapter_metadata.get("auth_type") + if auth_type not in _BEDROCK_VALID_AUTH_TYPES: + raise ValueError( + f"Unknown auth_type {auth_type!r}; expected one of " + f"{sorted(t for t in _BEDROCK_VALID_AUTH_TYPES if t)!r} or absent." + ) + + if auth_type == "iam_role": + for key in _BEDROCK_AWS_KEY_FIELDS: + validated.pop(key, 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'.") + 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. + for key in _BEDROCK_AWS_KEY_FIELDS: + if not validated.get(key): + validated.pop(key, None) + return validated + + class AWSBedrockLLMParameters(BaseChatCompletionParameters): """See https://docs.litellm.ai/docs/providers/bedrock.""" @@ -525,11 +590,12 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: result_metadata["thinking"] = thinking_config result_metadata["temperature"] = 1 - # Create validation metadata excluding control fields + # Create validation metadata excluding control fields. `auth_type` is + # a UI-only selector that drives form rendering; LiteLLM never sees it. validation_metadata = { k: v for k, v in result_metadata.items() - if k not in ("enable_thinking", "budget_tokens", "thinking") + if k not in ("enable_thinking", "budget_tokens", "thinking", "auth_type") } validated = AWSBedrockLLMParameters(**validation_metadata).model_dump() @@ -538,7 +604,11 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: if enable_thinking and "thinking" in result_metadata: validated["thinking"] = result_metadata["thinking"] - return validated + # Apply Bedrock auth semantics: IAM Role mode drops keys, Access + # Keys mode requires non-blank values, legacy (no auth_type) is + # lenient. Reads auth_type from result_metadata since validation_ + # metadata strips it before Pydantic. + return _resolve_bedrock_aws_credentials(result_metadata, validated) @staticmethod def validate_model(adapter_metadata: dict[str, "Any"]) -> str: @@ -959,8 +1029,10 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str: class AWSBedrockEmbeddingParameters(BaseEmbeddingParameters): """See https://docs.litellm.ai/docs/providers/bedrock.""" - aws_access_key_id: str | None - aws_secret_access_key: str | None + # Region is still mandatory — credentials are the only fields that + # may be absent (IAM Role / Instance Profile mode). + aws_access_key_id: str | None = None + aws_secret_access_key: str | None = None aws_region_name: str | None @staticmethod @@ -973,7 +1045,16 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: ): adapter_metadata["aws_region_name"] = adapter_metadata["region_name"] - return AWSBedrockEmbeddingParameters(**adapter_metadata).model_dump() + # `auth_type` is a UI-only selector; strip before LiteLLM kwargs. + validation_metadata = { + k: v for k, v in adapter_metadata.items() if k != "auth_type" + } + + validated = AWSBedrockEmbeddingParameters(**validation_metadata).model_dump() + + # Apply Bedrock auth semantics: IAM Role drops keys, Access Keys + # requires non-blank values, legacy (no auth_type) is lenient. + return _resolve_bedrock_aws_credentials(adapter_metadata, validated) @staticmethod def validate_model(adapter_metadata: dict[str, "Any"]) -> str: 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 b71028148c..070c5a9e10 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.json @@ -2,9 +2,7 @@ "title": "Bedrock Embeddings", "type": "object", "required": [ - "aws_secret_access_key", "region_name", - "aws_access_key_id", "model", "adapter_name" ], @@ -21,18 +19,6 @@ "default": "amazon.titan-embed-text-v2:0", "description": "Model name. Refer to [Bedrock's documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) for the list of available models." }, - "aws_access_key_id": { - "type": "string", - "title": "AWS Access Key ID", - "description": "Provide your AWS Access Key ID", - "format": "password" - }, - "aws_secret_access_key": { - "type": "string", - "title": "AWS Secret Access Key", - "description": "Provide your AWS Secret Access Key", - "format": "password" - }, "region_name": { "type": "string", "title": "AWS Region name", @@ -53,6 +39,60 @@ "title": "Timeout", "default": 900, "description": "Timeout in seconds" + }, + "auth_type": { + "type": "string", + "title": "Authentication Type", + "enum": [ + "access_keys", + "iam_role" + ], + "enumNames": [ + "Access Keys", + "IAM Role / Instance Profile (on-prem AWS only)" + ], + "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." + } + }, + "dependencies": { + "auth_type": { + "oneOf": [ + { + "properties": { + "auth_type": { + "enum": [ + "access_keys" + ] + }, + "aws_access_key_id": { + "type": "string", + "title": "AWS Access Key ID", + "description": "Provide your AWS Access Key ID.", + "format": "password" + }, + "aws_secret_access_key": { + "type": "string", + "title": "AWS Secret Access Key", + "description": "Provide your AWS Secret Access Key.", + "format": "password" + } + }, + "required": [ + "aws_access_key_id", + "aws_secret_access_key" + ] + }, + { + "properties": { + "auth_type": { + "enum": [ + "iam_role" + ] + } + } + } + ] } } } 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 9adb8b8ceb..03f23d1376 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json @@ -19,18 +19,6 @@ "default": "amazon.titan-text-express-v1", "description": "Model name. Refer to [Bedrock's documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) for the list of available models." }, - "aws_access_key_id": { - "type": "string", - "title": "AWS Access Key ID", - "description": "Provide your AWS Access Key ID. Leave empty if using AWS Profile or IAM role.", - "format": "password" - }, - "aws_secret_access_key": { - "type": "string", - "title": "AWS Secret Access Key", - "description": "Provide your AWS Secret Access Key. Leave empty if using AWS Profile or IAM role.", - "format": "password" - }, "region_name": { "type": "string", "title": "AWS Region Name", @@ -75,6 +63,60 @@ "title": "Enable Extended Thinking", "default": false, "description": "Enhance reasoning for complex tasks with step-by-step transparency. Available only for Claude 3.7 Sonnet." + }, + "auth_type": { + "type": "string", + "title": "Authentication Type", + "enum": [ + "access_keys", + "iam_role" + ], + "enumNames": [ + "Access Keys", + "IAM Role / Instance Profile (on-prem AWS only)" + ], + "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." + } + }, + "dependencies": { + "auth_type": { + "oneOf": [ + { + "properties": { + "auth_type": { + "enum": [ + "access_keys" + ] + }, + "aws_access_key_id": { + "type": "string", + "title": "AWS Access Key ID", + "description": "Provide your AWS Access Key ID.", + "format": "password" + }, + "aws_secret_access_key": { + "type": "string", + "title": "AWS Secret Access Key", + "description": "Provide your AWS Secret Access Key.", + "format": "password" + } + }, + "required": [ + "aws_access_key_id", + "aws_secret_access_key" + ] + }, + { + "properties": { + "auth_type": { + "enum": [ + "iam_role" + ] + } + } + } + ] } }, "allOf": [ diff --git a/unstract/sdk1/tests/test_bedrock_adapter.py b/unstract/sdk1/tests/test_bedrock_adapter.py new file mode 100644 index 0000000000..4d4b7c1cd7 --- /dev/null +++ b/unstract/sdk1/tests/test_bedrock_adapter.py @@ -0,0 +1,246 @@ +"""Unit tests for the AWS Bedrock LLM and embedding adapters. + +Covers the auth_type selector behaviour added alongside the IAM Role / +Instance Profile mode, plus backwards compatibility for legacy adapter +configurations stored before auth_type existed. +""" + +import pytest +from unstract.sdk1.adapters.base1 import ( + AWSBedrockEmbeddingParameters, + AWSBedrockLLMParameters, +) + +# ── LLM: validate auth_type semantics ──────────────────────────────────────── + + +def test_llm_legacy_no_auth_type_keeps_keys() -> None: + """Legacy adapters without auth_type must keep working unchanged.""" + out = AWSBedrockLLMParameters.validate( + { + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "secret", + } + ) + assert out["aws_access_key_id"] == "AKIAFAKE" + assert out["aws_secret_access_key"] == "secret" + assert out["aws_region_name"] == "us-east-1" + assert "auth_type" not in out + + +def test_llm_access_keys_mode_keeps_keys_and_strips_auth_type() -> 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", + } + ) + assert out["aws_access_key_id"] == "AKIAFAKE" + assert out["aws_secret_access_key"] == "secret" + assert "auth_type" not in out + + +def test_llm_iam_role_mode_drops_keys_even_when_present() -> None: + """IAM Role mode unconditionally drops access keys. + + A saved adapter switched into IAM mode must not silently leak the + previously stored long-lived credentials. + """ + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "iam_role", + "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "region_name": "us-east-1", + "aws_access_key_id": "STALE_KEY", + "aws_secret_access_key": "STALE_SECRET", + } + ) + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + assert out["aws_region_name"] == "us-east-1" + assert "auth_type" not in out + + +def test_llm_iam_role_mode_with_no_keys() -> None: + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "iam_role", + "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "region_name": "us-east-1", + } + ) + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + + +def test_llm_access_keys_mode_blank_keys_raises() -> None: + """Blank values must surface a clear error. + + Falling through to boto3's default chain would hide the user's + misconfiguration and authenticate with whatever ambient creds the + host happens to have. + """ + with pytest.raises(ValueError, match="aws_access_key_id is required"): + AWSBedrockLLMParameters.validate( + { + "auth_type": "access_keys", + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_access_key_id": "", + "aws_secret_access_key": "", + } + ) + + +def test_llm_access_keys_mode_whitespace_keys_raises() -> None: + with pytest.raises(ValueError, match="aws_secret_access_key is required"): + 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": " ", + } + ) + + +def test_llm_unknown_auth_type_raises() -> None: + """A typo or non-UI client must not silently fall through.""" + with pytest.raises(ValueError, match="Unknown auth_type"): + AWSBedrockLLMParameters.validate( + { + "auth_type": "access_key", # typo: missing 's' + "model": "anthropic.claude-3-haiku-20240307-v1:0", + "region_name": "us-east-1", + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "secret", + } + ) + + +def test_llm_other_params_preserved_through_strip() -> None: + """Non-credential params survive the auth-type handling. + + model_id, aws_profile_name, region, and thinking config must pass + through both the strip and the resolver unchanged. + """ + out = AWSBedrockLLMParameters.validate( + { + "auth_type": "iam_role", + "model": "anthropic.claude-3-7-sonnet-20250219-v1:0", + "region_name": "us-east-1", + "aws_profile_name": "dev-profile", + "model_id": ( + "arn:aws:bedrock:us-east-1:1234:application-inference-profile/abc" + ), + "enable_thinking": True, + "budget_tokens": 4096, + } + ) + assert out["aws_profile_name"] == "dev-profile" + assert out["aws_region_name"] == "us-east-1" + assert out["model_id"].endswith("application-inference-profile/abc") + assert out["thinking"] == {"type": "enabled", "budget_tokens": 4096} + + +# ── Embedding: same auth_type matrix ───────────────────────────────────────── + + +def test_embedding_legacy_no_auth_type_keeps_keys() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "secret", + } + ) + assert out["aws_access_key_id"] == "AKIAFAKE" + assert out["aws_secret_access_key"] == "secret" + assert "auth_type" not in out + + +def test_embedding_access_keys_mode_keeps_keys() -> 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", + } + ) + assert out["aws_access_key_id"] == "AKIAFAKE" + assert out["aws_secret_access_key"] == "secret" + assert "auth_type" not in out + + +def test_embedding_iam_role_mode_drops_stale_keys() -> None: + """Embedding-side parity with the LLM stale-key fix.""" + out = AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "iam_role", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_access_key_id": "STALE_KEY", + "aws_secret_access_key": "STALE_SECRET", + } + ) + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + assert out["aws_region_name"] == "us-east-1" + + +def test_embedding_iam_role_mode_with_no_keys() -> None: + out = AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "iam_role", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + } + ) + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + + +def test_embedding_access_keys_mode_blank_keys_raises() -> None: + with pytest.raises(ValueError, match="aws_access_key_id is required"): + AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "access_keys", + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + "aws_access_key_id": "", + "aws_secret_access_key": "", + } + ) + + +def test_embedding_unknown_auth_type_raises() -> None: + with pytest.raises(ValueError, match="Unknown auth_type"): + AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "iamrole", # typo + "model": "amazon.titan-embed-text-v2:0", + "region_name": "us-east-1", + } + ) + + +def test_embedding_region_required_when_absent() -> None: + """aws_region_name is still mandatory even though credentials are not.""" + from pydantic import ValidationError + + with pytest.raises(ValidationError): + AWSBedrockEmbeddingParameters.validate( + { + "auth_type": "iam_role", + "model": "amazon.titan-embed-text-v2:0", + } + )