feat(bedrock/rerank): migrate bedrock config to basererank config

litellm/__init__.py

@@ -872,6 +872,7 @@ from .llms.bedrock.common_utils import (
 from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import (
     AmazonAI21Config,
 )
+from .llms.bedrock.rerank.transformation import BedrockRerankConfig
 from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import (
     AmazonInvokeNovaConfig,
 )

litellm/llms/base_llm/rerank/transformation.py

@@ -86,8 +86,43 @@ class BaseRerankConfig(ABC):
     ) -> BaseLLMException:
         pass
 
-    @staticmethod
+    def calculate_total_queries(
+        self, query_tokens: int, document_tokens: list[int]
+    ) -> int:
+        """
+        Calculate the cost of a request based on token counts and document chunks.
+
+        Args:
+            query_tokens (int): Number of tokens in the query
+            document_tokens (list[int]): List of token counts for each document
+            cost_per_1000_queries (float): Cost per 1000 queries in dollars (default: $1.00)
+
+        Returns:
+            float: Total cost in dollars
+        """
+        TOKENS_PER_DOCUMENT = 512
+        CHUNKS_PER_QUERY = 100
+
+        # Validate query length
+        if query_tokens >= TOKENS_PER_DOCUMENT:
+            raise ValueError("Query tokens exceed maximum allowed tokens per document")
+
+        # Calculate total chunks needed
+        total_chunks = 0
+        available_tokens = TOKENS_PER_DOCUMENT - query_tokens
+
+        for doc_tokens in document_tokens:
+            # Calculate chunks needed for this document
+            chunks_needed = (doc_tokens + available_tokens - 1) // available_tokens
+            total_chunks += max(1, chunks_needed)
+
+        # Calculate total queries needed (rounded up to nearest multiple of CHUNKS_PER_QUERY)
+        total_queries = (total_chunks + CHUNKS_PER_QUERY - 1) // CHUNKS_PER_QUERY
+
+        return total_queries
+
     def calculate_rerank_cost(
+        self,
         model: str,
         custom_llm_provider: Optional[str] = None,
         num_queries: int = 1,

litellm/llms/bedrock/rerank/transformation.py

@@ -5,8 +5,12 @@ Why separate file? Make it easy to see how transformation works
 """
 
 import uuid
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
+import httpx
+
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.cohere.rerank.transformation import CohereRerankConfig
 from litellm.types.llms.bedrock import (
     BedrockRerankBedrockRerankingConfiguration,
     BedrockRerankConfiguration,
@@ -19,6 +23,7 @@ from litellm.types.llms.bedrock import (
     BedrockRerankTextQuery,
 )
 from litellm.types.rerank import (
+    OptionalRerankParams,
     RerankBilledUnits,
     RerankRequest,
     RerankResponse,
@@ -27,8 +32,36 @@ from litellm.types.rerank import (
     RerankTokens,
 )
 
+from ..common_utils import BedrockError
 
-class BedrockRerankConfig:
+
+class BedrockRerankConfig(CohereRerankConfig):
+    def get_supported_cohere_rerank_params(self, model: str) -> List:
+        if "cohere" in model.lower():
+            return super().get_supported_cohere_rerank_params(model)
+        else:  # amazon model supports restricted params
+            return ["query", "documents", "top_n"]
+
+    def map_cohere_rerank_params(
+        self,
+        non_default_params: Optional[dict],
+        model: str,
+        drop_params: bool,
+        query: str,
+        documents: List[Union[str, Dict[str, Any]]],
+        custom_llm_provider: Optional[str] = None,
+        top_n: Optional[int] = None,
+        rank_fields: Optional[List[str]] = None,
+        return_documents: Optional[bool] = True,
+        max_chunks_per_doc: Optional[int] = None,
+    ) -> OptionalRerankParams:
+        supported_params = self.get_supported_cohere_rerank_params(model)
+        optional_rerank_params = {}
+        if non_default_params:
+            for k, v in non_default_params.items():
+                if k in supported_params:
+                    optional_rerank_params[k] = v
+        return OptionalRerankParams(**optional_rerank_params)
 
     def _transform_sources(
         self, documents: List[Union[str, dict]]
@@ -117,3 +150,8 @@ class BedrockRerankConfig:
             results=_results,
             meta=rerank_meta,
         )  # Return response
+
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
+    ) -> BaseLLMException:
+        return BedrockError(status_code=status_code, message=error_message)

litellm/utils.py

@@ -6198,6 +6198,8 @@ class ProviderConfigManager:
             return litellm.AzureAIRerankConfig()
         elif litellm.LlmProviders.INFINITY == provider:
             return litellm.InfinityRerankConfig()
+        elif litellm.LlmProviders.BEDROCK == provider:
+            return litellm.BedrockRerankConfig()
         return litellm.CohereRerankConfig()
 
     @staticmethod

tests/llm_translation/test_rerank.py

@@ -385,3 +385,61 @@ def test_rerank_response_assertions():
     )
 
     assert_response_shape(r, custom_llm_provider="custom")
+
+
+def calculate_total_queries(query_tokens: int, document_tokens: list[int]) -> int:
+    """
+    Calculate the cost of a request based on token counts and document chunks.
+
+    Args:
+        query_tokens (int): Number of tokens in the query
+        document_tokens (list[int]): List of token counts for each document
+        cost_per_1000_queries (float): Cost per 1000 queries in dollars (default: $1.00)
+
+    Returns:
+        float: Total cost in dollars
+    """
+    TOKENS_PER_DOCUMENT = 512
+    CHUNKS_PER_QUERY = 100
+
+    # Validate query length
+    if query_tokens >= TOKENS_PER_DOCUMENT:
+        raise ValueError("Query tokens exceed maximum allowed tokens per document")
+
+    # Calculate total chunks needed
+    total_chunks = 0
+    available_tokens = TOKENS_PER_DOCUMENT - query_tokens
+
+    for doc_tokens in document_tokens:
+        # Calculate chunks needed for this document
+        chunks_needed = (doc_tokens + available_tokens - 1) // available_tokens
+        total_chunks += max(1, chunks_needed)
+
+    # Calculate total queries needed (rounded up to nearest multiple of CHUNKS_PER_QUERY)
+    total_queries = (total_chunks + CHUNKS_PER_QUERY - 1) // CHUNKS_PER_QUERY
+
+    return total_queries
+
+
+def test_rerank_cohere_api():
+    from litellm import token_counter
+
+    query_tokens = token_counter(text="hello", count_response_tokens=True)
+    document_tokens_1 = token_counter(text="world", count_response_tokens=True)
+    document_tokens_2 = token_counter(text="hello", count_response_tokens=True)
+
+    total_queries = calculate_total_queries(
+        query_tokens, [document_tokens_1, document_tokens_2] * 500
+    )
+
+    print("total_queries", total_queries)
+
+    response = litellm.rerank(
+        model="cohere/rerank-english-v3.0",
+        query="hello",
+        documents=["hello", "world"] * 500,
+        return_documents=True,
+        top_n=3,
+    )
+
+    print("rerank response", response)