From cb8ce09b0d1ce2a744cb20f20e5d0000545bd23f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 10 Feb 2026 13:37:17 +0530 Subject: [PATCH] Add support for langchain_aws via litellm passthrough --- docs/my-website/docs/pass_through/bedrock.md | 144 ++++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 27 ++++ .../llm_passthrough_endpoints.py | 50 +++++- .../proxy/auth/test_user_api_key_auth.py | 14 +- 4 files changed, 224 insertions(+), 11 deletions(-) diff --git a/docs/my-website/docs/pass_through/bedrock.md b/docs/my-website/docs/pass_through/bedrock.md index b8d20d77da..65c5d8caad 100644 --- a/docs/my-website/docs/pass_through/bedrock.md +++ b/docs/my-website/docs/pass_through/bedrock.md @@ -556,3 +556,147 @@ for event in response.get("completion"): print(completion) ``` + +## Using LangChain AWS SDK with LiteLLM + +You can use the [LangChain AWS SDK](https://python.langchain.com/docs/integrations/chat/bedrock/) with LiteLLM Proxy to get cost tracking, load balancing, and other LiteLLM features. + +### Quick Start + +**1. Install LangChain AWS**: + +```bash showLineNumbers +pip install langchain-aws +``` + +**2. Setup LiteLLM Proxy**: + +Create a `config.yaml`: + +```yaml showLineNumbers +model_list: + - model_name: claude-sonnet + litellm_params: + model: bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0 + aws_region_name: us-east-1 + custom_llm_provider: bedrock +``` + +Start the proxy: + +```bash showLineNumbers +export AWS_ACCESS_KEY_ID="your-access-key" +export AWS_SECRET_ACCESS_KEY="your-secret-key" + +litellm --config config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +**3. Use LangChain with LiteLLM**: + +```python showLineNumbers +from langchain_aws import ChatBedrockConverse +from langchain_core.messages import HumanMessage + +# Your LiteLLM API key +API_KEY = "Bearer sk-1234" + +# Initialize ChatBedrockConverse pointing to LiteLLM proxy +llm = ChatBedrockConverse( + model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + endpoint_url="http://localhost:4000/bedrock", + region_name="us-east-1", + aws_access_key_id=API_KEY, + aws_secret_access_key="bedrock" # Any non-empty value works +) + +# Invoke the model +messages = [HumanMessage(content="Hello, how are you?")] +response = llm.invoke(messages) + +print(response.content) +``` + +### Advanced Example: PDF Document Processing with Citations + +LangChain AWS SDK supports Bedrock's document processing features. Here's how to use it with LiteLLM: + +```python showLineNumbers +import os +import json +from langchain_aws import ChatBedrockConverse +from langchain_core.messages import HumanMessage + +# Your LiteLLM API key +API_KEY = "Bearer sk-1234" + +def get_llm() -> ChatBedrockConverse: + """Initialize LLM pointing to LiteLLM proxy""" + llm = ChatBedrockConverse( + model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + base_model_id="anthropic.claude-3-7-sonnet-20250219-v1:0", + endpoint_url="http://localhost:4000/bedrock", + region_name="us-east-1", + aws_access_key_id=API_KEY, + aws_secret_access_key="bedrock" + ) + return llm + +if __name__ == "__main__": + # Initialize the LLM + llm = get_llm() + + # Read PDF file as bytes (Converse API requires raw bytes) + with open("your-document.pdf", "rb") as file: + file_bytes = file.read() + + # Prepare messages with document attachment + messages = [ + HumanMessage(content=[ + {"text": "What is the policy number in this document?"}, + { + "document": { + "format": "pdf", + "name": "PolicyDocument", + "source": {"bytes": file_bytes}, + "citations": {"enabled": True} + } + } + ]) + ] + + # Invoke the LLM + response = llm.invoke(messages) + + # Print response with citations + print(json.dumps(response.content, indent=4)) +``` + +### Supported LangChain Features + +All LangChain AWS features work with LiteLLM: + +| Feature | Supported | Notes | +|---------|-----------|-------| +| Text Generation | ✅ | Full support | +| Streaming | ✅ | Use `stream()` method | +| Document Processing | ✅ | PDF, images, etc. | +| Citations | ✅ | Enable in document config | +| Tool Use | ✅ | Function calling support | +| Multi-modal | ✅ | Text + images + documents | + +### Troubleshooting + +**Issue**: `UnknownOperationException` error + +**Solution**: Make sure you're using the correct endpoint URL format: +- ✅ Correct: `http://localhost:4000/bedrock` +- ❌ Wrong: `http://localhost:4000/bedrock/v2` + +**Issue**: Authentication errors + +**Solution**: Ensure your API key is in the correct format: +```python +aws_access_key_id="Bearer sk-1234" # Include "Bearer " prefix +``` diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 05eeab3f61..42f10ff859 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -8,6 +8,7 @@ Returns a UserAPIKeyAuth object if the API key is valid """ import asyncio +import re import secrets from datetime import datetime, timezone from typing import List, Optional, Tuple, cast @@ -115,6 +116,18 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str: api_key = api_key.replace("Basic ", "") # handle langfuse input elif api_key.startswith("bearer "): api_key = api_key.replace("bearer ", "") + elif api_key.startswith("AWS4-HMAC-SHA256"): + # Handle AWS Signature V4 format from LangChain + # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... + # Extract the Bearer token from the Credential field + match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + if match: + api_key = match.group(1) + else: + # If no Bearer token found in Credential, try to extract just the credential value + match = re.search(r'Credential=([^/\s,]+)', api_key) + if match: + api_key = match.group(1) return api_key @@ -128,6 +141,20 @@ def _get_bearer_token( api_key = api_key.replace("Basic ", "") # handle langfuse input elif api_key.startswith("bearer "): api_key = api_key.replace("bearer ", "") + elif api_key.startswith("AWS4-HMAC-SHA256"): + # Handle AWS Signature V4 format from LangChain + # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... + # Extract the Bearer token from the Credential field + match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + if match: + api_key = match.group(1) + else: + # If no Bearer token found in Credential, try to extract just the credential value + match = re.search(r'Credential=([^/\s,]+)', api_key) + if match: + api_key = match.group(1) + else: + api_key = "" else: api_key = "" return api_key diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3dab6ea14f..81144ad9f3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -642,10 +642,10 @@ def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: by finding the action in the endpoint and extracting everything between "model" and the action. Args: - endpoint: The endpoint path (e.g., "/model/aws/anthropic/model-name/invoke") + endpoint: The endpoint path (e.g., "/model/aws/anthropic/model-name/invoke" or "v2/model/model-name/invoke") Returns: - The extracted model name (e.g., "aws/anthropic/model-name") + The extracted model name (e.g., "aws/anthropic/model-name" or "model-name") Raises: ValueError: If model cannot be extracted from endpoint @@ -657,7 +657,34 @@ def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: # Format: model/application-inference-profile/{profile-id}/{action} return "/".join(endpoint_parts[1:3]) - # Format: model/{modelId}/{action} + # Format: model/{modelId}/{action} or v2/model/{modelId}/{action} + # Find the index of "model" in the endpoint parts + model_index = None + for idx, part in enumerate(endpoint_parts): + if part == "model": + model_index = idx + break + + # If "model" keyword not found, try to extract model from the endpoint + # by finding the action and taking everything before it + if model_index is None: + # Find the index of the action in the endpoint parts + action_index = None + for idx, part in enumerate(endpoint_parts): + if part in BEDROCK_ENDPOINT_ACTIONS: + action_index = idx + break + + if action_index is not None and action_index > 1: + # Join all parts before the action (excluding empty strings) + model_parts = [p for p in endpoint_parts[1:action_index] if p] + if model_parts: + return "/".join(model_parts) + + raise ValueError( + f"'model' keyword not found and unable to extract model from endpoint. Expected format: /model/{{modelId}}/{{action}}. Got: {endpoint}" + ) + # Find the index of the action in the endpoint parts action_index = None for idx, part in enumerate(endpoint_parts): @@ -665,13 +692,22 @@ def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: action_index = idx break - if action_index is not None and action_index > 1: - # Join all parts between "model" and the action - return "/".join(endpoint_parts[1:action_index]) + if action_index is not None and action_index > model_index + 1: + # Join all parts between "model" and the action (excluding "model" itself) + return "/".join(endpoint_parts[model_index + 1:action_index]) # Fallback to taking everything after "model" if no action found - return "/".join(endpoint_parts[1:]) + model_parts = [p for p in endpoint_parts[model_index + 1:] if p] + if model_parts: + return "/".join(model_parts) + raise ValueError( + f"No model ID found after 'model' keyword. Expected format: /model/{{modelId}}/{{action}}. Got: {endpoint}" + ) + + except ValueError: + # Re-raise ValueError as-is + raise except Exception as e: raise ValueError( f"Model missing from endpoint. Expected format: /model/{{modelId}}/{{action}}. Got: {endpoint}" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 1f379f4371..9b7b7f4615 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -41,6 +41,12 @@ def test_get_api_key(): ("Basic sk-12345678", "sk-12345678", "Basic sk-12345678"), ("bearer sk-12345678", "sk-12345678", "bearer sk-12345678"), ("sk-12345678", "sk-12345678", "sk-12345678"), + # AWS Signature V4 format (LangChain AWS SDK) + ( + "AWS4-HMAC-SHA256 Credential=Bearer sk-12345678/20260210/us-east-1/bedrock/aws4_request, SignedHeaders=host, Signature=abc123", + "sk-12345678", + "AWS4-HMAC-SHA256 Credential=Bearer sk-12345678/20260210/us-east-1/bedrock/aws4_request, SignedHeaders=host, Signature=abc123", + ), ], ) def test_get_api_key_with_custom_litellm_key_header( @@ -243,10 +249,10 @@ async def test_proxy_admin_expired_key_from_cache(): Regression test for issue where PROXY_ADMIN keys from cache skipped expiration check. """ from datetime import datetime, timedelta, timezone - + from fastapi import Request from starlette.datastructures import URL - + from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, @@ -255,7 +261,7 @@ async def test_proxy_admin_expired_key_from_cache(): ) from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder from litellm.proxy.proxy_server import hash_token - + # Create an expired PROXY_ADMIN key api_key = "sk-test-proxy-admin-key" hashed_key = hash_token(api_key) @@ -368,7 +374,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): from user_obj attributes. """ from datetime import datetime - + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj