diff --git a/litellm/__init__.py b/litellm/__init__.py index cbbeca6750a..132cb4067e9 100644 --- a/litellm/__init__.py +++ b/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, ) diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 539009381af..ab7d668655c 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/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, diff --git a/litellm/llms/bedrock/rerank/transformation.py b/litellm/llms/bedrock/rerank/transformation.py index a5380febe9c..76690fa12db 100644 --- a/litellm/llms/bedrock/rerank/transformation.py +++ b/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) diff --git a/litellm/utils.py b/litellm/utils.py index d18cfed20db..fe95bbe1aaa 100644 --- a/litellm/utils.py +++ b/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 diff --git a/tests/llm_translation/test_rerank.py b/tests/llm_translation/test_rerank.py index 82efa92dfda..258a9e4eff6 100644 --- a/tests/llm_translation/test_rerank.py +++ b/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)