feat(bedrock): add amazon.titan-embed-g1-text-02 embedding model support - #29891
feat(bedrock): add amazon.titan-embed-g1-text-02 embedding model support#29891bhumikadangayach wants to merge 6 commits into
Conversation
- 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 BerriAI#29786 Related to VANDRANKI/litellm PR BerriAI#29790
|
|
Greptile SummaryThis PR adds support for the
Confidence Score: 4/5Safe to merge; changes are additive and follow the existing Titan embedding pattern throughout. The implementation mirrors the existing titan-embed-text-v1 setup end-to-end. The only nits are trailing whitespace on the new elif branches and a missing blank line before the new test function. No files require special attention beyond the minor style issues noted in the comments.
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/embed/embedding.py | Adds amazon.titan-embed-g1-text-02 to the provider allow-list and wires it to AmazonTitanG1Config for both request and response transformation; trailing whitespace on the two new elif branches and a stale error message in the unreachable else branch. |
| model_prices_and_context_window.json | Adds pricing entry for amazon.titan-embed-g1-text-02 with correct litellm_provider, mode, input cost, and vector size (1536), consistent with the existing titan-embed-text-v1 entry. |
| tests/llm_translation/test_bedrock_embedding.py | Extends the parametrized embedding test to cover the new model using the same mock pattern; adds a model-info unit test. The new test function is missing a blank line separator and the file ends without a trailing newline. |
Comments Outside Diff (1)
-
litellm/llms/bedrock/embed/embedding.py, line 472-486 (link)Trailing whitespace on both new
elifbranches. Also the error message in theelseblock no longer listsamazon.titan-embed-g1-text-02— while that branch is now unreachable, keeping the list accurate avoids confusion if the block is ever reached via future refactoring.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Reviews (1): Last reviewed commit: "feat(bedrock): add amazon.titan-embed-g1..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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 |
There was a problem hiding this comment.
| 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 | |
| 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 |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
"Hi, all required checks are passing. Happy to make any changes if needed!" |
|
Hi, just following up — this has been open about 3 weeks with no review yet. All required checks are passing aside from a patch-coverage threshold on 3 lines (a defensive branch in the error-message path). Happy to add a test for that if it'd help, or address anything else needed to move this forward. Thanks! |
2 participants
Problem
amazon.titan-embed-g1-text-02isn't in LiteLLM's routing table, so anyrequest to it fails with
Unable to map Bedrock request to provider.Fixes #29786
What I changed
amazon.titan-embed-g1-text-02to the provider routing allowlistAmazonTitanG1Config—same format as
titan-embed-text-v1since they share the same inputText schemamodel_prices_and_context_window.json— PR feat(bedrock): add amazon.titan-embed-g1-text-02 embedding model support #29790missed this, which meant cost tracking would silently return $0 for every call
Testing
Added
amazon.titan-embed-g1-text-02to the existing parametrizedtest_bedrock_embedding_modelstest, and a newtest_bedrock_titan_g1_text_02_model_infoto verify the pricing metadatais correct.