diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 30d3ed0364..1eb4d0eb94 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -94,6 +94,8 @@ def cost_per_token( "transcription", "aspeech", "speech", + "rerank", + "arerank", ] = "completion", ) -> Tuple[float, float]: """ @@ -488,6 +490,8 @@ def completion_cost( "transcription", "aspeech", "speech", + "rerank", + "arerank", ] = "completion", ### REGION ### custom_llm_provider=None, @@ -793,7 +797,9 @@ def response_cost_calculator( call_type=call_type, custom_llm_provider=custom_llm_provider, ) - elif isinstance(response_object, RerankResponse): + elif isinstance(response_object, RerankResponse) and ( + call_type == "arerank" or call_type == "rerank" + ): response_cost = rerank_cost( rerank_response=response_object, model=model, @@ -838,7 +844,7 @@ def rerank_cost( model: str, call_type: Literal["rerank", "arerank"], custom_llm_provider: Optional[str], -) -> Optional[float]: +) -> float: """ Returns - float or None: cost of response OR none if error. @@ -848,5 +854,11 @@ def rerank_cost( try: if custom_llm_provider == "cohere": return 0.002 + raise ValueError( + f"invalid custom_llm_provider for rerank model: {model}, custom_llm_provider: {custom_llm_provider}" + ) except Exception as e: + verbose_logger.exception( + f"litellm.cost_calculator.py::rerank_cost - Exception occurred - {str(e)}" + ) raise e diff --git a/litellm/rerank_api/types.py b/litellm/rerank_api/types.py index 3b95597c31..00cb32c180 100644 --- a/litellm/rerank_api/types.py +++ b/litellm/rerank_api/types.py @@ -23,7 +23,7 @@ class RerankResponse(BaseModel): id: str results: List[dict] # Contains index and relevance_score meta: dict # Contains api_version and billed_units - _hidden_params: Optional[dict] = {} + _hidden_params: dict = {} class Config: underscore_attrs_are_private = True