"""
Common utilities for NVIDIA NIM rerank provider.
"""


def get_nvidia_nim_rerank_config(model: str):
    """
    Get the appropriate NVIDIA NIM rerank config based on the model.
    
    Args:
        model: The model string (e.g., "nvidia/llama-3.2-nv-rerankqa-1b-v2" or "ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2")
    
    Returns:
        NvidiaNimRankingConfig if model starts with "ranking/", else NvidiaNimRerankConfig
    
    Example:
        - "ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2" -> NvidiaNimRankingConfig
        - "nvidia/llama-3.2-nv-rerankqa-1b-v2" -> NvidiaNimRerankConfig
    """
    from litellm.llms.nvidia_nim.rerank.ranking_transformation import (
        NvidiaNimRankingConfig,
    )
    from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig

    if model.startswith("ranking/"):
        return NvidiaNimRankingConfig()
    return NvidiaNimRerankConfig()