diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 27dc785bf57..f78095a6d51 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -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( diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b2836a096b7..127b1f20e2e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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", diff --git a/tests/llm_translation/test_bedrock_embedding.py b/tests/llm_translation/test_bedrock_embedding.py index 92c22f582d9..feda016bd9f 100644 --- a/tests/llm_translation/test_bedrock_embedding.py +++ b/tests/llm_translation/test_bedrock_embedding.py @@ -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 \ No newline at end of file