mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
feat(bedrock): add amazon.titan-embed-g1-text-02 embedding model support
- Add model to provider routing allowlist in embedding.py - Add request transformation using AmazonTitanG1Config - Add response transformation using AmazonTitanG1Config - Add pricing metadata to model_prices_and_context_window.json - Add unit tests for embedding and model info Fixes missing cost tracking reported in #29786 Related to VANDRANKI/litellm PR #29790
This commit is contained in:
parent
aaf1e2444b
commit
07c633be3a
3 changed files with 31 additions and 0 deletions
|
|
@ -224,6 +224,10 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
returned_response = AmazonTitanV2Config()._transform_response(
|
||||
response_list=response_list, model=model
|
||||
)
|
||||
elif model == "amazon.titan-embed-g1-text-02":
|
||||
returned_response = AmazonTitanG1Config()._transform_response(
|
||||
response_list=response_list, model=model
|
||||
)
|
||||
elif provider == "twelvelabs":
|
||||
returned_response = (
|
||||
TwelveLabsMarengoEmbeddingConfig()._transform_response(
|
||||
|
|
@ -447,6 +451,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
"amazon.titan-embed-image-v1",
|
||||
"amazon.titan-embed-text-v1",
|
||||
"amazon.titan-embed-text-v2:0",
|
||||
"amazon.titan-embed-g1-text-02",
|
||||
]:
|
||||
batch_data = []
|
||||
for i in input:
|
||||
|
|
@ -464,6 +469,10 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
transformed_request = AmazonTitanV2Config()._transform_request(
|
||||
input=i, inference_params=inference_params
|
||||
)
|
||||
elif model == "amazon.titan-embed-g1-text-02":
|
||||
transformed_request = AmazonTitanG1Config()._transform_request(
|
||||
input=i, inference_params=inference_params
|
||||
)
|
||||
else:
|
||||
raise Exception(
|
||||
"Unmapped model. Received={}. Expected={}".format(
|
||||
|
|
|
|||
|
|
@ -570,6 +570,15 @@
|
|||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1536
|
||||
},
|
||||
"amazon.titan-embed-g1-text-02": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1536
|
||||
},
|
||||
"amazon.titan-embed-text-v2:0": {
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "bedrock",
|
||||
|
|
|
|||
|
|
@ -34,6 +34,11 @@ img_base_64 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkBAMAAACCzIh
|
|||
"text",
|
||||
titan_embedding_response,
|
||||
), # V2 text model
|
||||
(
|
||||
"bedrock/amazon.titan-embed-g1-text-02",
|
||||
"text",
|
||||
titan_embedding_response,
|
||||
), # G1 text model
|
||||
(
|
||||
"bedrock/amazon.titan-embed-image-v1",
|
||||
"image",
|
||||
|
|
@ -459,3 +464,11 @@ def test_bedrock_embedding_region_bug_reproduction():
|
|||
os.environ["AWS_REGION_NAME"] = original_region_name
|
||||
else:
|
||||
os.environ.pop("AWS_REGION_NAME", None)
|
||||
def test_bedrock_titan_g1_text_02_model_info():
|
||||
"""Test that amazon.titan-embed-g1-text-02 has correct pricing metadata"""
|
||||
model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02")
|
||||
assert model_info is not None, "Model info should not be None"
|
||||
assert model_info["litellm_provider"] == "bedrock"
|
||||
assert model_info["mode"] == "embedding"
|
||||
assert model_info["input_cost_per_token"] == 1e-07
|
||||
assert model_info["max_input_tokens"] == 8192
|
||||
Loading…
Add table
Reference in a new issue