litellm/tests/unit/proxy/client/test_chat.py

import importlib
import importlib.util
from importlib.machinery import PathFinder
import json
import time
import site
import sys
from typing import Final

import pytest
import requests
from pydantic import JsonValue, ValidationError
from litellm.proxy.client.chat import ChatClient
from litellm.proxy.client.exceptions import StreamError, UnauthorizedError


def _load_http_mocking_responses():
    """Load the third-party `responses` package even if test collection creates
    a top-level `responses` namespace package from `tests/unit/responses`.
    """
    module = importlib.import_module("responses")
    if hasattr(module, "activate"):
        return module

    for module_name in list(sys.modules):
        if module_name == "responses" or module_name.startswith("responses."):
            sys.modules.pop(module_name, None)

    search_paths = []
    try:
        search_paths.extend(site.getsitepackages())
    except AttributeError:
        pass
    user_site = site.getusersitepackages()
    if isinstance(user_site, str):
        search_paths.append(user_site)
    else:
        search_paths.extend(user_site)

    spec = PathFinder.find_spec("responses", search_paths)
    if spec is None or spec.loader is None:
        raise ImportError("Unable to load the third-party `responses` package")
    module = importlib.util.module_from_spec(spec)
    sys.modules["responses"] = module
    spec.loader.exec_module(module)

    if not hasattr(module, "activate"):
        raise ImportError("Unable to load the third-party `responses` package")
    return module


responses = _load_http_mocking_responses()


@pytest.fixture
def base_url():
    return "http://localhost:8000"


@pytest.fixture
def api_key():
    return "test-api-key"


@pytest.fixture
def client(base_url, api_key):
    return ChatClient(base_url=base_url, api_key=api_key)


@pytest.fixture
def sample_messages():
    return [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Name 3 countries"},
    ]


def test_client_initialization(base_url, api_key):
    """Test that the ChatClient is properly initialized"""
    client = ChatClient(base_url=base_url, api_key=api_key)

    assert client._base_url == base_url
    assert client._api_key == api_key


def test_client_initialization_strips_trailing_slash():
    """Test that the client properly strips trailing slashes from base_url during initialization"""
    base_url = "http://localhost:8000/////"
    client = ChatClient(base_url=base_url)

    assert client._base_url == "http://localhost:8000"


def test_client_without_api_key(base_url):
    """Test that the client works without an API key"""
    client = ChatClient(base_url=base_url)

    assert client._api_key is None


def test_completions_request_creation(client, base_url, api_key, sample_messages):
    """Test that completions creates a request with correct URL, headers, and body"""
    request = client.completions(
        model="gpt-4",
        messages=sample_messages,
        temperature=0.7,
        max_tokens=100,
        return_request=True,
    )

    # Check request method and URL
    assert request.method == "POST"
    assert request.url == f"{base_url}/chat/completions"

    # Check headers
    assert request.headers["Content-Type"] == "application/json"
    assert request.headers["Authorization"] == f"Bearer {api_key}"

    # Check request body
    assert request.json == {
        "model": "gpt-4",
        "messages": sample_messages,
        "temperature": 0.7,
        "max_tokens": 100,
    }


def test_completions_minimal_request(client, sample_messages):
    """Test that completions works with only required parameters"""
    request = client.completions(
        model="gpt-4", messages=sample_messages, return_request=True
    )

    # Check request body has only required fields
    assert request.json == {"model": "gpt-4", "messages": sample_messages}


def test_completions_all_parameters(client, sample_messages):
    """Test that completions accepts all optional parameters"""
    request = client.completions(
        model="gpt-4",
        messages=sample_messages,
        temperature=0.7,
        top_p=0.9,
        n=2,
        max_tokens=100,
        presence_penalty=0.5,
        frequency_penalty=-0.5,
        user="test-user",
        return_request=True,
    )

    # Check all parameters are included in request body
    assert request.json == {
        "model": "gpt-4",
        "messages": sample_messages,
        "temperature": 0.7,
        "top_p": 0.9,
        "n": 2,
        "max_tokens": 100,
        "presence_penalty": 0.5,
        "frequency_penalty": -0.5,
        "user": "test-user",
    }


@responses.activate
def test_completions_mock_response(client, sample_messages):
    """Test completions with a mocked successful response"""
    mock_response = {
        "id": "chatcmpl-123",
        "object": "chat.completion",
        "created": 1677858242,
        "model": "gpt-4",
        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": "Hello! How can I help you today?",
                },
                "finish_reason": "stop",
                "index": 0,
            }
        ],
    }

    # Mock the POST request
    responses.add(
        responses.POST,
        f"{client._base_url}/chat/completions",
        json=mock_response,
        status=200,
    )

    response = client.completions(model="gpt-4", messages=sample_messages)

    assert response == mock_response
    assert (
        response["choices"][0]["message"]["content"]
        == "Hello! How can I help you today?"
    )


@responses.activate
def test_completions_unauthorized_error(client, sample_messages):
    """Test that completions raises UnauthorizedError for 401 responses"""
    # Mock a 401 response
    responses.add(
        responses.POST,
        f"{client._base_url}/chat/completions",
        status=401,
        json={"error": "Unauthorized"},
    )

    with pytest.raises(UnauthorizedError):
        client.completions(model="gpt-4", messages=sample_messages)


@responses.activate
def test_completions_other_errors(client, sample_messages):
    """Test that completions raises HTTPError for other error responses"""
    # Mock a 500 response
    responses.add(
        responses.POST,
        f"{client._base_url}/chat/completions",
        status=500,
        json={"error": "Internal Server Error"},
    )

    with pytest.raises(requests.exceptions.HTTPError) as exc_info:
        client.completions(model="gpt-4", messages=sample_messages)
    assert exc_info.value.response.status_code == 500


