"""
Custom cost calculator for Cohere rerank models
"""

from typing import Tuple

from litellm.utils import get_model_info


def cost_per_query(model: str, num_queries: int = 1) -> Tuple[float, float]:
    """
    Calculates the cost per query for a given rerank model.

    Input:
        - model: str, the model name without provider prefix

    Returns:
        Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
    """

    model_info = get_model_info(model=model, custom_llm_provider="cohere")

    if (
        "input_cost_per_query" not in model_info
        or model_info["input_cost_per_query"] is None
    ):
        return 0.0, 0.0

    prompt_cost = model_info["input_cost_per_query"] * num_queries

    return prompt_cost, 0.0