litellm/tests/integration/streaming/test_decision_model_buffering.py

import json
import uuid
from pathlib import Path
from typing import Final, Literal

import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value, string_value
from integration._support.process import owned_proxy
from integration._support.upstream import delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse, SseResponse
from pydantic import JsonValue


@pytest.mark.parametrize("endpoint", ["chat/completions", "messages", "responses"])
@pytest.mark.parametrize(
    ("buffering", "probability", "judge_status", "fallback"),
    [
        (True, 1.0, 200, "fail_closed"),
        (True, 0.0, 200, "fail_closed"),
        (False, 1.0, 200, "fail_closed"),
        (None, 1.0, 200, "fail_closed"),
        (True, 0.0, 500, "fail_closed"),
        (True, 0.0, 500, "fail_open"),
    ],
    ids=["buffered-block", "buffered-allow", "unbuffered-block", "default-block", "outage-closed", "outage-open"],
)
def test_decision_model_streaming_buffering(
    gateway: Gateway,
    tmp_path: Path,
    buffering: bool | None,
    probability: float,
    endpoint: str,
    judge_status: int,
    fallback: Literal["fail_closed", "fail_open"],
) -> None:
    identity: Final = uuid.uuid4().hex
    frames: Final = _frames(identity, endpoint)
    chat: Final = register_scenario(
        identity,
        SseResponse(content_type="text/event-stream", frames=frames),
        control_url=gateway.upstream_url,
    )
    try:
        decider: Final = register_scenario(
            identity + "-decider",
            JsonResponse(
                content_type="application/json",
                status=judge_status,
                body={
                    "model": "scripted-decider",
                    "answers": [{"type": "predicate", "name": "bananas", "probability": probability}],
                    "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
                },
            ),
            control_url=gateway.upstream_url,
        )
        try:
            config: Final = {
                **JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())),
                "model_list": [
                    {
                        "model_name": "buffering-chat",
                        "litellm_params": {
                            "model": "anthropic/scripted-chat" if endpoint == "messages" else "openai/scripted-chat",
                            "api_key": "synthetic",
                            "api_base": chat.api_base(),
                        },
                    },
                    {
                        "model_name": "buffering-decider",
                        "litellm_params": {
                            "model": "openai/scripted-decider",
                            "api_key": "synthetic",
                            "api_base": decider.api_base(),
                        },
                    },
                ],
                "guardrails": [
                    {
                        "guardrail_name": "banana-policy",
                        "litellm_params": {
                            "guardrail": "decision_model",
                            "mode": "post_call",
                            "default_on": True,
                            "decision_model": "buffering-decider",
                            "unreachable_fallback": fallback,
                            **({"streaming_buffer_until_moderated": buffering} if buffering is not None else {}),
                            "checks": [
                                {
                                    "name": "bananas",
                                    "instructions": "Does the text mention bananas?",
                                    "action": "block",
                                    "threshold": 0.5,
                                }
                            ],
                        },
                    }
                ],
                "general_settings": {"master_key": gateway.key},
            }
            config_path: Final = tmp_path / "decision-buffering.yaml"
            config_path.write_text(yaml.safe_dump(config))
            with owned_proxy(gateway, tmp_path, {}, config=config_path, workers=2) as candidate:
                response: Final = candidate.request(
                    "POST",
                    f"/v1/{endpoint}",
                    {
                        "model": "buffering-chat",
                        **(
                            {"input": "Tell me about bananas"}
                            if endpoint == "responses"
                            else {"messages": [{"role": "user", "content": "Tell me about bananas"}]}
                        ),
                        **({"max_tokens": 100} if endpoint == "messages" else {}),
                        "stream": True,
                        "cache": {"no-cache": True, "no-store": True},
                    },
                )
                events: Final = tuple(
                    JSON_OBJECT.validate_json(line.removeprefix("data: "))
                    for line in response.text.splitlines()
                    if line.startswith("data: {")
                )
                content: Final = "".join(_event_text(event, endpoint) for event in events)
                if judge_status != 200 and fallback == "fail_closed":
                    assert "Decision model guardrail failed to respond" in response.text, response.text
                    assert content == "", response.text
                elif probability == 0:
                    assert content == "Bananas are fruit.", response.text
                    assert '"error":' not in response.text, response.text
                else:
                    assert "Violated decision model guardrail policy" in response.text, response.text
                    if buffering:
                        assert content == "", response.text
                    else:
                        assert content and "Bananas are fruit.".startswith(content), response.text
        finally:
            delete_scenario(decider)
    finally:
        delete_scenario(chat)


def _event_text(event: dict[str, JsonValue], endpoint: str) -> str:
    if endpoint == "responses":
        return string_value(event.get("delta", "")) if event.get("type") == "response.output_text.delta" else ""
    if endpoint == "messages":
        return string_value(object_value(event.get("delta", {})).get("text", ""))
    choices: Final = list_value(event.get("choices", [{}]))
    return string_value(object_value(object_value(choices[0]).get("delta", {})).get("content", ""))


def _frames(identity: str, endpoint: str) -> tuple[str, ...]:
    pieces: Final = ("Ban", "anas", " are", " fruit", ".")
    if endpoint == "messages":
        events: Final = (
            {
                "type": "message_start",
                "message": {
                    "id": identity,
                    "type": "message",
                    "role": "assistant",
                    "content": [],
                    "model": "scripted-chat",
                    "stop_reason": None,
                    "usage": {"input_tokens": 1},
                },
            },
            {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
            *(
                {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}
                for text in pieces
            ),
            {"type": "content_block_stop", "index": 0},
            {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}},
            {"type": "message_stop"},
        )
        return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in events)
    if endpoint == "responses":
        message: Final = {
            "id": "msg_" + identity,
            "type": "message",
            "role": "assistant",
            "status": "completed",
            "content": [{"type": "output_text", "text": "".join(pieces), "annotations": []}],
        }
        response: Final = {
            "id": "resp_" + identity,
            "object": "response",
            "created_at": 1,
            "status": "completed",
            "model": "scripted-chat",
            "output": [message],
            "usage": {"input_tokens": 1, "output_tokens": 5, "total_tokens": 6},
        }
        response_events: Final = (
            {
                "type": "response.created",
                "sequence_number": 0,
                "response": {**response, "status": "in_progress", "output": []},
            },
            {
                "type": "response.output_item.added",
                "sequence_number": 1,
                "output_index": 0,
                "item": {**message, "content": []},
            },
            *(
                {
                    "type": "response.output_text.delta",
                    "sequence_number": index + 2,
                    "item_id": message["id"],
                    "output_index": 0,
                    "content_index": 0,
                    "delta": text,
                }
                for index, text in enumerate(pieces)
            ),
            {"type": "response.output_item.done", "sequence_number": 7, "output_index": 0, "item": message},
            {"type": "response.completed", "sequence_number": 8, "response": response},
        )
        return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in response_events)
    return (
        *(
            f"data: {
                json.dumps(
                    {
                        'id': identity,
                        'object': 'chat.completion.chunk',
                        'created': 1,
                        'model': 'scripted-chat',
                        'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}],
                    }
                )
            }"
            for text in pieces
        ),
        "data: [DONE]",
    )