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

import importlib
import importlib.util
from importlib.machinery import PathFinder
import time
import site
import sys

import pytest
import requests
from pydantic import ValidationError
from litellm.proxy.client.chat import ChatClient
from litellm.proxy.client.exceptions import 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


@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}]