scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(test_bedrock_completion.py): add testing for bedrock cohere rerank

Browse source
This commit is contained in:
Krrish Dholakia 2025-02-20 14:12:51 -08:00
parent d7efe22d6e
commit 4e7d3523af
1 changed files with 10 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
tests/llm_translation/test_bedrock_completion.py
View file

@ -2186,6 +2186,16 @@ class TestBedrockRerank(BaseLLMRerankTest):
}
class TestBedrockCohereRerank(BaseLLMRerankTest):
def get_custom_llm_provider(self) -> litellm.LlmProviders:
return litellm.LlmProviders.BEDROCK
def get_base_rerank_call_args(self) -> dict:
return {
"model": "bedrock/arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0",
}
@pytest.mark.parametrize(
"messages, continue_message_index",
[