UN-3152 [FEAT] Add AWS_BEARER_TOKEN_BEDROCK auth to Bedrock LLM and Embedding adapters#1952
Open
jaseemjaskp wants to merge 5 commits into
Open
UN-3152 [FEAT] Add AWS_BEARER_TOKEN_BEDROCK auth to Bedrock LLM and Embedding adapters#1952jaseemjaskp wants to merge 5 commits into
jaseemjaskp wants to merge 5 commits into
Conversation
… adapters Adds a third `auth_type` option, `bearer_token`, to the AWS Bedrock LLM and Embedding adapters alongside the existing `access_keys` and `iam_role` modes. The user pastes a Bedrock API key into the form; the shared resolver translates `aws_bearer_token` to LiteLLM's `api_key` kwarg, which LiteLLM v1.82.3+ recognises as the bearer token and sends as `Authorization: Bearer <token>` instead of doing SigV4 signing. The resolver now also strips stale bearer tokens when other auth modes are selected, mirroring the existing iam_role stale-key fix and keeping the LLM and embedding paths symmetric.
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
Summary by CodeRabbit
WalkthroughAdds explicit Bedrock bearer-token auth mode. UI schemas, parameter models, and credential resolver were extended so ChangesBedrock Bearer Token Authentication Support
Sequence DiagramsequenceDiagram
actor UI as User (UI)
participant Schema as JSON Schema
participant Params as Params Model
participant Resolver as Credential Resolver
participant LiteLLM as LiteLLM Config
UI->>Schema: Selects auth_type="bearer_token" and provides aws_bearer_token
Schema-->>UI: Schema-level validation (required/password)
UI->>Params: Construct AWSBedrock*Parameters (auth_type, aws_bearer_token, ...)
Params->>Resolver: validate() → _resolve_bedrock_aws_credentials(auth_type, aws_bearer_token, ...)
alt auth_type == "bearer_token"
Resolver->>Resolver: Trim and require non-blank aws_bearer_token
Resolver->>Resolver: Drop access key fields
Resolver->>Resolver: Set `api_key` = stripped aws_bearer_token
else auth_type == "access_keys"
Resolver->>Resolver: Require non-blank access keys
Resolver->>Resolver: Drop aws_bearer_token
else auth_type == "iam_role"
Resolver->>Resolver: Drop access keys and aws_bearer_token
else auth_type is None (legacy)
Resolver->>Resolver: Leniently strip blanks and drop aws_bearer_token (do not promote)
end
Resolver-->>Params: Return validated/translated kwargs (api_key present only for bearer_token)
Params->>LiteLLM: Pass final kwargs (may include `api_key`)
LiteLLM-->>Params: Config ready for Bedrock calls
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
unstract/sdk1/src/unstract/sdk1/adapters/base1.py (1)
501-505: ⚡ Quick winNormalize stripped credentials before returning them
These branches use
.strip()only to test for non-blank values, but they keep the raw string invalidated. A key/token pasted with a trailing newline or space will pass validation here and then fail authentication downstream; in the legacy branch, whitespace-only access keys are also retained instead of being dropped. Strip and write back the normalized value, and use the same blank check in theauth_type is Nonepath.Suggested fix
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(): + if not isinstance(value, str): raise ValueError(f"{key} is required when auth_type is 'access_keys'.") + value = value.strip() + if not value: + raise ValueError(f"{key} is required when auth_type is 'access_keys'.") + validated[key] = value @@ token = validated.pop(_BEDROCK_BEARER_TOKEN_FIELD, None) - if isinstance(token, str) and token.strip(): - validated[_BEDROCK_LITELLM_BEARER_KWARG] = token - return + if isinstance(token, str): + token = token.strip() + if token: + validated[_BEDROCK_LITELLM_BEARER_KWARG] = token + return @@ for key in _BEDROCK_AWS_KEY_FIELDS: - if not validated.get(key): + value = validated.get(key) + if not isinstance(value, str) or not value.strip(): validated.pop(key, None) + else: + validated[key] = value.strip()Also applies to: 517-520, 582-585
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@unstract/sdk1/src/unstract/sdk1/adapters/base1.py` around lines 501 - 505, The validation currently only uses .strip() to test for non-empty strings but leaves the raw values in the validated dict; update _require_bedrock_access_keys (and the similar validation blocks referenced for the legacy and auth_type is None branches) to strip each credential value and write the normalized string back into validated before further checks, and use the same blank-check (e.g., if not value) on the stripped value so trailing/leading whitespace is removed and purely-whitespace values are treated as missing.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@unstract/sdk1/src/unstract/sdk1/adapters/base1.py`:
- Around line 521-525: The current check around _BEDROCK_BEARER_TOKEN_FIELD (in
the code path for auth_type == "bearer_token") unconditionally raises when
required and the metadata field is missing, which blocks env-injected
short-lived tokens; change the logic so you only error when the field is
explicitly required and neither a non-blank metadata token nor the environment
fallback (AWS_BEARER_TOKEN_BEDROCK) is available, and only translate/set
aws_bearer_token when a non-empty value is provided in metadata; apply the same
fix to the second occurrence handling at the other block (the lines around where
similar ValueError is raised).
---
Nitpick comments:
In `@unstract/sdk1/src/unstract/sdk1/adapters/base1.py`:
- Around line 501-505: The validation currently only uses .strip() to test for
non-empty strings but leaves the raw values in the validated dict; update
_require_bedrock_access_keys (and the similar validation blocks referenced for
the legacy and auth_type is None branches) to strip each credential value and
write the normalized string back into validated before further checks, and use
the same blank-check (e.g., if not value) on the stripped value so
trailing/leading whitespace is removed and purely-whitespace values are treated
as missing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: fbfb14f6-4efe-40d2-a8b5-7147404872bb
📒 Files selected for processing (4)
unstract/sdk1/src/unstract/sdk1/adapters/base1.pyunstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.jsonunstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.jsonunstract/sdk1/tests/test_bedrock_adapter.py
Contributor
|
| Filename | Overview |
|---|---|
| unstract/sdk1/src/unstract/sdk1/adapters/base1.py | Core credential resolver extended with bearer_token branch; token is correctly stripped before storage, stale keys are purged symmetrically across all three modes, and new helpers are well-scoped. No issues found. |
| unstract/sdk1/src/unstract/sdk1/llm.py | Exception handler broadened from ValidationError to ValueError; safe because pydantic.ValidationError subclasses ValueError in both Pydantic v1 and v2, and new resolver ValueErrors are now also caught. |
| unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json | New bearer_token oneOf branch added with minLength: 1, format: password, and required: ["aws_bearer_token"]; schema is consistent with the LLM enum/enumNames ordering. |
| unstract/sdk1/src/unstract/sdk1/adapters/embedding1/static/bedrock.json | Same bearer_token oneOf branch added with minLength: 1 and format: password; mirrors the LLM schema correctly. |
| unstract/sdk1/tests/test_bedrock_adapter.py | 14 new test cases cover bearer-token translation, stale-key eviction, whitespace stripping, blank/missing token errors for both LLM and embedding paths, and legacy mode. Parity between LLM and embedding suites is complete. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["adapter_metadata from UI"] --> B["validate()"]
B --> C["Pydantic model_dump → validated dict"]
C --> D["_resolve_bedrock_aws_credentials"]
D --> F{"auth_type"}
F -- "iam_role" --> G["drop access keys + bearer token\nboto3 chain takes over"]
F -- "access_keys" --> H["require non-blank keys\ndrop bearer token"]
F -- "bearer_token" --> I["drop access keys\ntranslate token to api_key"]
F -- "None legacy" --> J["strip blank keys\ndrop bearer token"]
F -- "unknown" --> K["raise ValueError"]
I --> L["api_key equals token.strip()"]
G --> M["validated dict to LiteLLM"]
H --> M
L --> M
J --> M
Reviews (5): Last reviewed commit: "[FIX] Trim comments per review — drop PR..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
jaseemjaskp
commented
May 11, 2026
Contributor
Author
jaseemjaskp
left a comment
jaseemjaskp
left a comment
There was a problem hiding this comment.
PR Review Toolkit — multi-agent review
Findings from code-reviewer, comment-analyzer, pr-test-analyzer, silent-failure-hunter, type-design-analyzer, and code-simplifier. Existing greptile/coderabbit notes (base1.py:520/525, test_bedrock_adapter.py:396, minLength on both schemas at :134/:113) are not re-posted.
Highlights
- P1 — LLM error-wrapping asymmetry (out-of-diff, surfaced below at
base1.py:522).embedding.py:106catches(ValidationError, ValueError)and wraps inSdkError;llm.py:195catches onlyValidationError. Every newValueErrorthis PR adds (unknown auth_type, blank bearer, blank access keys) will escape the LLM path unwrapped and bypass downstreamexcept SdkError:blocks (llm.py:242/343/436/511/580). Recommend changingllm.py:195toexcept (ValidationError, ValueError) as e:for parity. - P1 — Legacy silent auto-translation of a hand-stored
aws_bearer_token(noauth_type) intoapi_keycan override env-injected credentials with no log line. Seebase1.py:585. - P1 — Test parity drift is structurally guaranteed by the duplicated LLM/Embedding test blocks. Recommend parametrizing across
[AWSBedrockLLMParameters, AWSBedrockEmbeddingParameters]. Seetest_bedrock_adapter.py:160. - P2 — Multiple silent-drop paths (type confusion on non-string token, silent strip of conflicting access keys submitted alongside bearer) — see inline.
- P2 — Type-design opportunities:
StrEnumforauth_type(4 sites + 2 JSON schemas duplicate the strings);Annotated[str, StringConstraints(min_length=1, strip_whitespace=True)]foraws_bearer_tokento push the existingminLength/stripinvariants into the type system.
Positive
_resolve_bedrock_aws_credentialsdocstring is exemplary (enumerates each branch + the why).- Cross-mode stale-credential scrubbing is well-tested in both directions (bearer→keys/iam and keys/iam→bearer).
- Error messages use
match=regex (not substring) and tests target the publicvalidate()API rather than the private helpers — survives refactoring. auth_typeis correctly stripped fromvalidatedbefore reaching LiteLLM kwargs.
…llm.py ValueError, schema minLength - Strip whitespace before storing the bearer token so a copy-pasted key with stray spaces does not produce `Authorization: Bearer <tok> ` (greptile P1). - Drop the legacy auto-translation of `aws_bearer_token` in the no- auth_type branch. The field is brand new in this PR, so any value there came from manual DB editing; auto-promoting it would silently override env-injected `AWS_BEARER_TOKEN_BEDROCK` or boto3 default- chain credentials with no log line. Users must opt in by setting `auth_type='bearer_token'` explicitly. - Catch `ValueError` (in addition to `ValidationError`) at `llm.py:195` so the new resolver errors surface as `SdkError` instead of escaping past the downstream `except SdkError:` blocks. Embedding side already catches both. - Add `minLength: 1` to the `aws_bearer_token` field in both LLM and embedding schemas so the form renderer can flag a blank token client- side. - Replace the subjective "recommended for developers and CI" framing in the auth_type description with a region-scoping warning — Bedrock API keys are region-scoped and a mismatch surfaces as an opaque 401. - Improve `aws_bearer_token` field comment in both parameter classes from `# AWS_BEARER_TOKEN_BEDROCK` to a descriptive note matching the sibling-field style. Tests: add missing-token, whitespace-strip, embedding-whitespace, and bearer-typo cases on both sides; flip the legacy translation tests to assert the new "drop silently" behavior. 35 passed (was 28).
jaseemjaskp
changed the title
…r in except Pydantic v2 `ValidationError` inherits from `ValueError`, so listing both in the same except clause is redundant and Sonar flags it. Simplify to `except ValueError` and drop the now-unused import; behaviour is unchanged (Pydantic validation errors and the Bedrock auth resolver's explicit ValueErrors are both caught and re-raised as `SdkError`).
…oken-auth # Conflicts: # unstract/sdk1/src/unstract/sdk1/llm.py
chandrasekharan-zipstack
approved these changes
May 11, 2026
Contributor
chandrasekharan-zipstack
left a comment
chandrasekharan-zipstack
left a comment
There was a problem hiding this comment.
LGTM but left a note on the comments itself
pk-zipstack
approved these changes
May 11, 2026
…ational filler Per @chandrasekharan-zipstack: comments should make sense to a future reader of the merged code, not reference the conversation that produced them. Audited every comment/docstring added in this PR: - base1.py constants block: drop the `see litellm/llms/bedrock/base_aws_llm.py` file pointer (rots if upstream restructures) — keep only the WHY (UI uses `aws_bearer_token`, LiteLLM uses `api_key`). - `_translate_bedrock_bearer_token` docstring: collapse to two sentences; the AWS opaque-401 rationale for stripping whitespace is the only non-obvious bit worth keeping. - `_resolve_bedrock_aws_credentials` docstring: lift the shared safety invariant (each branch drops *other* modes' creds; explicit modes raise on blank values) into one paragraph above the per-mode bullets, so each bullet stays one line. - Legacy-branch inline comment: drop "the field is new in this PR" — describe the rule (bearer auth must be opted into explicitly), not its history. - Field comments on both parameter classes: drop the "Short-lived Bedrock API key" prefix (the field name says it). - llm.py:195: drop the Bedrock-resolver-specific reference; keep only the load-bearing invariant (`ValidationError` ⊂ `ValueError`). - Test docstrings: drop two redundant ones (test names already convey the intent); rewrite the legacy-stealth-promotion docstring to describe the rule rather than its PR provenance. 35/35 tests still pass.
|
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds a third Authentication Type option — Bedrock API Key (Bearer Token) — to the AWS Bedrock LLM and Embedding adapters, alongside the existing Access Keys and IAM Role / Instance Profile modes.
Why
AWS now issues short-lived Bedrock API keys (
AWS_BEARER_TOKEN_BEDROCK) that authenticate via a singleAuthorization: Bearer <token>header instead of full SigV4 signing with long-lived IAM access keys. This is the AWS-recommended path for developer and CI workflows that need rotatable credentials without provisioning IAM users. Until now Unstract had no way to surface this option to end users.How
_resolve_bedrock_aws_credentials()helper inunstract/sdk1/src/unstract/sdk1/adapters/base1.pyto recogniseauth_type == "bearer_token"and translate the UI'saws_bearer_tokenfield to LiteLLM'sapi_keykwarg (whichlitellm/llms/bedrock/base_aws_llm.pyconsumes as the bearer token).aws_bearer_token: str | None = Noneto bothAWSBedrockLLMParametersandAWSBedrockEmbeddingParameters.oneOfbranch with a password-formattedaws_bearer_tokenfield to bothbedrock.jsonschemas (LLM + embedding)._drop_bedrock_access_keys,_require_bedrock_access_keys,_translate_bedrock_bearer_token) so the dispatch stays under the project's complexity threshold.Can this PR break any existing features. If yes, please list possible items. If no, please explain why.
No regressions expected. The existing
auth_typematrix (None,access_keys,iam_role) takes the same paths as before — verified by the original 14 tests, all of which still pass unchanged. The new logic only kicks in when a user explicitly selects the Bearer Token mode in the UI. Pre-existing adapters with noauth_typeretain their lenient behaviour, and the new field defaults toNoneso loaded Pydantic models stay compatible.Database Migrations
Env Config
AWS_BEARER_TOKEN_BEDROCKenv var directly if they prefer LiteLLM's env-fallback path, but Unstract does not depend on it.Relevant Docs
Related Issues or PRs
auth_typeselector that introduced IAM Role / Instance Profile mode.Dependencies Versions
unstract/sdk1/pyproject.tomlvia the Zipstack fork) — supports bearer-token auth out of the box.Notes on Testing
Unit tests live in
unstract/sdk1/tests/test_bedrock_adapter.py. New cases cover: bearer-token translation toapi_key, stale access-key drop in bearer mode, blank/whitespace token rejection, stale-token drop wheniam_roleoraccess_keysis selected, and the legacy hand-saved-token translation path. Run locally with:Result: 28 passed (14 pre-existing + 14 new).
Manual end-to-end smoke check (optional, requires a real Bedrock API key):
unstract/sdk1.Authorization: Bearer …rather than the SigV4AWS4-HMAC-SHA256 …header.Screenshots
N/A — backend + JSON-schema-driven form changes only. The new dropdown option will render automatically from the updated
bedrock.jsonschemas.Checklist
I have read and understood the Contribution Guidelines.