Skip to content

Commit 4c983f9

Browse files
ishaan-jaffishaan-jaff
authored
[Feat] Add Bedrock Twelve Labs embedding provider support (#14697)
* fix: add 12 labs to bedrock embedding * fix: get_bedrock_embedding_provider * test: test_text_embedding * fix: 12 labs embedding transform * fix: refactor 12 labs transform logic * fix: test_e2e_bedrock_embedding * fix: test_e2e_bedrock_embedding * feat: add bedrock twelvelabs pricing * DOCS: docs bedrock embedding * DOCS: 12 labs bedrock overview * fix: bedrock embeddings 12 labs
1 parent ec61a71 commit 4c983f9

File tree

11 files changed

+419
-105
lines changed

11 files changed

+419
-105
lines changed

‎docs/my-website/docs/providers/bedrock.md‎

Lines changed: 0 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1821,40 +1821,6 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
18211821
| Mistral 7B Instruct | `completion(model='bedrock/mistral.mistral-7b-instruct-v0:2', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
18221822
| Mixtral 8x7B Instruct | `completion(model='bedrock/mistral.mixtral-8x7b-instruct-v0:1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
18231823

1824-
## Bedrock Embedding
1825-
1826-
### API keys
1827-
This can be set as env variables or passed as **params to litellm.embedding()**
1828-
```python
1829-
import os
1830-
os.environ["AWS_ACCESS_KEY_ID"] = "" # Access key
1831-
os.environ["AWS_SECRET_ACCESS_KEY"] = "" # Secret access key
1832-
os.environ["AWS_REGION_NAME"] = "" # us-east-1, us-east-2, us-west-1, us-west-2
1833-
```
1834-
1835-
### Usage
1836-
```python
1837-
from litellm import embedding
1838-
response = embedding(
1839-
model="bedrock/amazon.titan-embed-text-v1",
1840-
input=["good morning from litellm"],
1841-
)
1842-
print(response)
1843-
```
1844-
1845-
## Supported AWS Bedrock Embedding Models
1846-
1847-
| Model Name | Usage | Supported Additional OpenAI params |
1848-
|----------------------|---------------------------------------------|-----|
1849-
| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py#L59) |
1850-
| Titan Embeddings - V1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py#L53)
1851-
| Titan Multimodal Embeddings | `embedding(model="bedrock/amazon.titan-embed-image-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py#L28) |
1852-
| Cohere Embeddings - English | `embedding(model="bedrock/cohere.embed-english-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
1853-
| Cohere Embeddings - Multilingual | `embedding(model="bedrock/cohere.embed-multilingual-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
1854-
1855-
### Advanced - [Drop Unsupported Params](https://docs.litellm.ai/docs/completion/drop_params#openai-proxy-usage)
1856-
1857-
### Advanced - [Pass model/provider-specific Params](https://docs.litellm.ai/docs/completion/provider_specific_params#proxy-usage)
18581824

18591825
## Image Generation
18601826
Use this for stable diffusion, and amazon nova canvas on bedrock

‎docs/my-website/docs/providers/bedrock_embedding.md‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
## Bedrock Embedding
2+
3+
## Supported Embedding Models
4+
5+
| Provider | LiteLLM Route | AWS Documentation |
6+
|----------|---------------|-------------------|
7+
| Amazon Titan | `bedrock/amazon.*` | [Amazon Titan Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html) |
8+
| Cohere | `bedrock/cohere.*` | [Cohere Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-embed.html) |
9+
| TwelveLabs | `bedrock/us.twelvelabs.*` | [TwelveLabs](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-twelvelabs.html) |
10+
11+
### API keys
12+
This can be set as env variables or passed as **params to litellm.embedding()**
13+
```python
14+
import os
15+
os.environ["AWS_ACCESS_KEY_ID"] = "" # Access key
16+
os.environ["AWS_SECRET_ACCESS_KEY"] = "" # Secret access key
17+
os.environ["AWS_REGION_NAME"] = "" # us-east-1, us-east-2, us-west-1, us-west-2
18+
```
19+
20+
## Usage
21+
### LiteLLM Python SDK
22+
```python
23+
from litellm import embedding
24+
response = embedding(
25+
model="bedrock/amazon.titan-embed-text-v1",
26+
input=["good morning from litellm"],
27+
)
28+
print(response)
29+
```
30+
31+
### LiteLLM Proxy Server
32+
33+
#### 1. Setup config.yaml
34+
```yaml
35+
model_list:
36+
- model_name: titan-embed-v1
37+
litellm_params:
38+
model: bedrock/amazon.titan-embed-text-v1
39+
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
40+
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
41+
aws_region_name: us-east-1
42+
- model_name: titan-embed-v2
43+
litellm_params:
44+
model: bedrock/amazon.titan-embed-text-v2:0
45+
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
46+
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
47+
aws_region_name: us-east-1
48+
```
49+
50+
#### 2. Start Proxy
51+
```bash
52+
litellm --config /path/to/config.yaml
53+
```
54+
55+
#### 3. Use with OpenAI Python SDK
56+
```python
57+
import openai
58+
client = openai.OpenAI(
59+
api_key="anything",
60+
base_url="http://0.0.0.0:4000"
61+
)
62+
63+
response = client.embeddings.create(
64+
input=["good morning from litellm"],
65+
model="titan-embed-v1"
66+
)
67+
print(response)
68+
```
69+
70+
#### 4. Use with LiteLLM Python SDK
71+
```python
72+
import litellm
73+
response = litellm.embedding(
74+
model="titan-embed-v1", # model alias from config.yaml
75+
input=["good morning from litellm"],
76+
api_base="http://0.0.0.0:4000",
77+
api_key="anything"
78+
)
79+
print(response)
80+
```
81+
82+
## Supported AWS Bedrock Embedding Models
83+
84+
| Model Name | Usage | Supported Additional OpenAI params |
85+
|----------------------|---------------------------------------------|-----|
86+
| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py#L59) |
87+
| Titan Embeddings - V1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py#L53)
88+
| Titan Multimodal Embeddings | `embedding(model="bedrock/amazon.titan-embed-image-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py#L28) |
89+
| TwelveLabs Marengo Embed 2.7 | `embedding(model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", input=input)` | Supports multimodal input (text, video, audio, image) |
90+
| Cohere Embeddings - English | `embedding(model="bedrock/cohere.embed-english-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
91+
| Cohere Embeddings - Multilingual | `embedding(model="bedrock/cohere.embed-multilingual-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
92+
93+
### Advanced - [Drop Unsupported Params](https://docs.litellm.ai/docs/completion/drop_params#openai-proxy-usage)
94+
95+
### Advanced - [Pass model/provider-specific Params](https://docs.litellm.ai/docs/completion/provider_specific_params#proxy-usage)

‎docs/my-website/sidebars.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,7 @@ const sidebars = {
411411
label: "Bedrock",
412412
items: [
413413
"providers/bedrock",
414+
"providers/bedrock_embedding",
414415
"providers/bedrock_agents",
415416
"providers/bedrock_batches",
416417
"providers/bedrock_vector_store",

‎litellm/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
bedrock_embedding_models,
6868
known_tokenizer_config,
6969
BEDROCK_INVOKE_PROVIDERS_LITERAL,
70+
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
7071
BEDROCK_CONVERSE_MODELS,
7172
DEFAULT_MAX_TOKENS,
7273
DEFAULT_SOFT_BUDGET,

‎litellm/constants.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -769,6 +769,12 @@
769769
"deepseek_r1",
770770
]
771771

772+
BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[
773+
"cohere",
774+
"amazon",
775+
"twelvelabs",
776+
]
777+
772778
BEDROCK_CONVERSE_MODELS = [
773779
"openai.gpt-oss-20b-1:0",
774780
"openai.gpt-oss-120b-1:0",

‎litellm/llms/bedrock/base_aws_llm.py‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,11 @@
2020

2121
from litellm._logging import verbose_logger
2222
from litellm.caching.caching import DualCache
23-
from litellm.constants import BEDROCK_INVOKE_PROVIDERS_LITERAL, BEDROCK_MAX_POLICY_SIZE
23+
from litellm.constants import (
24+
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
25+
BEDROCK_INVOKE_PROVIDERS_LITERAL,
26+
BEDROCK_MAX_POLICY_SIZE,
27+
)
2428
from litellm.litellm_core_utils.dd_tracing import tracer
2529
from litellm.secret_managers.main import get_secret, get_secret_str
2630

@@ -327,6 +331,40 @@ def get_bedrock_invoke_provider(
327331
return provider
328332
return None
329333

334+
@staticmethod
335+
def get_bedrock_embedding_provider(
336+
model: str,
337+
) -> Optional[BEDROCK_EMBEDDING_PROVIDERS_LITERAL]:
338+
"""
339+
Helper function to get the bedrock embedding provider from the model
340+
341+
Handles scenarios like:
342+
1. model=cohere.embed-english-v3:0 -> Returns `cohere`
343+
2. model=amazon.titan-embed-text-v1 -> Returns `amazon`
344+
3. model=us.twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs`
345+
4. model=twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs`
346+
"""
347+
# Handle regional models like us.twelvelabs.marengo-embed-2-7-v1:0
348+
if "." in model:
349+
parts = model.split(".")
350+
# Check if the second part (after potential region) is a known provider
351+
if len(parts) >= 2:
352+
potential_provider = parts[1] # e.g., "twelvelabs" from "us.twelvelabs.marengo-embed-2-7-v1:0"
353+
if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
354+
return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider)
355+
356+
# Check if the first part is a known provider (standard format)
357+
potential_provider = parts[0] # e.g., "cohere" from "cohere.embed-english-v3:0"
358+
if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
359+
return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider)
360+
361+
# Fallback: check if any provider name appears in the model string
362+
for provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
363+
if provider in model:
364+
return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, provider)
365+
366+
return None
367+
330368
def _get_aws_region_name(
331369
self,
332370
optional_params: dict,

0 commit comments

Comments
 (0)