def test_completions_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
    """
    A proxy that accepts the connection but never answers used to pin the caller's
    process forever, since the request carried no timeout at all.
    """
    client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1)

    started = time.monotonic()
    with pytest.raises(requests.exceptions.Timeout):
        client.completions(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}])

    assert time.monotonic() - started < 10


def test_completions_stream_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
    """
    The streaming call opens the response before reading chunks, so a proxy that never
    sends its headers used to hang here forever too.
    """
    client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1)

    started = time.monotonic()
    with pytest.raises(requests.exceptions.Timeout):
        next(client.completions_stream(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}]))

    assert time.monotonic() - started < 10


@pytest.mark.parametrize(
    "body, expected",
    [
        (
            b'data: {"choices": [{"delta": {"content": "Hel"}}]}\n\n'
            b'data: {"choices": [{"delta": {"content": "lo"}}]}\n\n'
            b"data: [DONE]\n\n",
            [{"choices": [{"delta": {"content": "Hel"}}]}, {"choices": [{"delta": {"content": "lo"}}]}],
        ),
        (
            b': keep-alive\n\nevent: ping\n\ndata: not json\n\ndata: {"id": "caf\xc3\xa9"}\n\ndata: 7\n\n',
            [{"id": "café"}, 7],
        ),
        (b'data: {"id": 1}\n\ndata:  [DONE] \n\ndata: {"id": 2}\n\n', [{"id": 1}]),
        (b"", []),
    ],
)
@responses.activate
def test_completions_stream_yields_parsed_sse_chunks(client, base_url, sample_messages, body, expected):
    responses.add(responses.POST, f"{base_url}/chat/completions", body=body, status=200)

    assert list(client.completions_stream(model="gpt-4", messages=sample_messages)) == expected


@pytest.mark.parametrize(
    "error_frame, error, message",
    [
        (
            b'data: {"error": {"message": "blocked by guardrail", "type": "None", "code": "400"}}',
            {"message": "blocked by guardrail", "type": "None", "code": "400"},
            "blocked by guardrail",
        ),
        (b'data: {"error": "upstream cut"}', "upstream cut", "upstream cut"),
        (b'data: {"error": {"code": "500"}}', {"code": "500"}, "{'code': '500'}"),
    ],
)
@responses.activate
def test_completions_stream_raises_on_an_error_frame_after_partial_content(
    client: ChatClient,
    base_url: str,
    sample_messages: list[dict[str, str]],
    error_frame: bytes,
    error: JsonValue,
    message: str,
) -> None:
    body: Final = b'data: {"choices": [{"delta": {"content": "Hel"}}]}\n\n' + error_frame + b"\n\ndata: [DONE]\n\n"
    responses.add(responses.POST, f"{base_url}/chat/completions", body=body, status=200)
    stream: Final = client.completions_stream(model="gpt-4", messages=sample_messages)

    assert next(stream) == {"choices": [{"delta": {"content": "Hel"}}]}
    with pytest.raises(StreamError) as raised:
        next(stream)
    assert raised.value.error == error
    assert str(raised.value) == message


@pytest.mark.parametrize("error", [None, False, 0, "", {}])
@responses.activate
def test_completions_stream_yields_a_chunk_whose_error_is_empty(
    client: ChatClient, base_url: str, sample_messages: list[dict[str, str]], error: JsonValue
) -> None:
    chunk: Final = {"choices": [{"delta": {"content": "Hel"}}], "error": error}
    body: Final = b"data: " + json.dumps(chunk).encode() + b"\n\ndata: [DONE]\n\n"
    responses.add(responses.POST, f"{base_url}/chat/completions", body=body, status=200)

    assert list(client.completions_stream(model="gpt-4", messages=sample_messages)) == [chunk]


def test_stream_error_redacts_a_key_in_its_message() -> None:
    key: Final = "sk-" + "a1b2c3d4" * 4

    error: Final = StreamError({"message": f"upstream rejected {key}"})

    assert key not in str(error)
    assert str(error).startswith("upstream rejected ")


@responses.activate
def test_completions_stream_rejects_bytes_that_are_not_utf8(client, base_url, sample_messages):
    responses.add(responses.POST, f"{base_url}/chat/completions", body=b"data: \xff\n\n", status=200)

    with pytest.raises(UnicodeDecodeError):
        list(client.completions_stream(model="gpt-4", messages=sample_messages))


@pytest.mark.parametrize("line", ['data: {"id": 1}', 5, memoryview(b'data: {"id": 1}')])
@responses.activate
def test_completions_stream_rejects_lines_that_are_not_bytes(client, base_url, sample_messages, monkeypatch, line):
    responses.add(responses.POST, f"{base_url}/chat/completions", body=b"", status=200)
    monkeypatch.setattr(requests.Response, "iter_lines", lambda self: iter([line]))

    with pytest.raises(ValidationError):
        list(client.completions_stream(model="gpt-4", messages=sample_messages))


@responses.activate
def test_completions_stream_accepts_bytearray_lines(client, base_url, sample_messages, monkeypatch):
    responses.add(responses.POST, f"{base_url}/chat/completions", body=b"", status=200)
    monkeypatch.setattr(requests.Response, "iter_lines", lambda self: iter([bytearray(b'data: {"id": 1}')]))

    assert list(client.completions_stream(model="gpt-4", messages=sample_messages)) == [{"id": 1}]