litellm/tests/integration/configuration/test_callback_settings_boot.py

import json
import uuid
from collections.abc import Mapping
from pathlib import Path
from typing import Final

import pytest
from pydantic import JsonValue

from tests.integration._support.client import Gateway
from tests.integration._support.process import owned_proxy

SERVING_CONSUMERS: Final = {
    "compression_interception": "CompressionInterceptionLogger",
    "code_interpreter_interception": "CodeInterpreterInterceptionLogger",
    "websearch_interception": "WebSearchInterceptionLogger",
}
OTEL_CONSUMER: Final = {"otel": "OpenTelemetry"}
GUARDRAIL_CONSUMERS: Final = {
    "presidio": "_OPTIONAL_PresidioPIIMasking",
    "lakera_prompt_injection": "lakeraAI_Moderation",
}

TOP_LEVEL_SHAPES: Final = (
    pytest.param({}, id="empty-object"),
    pytest.param(None, id="null"),
    pytest.param("otel", id="string"),
    pytest.param(["otel"], id="list"),
    pytest.param(True, id="bool"),
    pytest.param(0, id="zero"),
)

CONSUMER_SHAPES: Final = (
    pytest.param({}, id="empty-object"),
    pytest.param(None, id="null"),
    pytest.param("on", id="string"),
    pytest.param(True, id="bool"),
    pytest.param([], id="empty-list"),
    pytest.param(["on"], id="list"),
    pytest.param(7, id="int"),
)

TOP_LEVEL_BOOT_CRASH: Final = (
    "BUG: a non-object callback_settings is stored verbatim and proxy startup crashes calling .get on it"
)
TOP_LEVEL_BOOT_CRASH_IDS: Final = frozenset({"string", "list", "bool"})
OTEL_DROPPED: Final = (
    "BUG: a non-object callback_settings.otel fails dict() and the otel callback is silently not registered"
)
OTEL_DROPPED_IDS: Final = frozenset({"null", "string", "bool", "int"})


def _write_config(
    directory: Path, upstream_url: str, model: str, callbacks: tuple[str, ...], callback_settings: JsonValue
) -> Path:
    config: Final = directory / f"callback_settings_{uuid.uuid4().hex}.yaml"
    config.write_text(
        json.dumps(
            {
                "model_list": [
                    {
                        "model_name": model,
                        "litellm_params": {
                            "model": f"openai/{model}",
                            "api_base": f"{upstream_url}/v1",
                            "api_key": "integration-provider-key",
                        },
                    }
                ],
                "litellm_settings": {"callbacks": list(callbacks)},
                "callback_settings": callback_settings,
                "general_settings": {
                    "master_key": "os.environ/LITELLM_MASTER_KEY",
                    "database_url": "os.environ/DATABASE_URL",
                },
            }
        )
    )
    return config


def _assert_registered(candidate: Gateway, consumers: Mapping[str, str]) -> None:
    response: Final = candidate.request("GET", "/active/callbacks")
    assert response.status_code == 200, response.text
    missing: Final = sorted(name for name, class_name in consumers.items() if class_name not in response.text)
    assert missing == [], response.text


@pytest.mark.parametrize("callback_settings", TOP_LEVEL_SHAPES)
def test_top_level_callback_settings_shape_boots_registers_and_serves_chat(
    gateway: Gateway, tmp_path: Path, callback_settings: JsonValue, request: pytest.FixtureRequest
) -> None:
    if request.node.callspec.id in TOP_LEVEL_BOOT_CRASH_IDS:
        pytest.skip(TOP_LEVEL_BOOT_CRASH)
    consumers: Final = {**SERVING_CONSUMERS, **OTEL_CONSUMER}
    model: Final = f"integration-callback-settings-{uuid.uuid4().hex}"
    config: Final = _write_config(tmp_path, gateway.upstream_url, model, tuple(consumers), callback_settings)
    with owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config) as candidate:
        _assert_registered(candidate, consumers)
        reply: Final = candidate.chat(model, text=f"callback settings {uuid.uuid4().hex}")
        assert reply["model"] == model, reply


@pytest.mark.parametrize("value", CONSUMER_SHAPES)
def test_serving_consumer_settings_shape_boots_registers_and_serves_chat(
    gateway: Gateway, tmp_path: Path, value: JsonValue
) -> None:
    model: Final = f"integration-callback-settings-{uuid.uuid4().hex}"
    config: Final = _write_config(
        tmp_path,
        gateway.upstream_url,
        model,
        tuple(SERVING_CONSUMERS),
        {consumer: value for consumer in SERVING_CONSUMERS},
    )
    with owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config) as candidate:
        _assert_registered(candidate, SERVING_CONSUMERS)
        reply: Final = candidate.chat(model, text=f"callback settings {uuid.uuid4().hex}")
        assert reply["model"] == model, reply


@pytest.mark.parametrize("value", CONSUMER_SHAPES)
def test_otel_settings_shape_boots_registers_and_serves_chat(
    gateway: Gateway, tmp_path: Path, value: JsonValue, request: pytest.FixtureRequest
) -> None:
    if request.node.callspec.id in OTEL_DROPPED_IDS:
        pytest.skip(OTEL_DROPPED)
    model: Final = f"integration-callback-settings-{uuid.uuid4().hex}"
    config: Final = _write_config(tmp_path, gateway.upstream_url, model, tuple(OTEL_CONSUMER), {"otel": value})
    with owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config) as candidate:
        _assert_registered(candidate, OTEL_CONSUMER)
        reply: Final = candidate.chat(model, text=f"callback settings {uuid.uuid4().hex}")
        assert reply["model"] == model, reply


@pytest.mark.parametrize("value", CONSUMER_SHAPES)
def test_guardrail_consumer_settings_shape_boots_and_registers(
    gateway: Gateway, tmp_path: Path, value: JsonValue
) -> None:
    model: Final = f"integration-callback-settings-{uuid.uuid4().hex}"
    config: Final = _write_config(
        tmp_path,
        gateway.upstream_url,
        model,
        tuple(GUARDRAIL_CONSUMERS),
        {consumer: value for consumer in GUARDRAIL_CONSUMERS},
    )
    environment: Final = {
        "STORE_MODEL_IN_DB": "False",
        "PRESIDIO_ANALYZER_API_BASE": gateway.upstream_url,
        "PRESIDIO_ANONYMIZER_API_BASE": gateway.upstream_url,
    }
    with owned_proxy(gateway, tmp_path, environment, config=config) as candidate:
        _assert_registered(candidate, GUARDRAIL_CONSUMERS)