Repository: meta-llama/llama-stack
File: tests/integration/inference/test_openai_embeddings.py (Python)

Code before:
    provider = provider_from_model(client_with_models, model_id)
    provider_type = provider.provider_type

    if provider_type not in ASYMMETRIC_EMBEDDING_MODELS_BY_PROVIDER:
        return False

    return model_id in ASYMMETRIC_EMBEDDING_MODELS_BY_PROVIDER[provider_type]


def get_extra_body_for_model(client_with_models, model_id, input_type="query"):
    provider = provider_from_model(client_with_models, model_id)

    if provider.provider_type == "remote::nvidia":
        return {"input_type": input_type}

    return None


def verify_and_get_extra_body(compat_client, client_with_models, model_id, **kwargs):
    """
    For asymmetric models, verify that extra_body is required by testing that
    calling without it raises an error. Return the extra_body.
    For other models, return None.
    """
    is_asymmetric = is_asymmetric_model(client_with_models, model_id)
    if is_asymmetric:
        with pytest.raises(Exception):  # noqa: B017
            compat_client.embeddings.create(**kwargs)
        return get_extra_body_for_model(client_with_models, model_id)
    return None


def skip_if_model_doesnt_support_user_param(client, model_id):
    provider = provider_from_model(client, model_id)
    if provider.provider_type in (
        "remote::together",  # service returns 400
        "remote::fireworks",  # service returns 400 malformed input
    ):
        pytest.skip(f"Model {model_id} hosted by {provider.provider_type} does not support user param.")


def skip_if_model_doesnt_support_encoding_format_base64(client, model_id):
    provider = provider_from_model(client, model_id)

    should_skip = provider.provider_type in (
        "remote::databricks",  # param silently ignored, always returns floats
        "remote::fireworks",  # param silently ignored, always returns list of floats
        "remote::ollama",  # param silently ignored, always returns list of floats
    ) or (
        provider.provider_type == "remote::nvidia"
        and model_id

Review comment (refactor):
nah I don't think we need to be that defensive. i'd just simplify this a bunch more. simply make the callsites be: 

```client.embeddings.create(..., extra_body=get_extra_body())```

that's it! 

Code after:
    provider_type = provider.provider_type

    if provider_type not in ASYMMETRIC_EMBEDDING_MODELS_BY_PROVIDER:
        return False

    return model_id in ASYMMETRIC_EMBEDDING_MODELS_BY_PROVIDER[provider_type]


def get_extra_body_for_model(client_with_models, model_id, input_type="query"):
    if not is_asymmetric_model(client_with_models, model_id):
        return None

    provider = provider_from_model(client_with_models, model_id)

    if provider.provider_type == "remote::nvidia":
        return {"input_type": input_type}

    return None


def skip_if_model_doesnt_support_user_param(client, model_id):
    provider = provider_from_model(client, model_id)
    if provider.provider_type in (
        "remote::together",  # service returns 400
        "remote::fireworks",  # service returns 400 malformed input
    ):
        pytest.skip(f"Model {model_id} hosted by {provider.provider_type} does not support user param.")


def skip_if_model_doesnt_support_encoding_format_base64(client, model_id):
    provider = provider_from_model(client, model_id)

    should_skip = provider.provider_type in (
        "remote::databricks",  # param silently ignored, always returns floats
        "remote::fireworks",  # param silently ignored, always returns list of floats
        "remote::ollama",  # param silently ignored, always returns list of floats
    ) or (
        provider.provider_type == "remote::nvidia"
        and model_id
        in [
            "nvidia/nv-embedqa-e5-v5",
            "nvidia/nv-embedqa-mistral-7b-v2",
            "snowflake/arctic-embed-l",
        ]
    )

    if should_skip:
        pytest.skip(f"Model {model_id} hosted by {provider.provider_type} does not support encoding_format='base64'.")


def skip_if_model_doesnt_support_variable_dimensions(client_with_models, model_id):