From 4e7d3523af5a4383ea78050f8f623bc8ad109c5d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 20 Feb 2025 14:12:51 -0800 Subject: [PATCH] test(test_bedrock_completion.py): add testing for bedrock cohere rerank --- tests/llm_translation/test_bedrock_completion.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 685c5e5409c..7fca7b5f1ae 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -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", [