ci(unit): cut unit check time with cheaper per-test resets and fewer non-unit tests (#51864)

* test(unit): reset module state only when a test changed it

The proxy conftest deep-copied every snapshotted litellm attribute before each
test, and the copy of litellm.model_cost alone was a median 0.12 s for each of
56,031 proxy tests on the runner. The reset now copies an attribute back only
when it no longer equals the snapshot.

The enterprise and MCP server conftests reloaded litellm before every test,
3.27 s and 0.21 s per test. They now use the same reset.

* test(unit): remove six flaky tests

They failed on main without a code change behind them between 2026-11-05 and
2026-11-14: the vector store hook SDK fallback, the concurrent user config TPM
count, three Prometheus metrics server port tests and the OpenCode list cost
billing test.

* ci: run four heavy unit shards on 16-core runners and keep two smoke jobs off pull requests

proxy-feature-endpoints, proxy-endpoints, root and llms-providers ran 1.9 to 2.2
times faster with 8 workers on a 16-core runner in a benchmark run.

windows-sdk and proxy-build-from-pip still run on pushes to main and manual
runs. proxy-build-from-pip installs a published release and never runs the
pull request's code.

* test: move 109 tests that start a process or open a socket into tests/integration

tests/unit is in-process by its own rules, and these tests spawn a Python
interpreter, a shell or a local TCP or HTTP server. They took 921 s of unit
shard time per run, 396 s of it in the MCP client transport tests. Each one now
lives in a *_real_io.py file under the integration directory for its domain and
runs with that directory's group.

The tests themselves are unchanged apart from one repository-root path in the
CLI test.

* test(unit): copy back only the dict entries a test changed

A test module that registers models at import leaves litellm.model_cost four
entries larger than the snapshot before every test in its directory, so the
reset still deep-copied the whole price map for each of them. It now rebuilds a
changed dict from the live entries that still match and copies only the rest.

* ci(smoke): keep the two jobs that skip on pull requests out of the smoke gate

A skipped job passing the gate is what test_unit_passed_gate guards against,
so the gate goes back to requiring success from every job it needs and
windows-sdk and proxy-build-from-pip stop being part of it. They still run
and report on pushes to main and manual runs.
This commit is contained in:
yuneng-jiang 2026-11-14 15:41:30 -08:00 • committed by GitHub
parent ac23812b4b
commit da55893656
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
68 changed files with 5143 additions and 3821 deletions

View file

@ -178,6 +178,7 @@ jobs:
windows-sdk:
name: windows-sdk
if: github.event_name != 'pull_request'
runs-on: windows-2025
timeout-minutes: 30
env:
@ -213,6 +214,7 @@ jobs:
proxy-build-from-pip:
name: proxy-build-from-pip
if: github.event_name != 'pull_request'
runs-on: ubuntu-24.04
timeout-minutes: 10
services:
@ -352,7 +354,7 @@ jobs:
smoke-passed:
name: smoke passed
needs: [dashboard-build, core-checks, all-extras, windows-sdk, proxy-build-from-pip]
needs: [dashboard-build, core-checks, all-extras]
if: always()
runs-on: ubuntu-latest
timeout-minutes: 2

View file

@ -126,7 +126,7 @@ jobs:
name: ${{ matrix.shard }}
needs: [rust-bridge, assert-shard-coverage]
if: ${{ !cancelled() }}
runs-on: ubuntu-latest
runs-on: ${{ matrix.runner || 'ubuntu-latest' }}
timeout-minutes: ${{ matrix.job-timeout-minutes }}
permissions:
contents: read
@ -220,6 +220,7 @@ jobs:
dist: loadscope
- shard: llms-providers
runner: ubuntu-latest-16-cores
test-path: |-
tests/unit/llms/test_api_base_trailing_slash_routes.py
tests/unit/llms/test_cache_control_and_reasoning.py
@ -380,7 +381,7 @@ jobs:
tests/unit/llms/zai
--ignore=tests/unit/llms/base_llm/batches/base_batches_config_test.py
python-version: "3.12"
workers: 4
workers: 8
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
@ -398,10 +399,11 @@ jobs:
dist: loadscope
- shard: root
runner: ubuntu-latest-16-cores
test-path: >-
tests/unit/test_*.py
python-version: "3.12"
workers: 4
workers: 8
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
@ -522,6 +524,7 @@ jobs:
dist: loadscope
- shard: proxy-endpoints
runner: ubuntu-latest-16-cores
test-path: |-
tests/unit/proxy/management_endpoints
tests/unit/proxy/management_helpers
@ -539,13 +542,14 @@ jobs:
--ignore=tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py
--ignore=tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
python-version: "3.12"
workers: 4
workers: 8
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
dist: loadscope
- shard: proxy-feature-endpoints
runner: ubuntu-latest-16-cores
test-path: |-
tests/unit/proxy/guardrails
--ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py
@ -578,7 +582,7 @@ jobs:
tests/unit/proxy/config_resolvers
tests/unit/proxy/utils
python-version: "3.12"
workers: 4
workers: 8
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60

107
tests/_module_state.py Normal file
View file

@ -0,0 +1,107 @@
import copy
import inspect
import itertools
import warnings
from collections.abc import Mapping
from types import ModuleType
from typing import Final
_SNAPSHOT_TYPES: Final = (list, dict, set, tuple, str, int, float, bool, bytes)
_MUTABLE_CONTAINERS: Final = (list, dict, set, bytearray)
_ABSENT: Final = object()
def _holds_mutable_container(value: object) -> bool:
if isinstance(value, _MUTABLE_CONTAINERS):
return True
if isinstance(value, tuple):
return any(_holds_mutable_container(element) for element in value)
return False
def _restored_value(value: object) -> object:
return copy.deepcopy(value) if _holds_mutable_container(value) else value
def _readable_attributes(module: ModuleType) -> tuple[tuple[str, object], ...]:
return tuple((attr, value) for attr, value in tuple(vars(module).items()) if not attr.startswith("_"))
def _snapshot_of(module: ModuleType, attr: str, value: object) -> object:
try:
return _restored_value(value)
except Exception as exc:
warnings.warn(f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", stacklevel=2)
return _ABSENT
def snapshot_mutable_state(module: ModuleType) -> Mapping[str, object]:
candidates: Final = (
(attr, _snapshot_of(module, attr, value))
for attr, value in _readable_attributes(module)
if value is None or isinstance(value, _SNAPSHOT_TYPES)
)
return {attr: value for attr, value in candidates if value is not _ABSENT}
def _still_matches(current: object, default: object) -> bool:
if current is default:
return True
if type(current) is not type(default):
return False
try:
return bool(current == default)
except Exception:
return False
def _restored_entries(current: dict[object, object], default: dict[object, object]) -> dict[object, object]:
return {
key: current[key]
if key in current and type(current[key]) is type(value) and current[key] == value
else _restored_value(value)
for key, value in default.items()
}
def _restored_from(current: object, default: object) -> object:
if type(current) is dict and type(default) is dict:
return _restored_entries(current, default)
return _restored_value(default)
def restore_mutable_state(module: ModuleType, snapshot: Mapping[str, object]) -> None:
live: Final = vars(module)
for attr, default in snapshot.items():
current = live.get(attr, _ABSENT)
if _still_matches(current, default):
continue
try:
setattr(module, attr, _restored_from(current, default))
except Exception as exc:
warnings.warn(f"conftest: could not restore {module.__name__}.{attr}: {exc}", stacklevel=2)
def _is_flushable(value: object) -> bool:
if inspect.isclass(value) or inspect.ismodule(value):
return False
return callable(getattr(value, "flush_cache", None))
def _flushable_attributes(module: ModuleType) -> tuple[tuple[ModuleType, str], ...]:
return tuple((module, attr) for attr, value in _readable_attributes(module) if _is_flushable(value))
def collect_flushable_caches(modules: tuple[ModuleType, ...]) -> tuple[tuple[ModuleType, str], ...]:
return tuple(itertools.chain.from_iterable(_flushable_attributes(module) for module in modules))
def flush_caches(targets: tuple[tuple[ModuleType, str], ...]) -> None:
for module, attr in targets:
flush = getattr(vars(module).get(attr), "flush_cache", None)
if not callable(flush):
continue
try:
flush()
except Exception as exc:
warnings.warn(f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", stacklevel=2)

View file

@ -11,7 +11,7 @@ _WORKFLOWS_DIR: Final = Path(__file__).resolve().parents[2] / ".github" / "workf
_TIERS: Final = (
("test-linting.yml", "lint-passed", frozenset()),
("test-unit.yml", "unit-passed", frozenset({"coverage"})),
("test-merge-smoke.yml", "smoke-passed", frozenset()),
("test-merge-smoke.yml", "smoke-passed", frozenset({"windows-sdk", "proxy-build-from-pip"})),
)
_Tier = tuple[str, str, frozenset[str]]

View file

@ -26,6 +26,10 @@ directory that a `GROUPS` entry in `run.py` selects. There is no manifest and no
tests. A product bug the test exposes is `pytest.skip("BUG: <symptom>")` at the top of the body, not a
fix in the test and not a deletion. Needs no network, proxy, DB or Redis: `tests/unit`
A test that starts its own process or opens its own socket but needs no proxy also lives here, in a
`*_real_io.py` file next to the domain it covers, because `tests/unit` is in-process. It asks for none of
the gateway fixtures, so it runs as plain pytest inside its group
GitHub Actions integration groups run every file here except the OCI and sandbox tests, which no CI runs, and
the files in `GITHUB_FILES` in `run.py`. `test-integration.yml` runs the same groups through the same
`run_integration.sh` on GitHub Actions, on main every three hours from 8am to 8pm Pacific Monday to Saturday,

View file

@ -0,0 +1,124 @@
import re
import shutil
import subprocess
import sys
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import pytest
from litellm.proxy.auth.master_key_boot_check import (
GENERATE_MASTER_KEY_COMMAND,
MASTER_KEY_ENV_VAR,
MIGRATE_FROM_MASTER_KEY_ENV_VAR,
PRINT_NEW_MASTER_KEY_COMMAND,
ConfigFileSource,
EnvironmentSource,
MasterKeyBootVerdict,
SafeMasterKey,
StoredSecretsMigration,
UnsafeMasterKeyReason,
UnsafeMasterKeyRefused,
announce_on_stderr_at_exit,
master_key_boot_verdict,
render_refusal,
)
PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234"
def _verdict(
master_key: str | None,
general_settings: Mapping[str, object] | None = None,
*,
environment_master_key: str | None = None,
config_file_path: str | None = None,
override_env_is_on: bool = False,
salt_key_is_set: bool = False,
database_is_configured: bool = False,
) -> MasterKeyBootVerdict:
return master_key_boot_verdict(
master_key=master_key,
environment_master_key=environment_master_key,
general_settings=general_settings or {},
config_file_path=config_file_path,
override_env_is_on=override_env_is_on,
salt_key_is_set=salt_key_is_set,
database_is_configured=database_is_configured,
)
def _refusal(
reason: UnsafeMasterKeyReason = UnsafeMasterKeyReason.PUBLICLY_KNOWN,
source: ConfigFileSource | EnvironmentSource = EnvironmentSource(),
environment_variable_is_set: bool = False,
migration: StoredSecretsMigration | None = None,
) -> UnsafeMasterKeyRefused:
return UnsafeMasterKeyRefused(
reason=reason,
source=source,
environment_variable_is_set=environment_variable_is_set,
migration=migration,
)
_MIGRATION = StoredSecretsMigration(from_master_key=PUBLICLY_KNOWN_KEY, encrypted_value_count=3)
@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl")
@pytest.mark.parametrize("from_master_key", [PUBLICLY_KNOWN_KEY, "", " " + PUBLICLY_KNOWN_KEY])
def test_printed_migration_commands_save_both_keys_to_the_env_file(tmp_path: Path, from_master_key: str):
from dotenv import dotenv_values
text = render_refusal(
_refusal(migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1))
)
commands = [line.strip() for line in text.splitlines() if line.strip().startswith("echo ")]
subprocess.run(["bash", "-c", "\n".join(commands)], cwd=tmp_path, capture_output=True, text=True, check=True)
saved = dotenv_values(tmp_path / ".env")
assert saved[MIGRATE_FROM_MASTER_KEY_ENV_VAR] == from_master_key
assert _verdict(saved[MASTER_KEY_ENV_VAR]) == SafeMasterKey()
assert len(commands) == 2
@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl")
def test_printed_command_saves_a_key_the_boot_check_accepts(tmp_path: Path):
completed = subprocess.run(
["bash", "-c", GENERATE_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True
)
saved = (tmp_path / ".env").read_text()
match = re.fullmatch(rf"{MASTER_KEY_ENV_VAR}=(sk-[0-9a-f]{{64}})\n", saved)
assert match is not None
assert completed.stdout == saved
assert _verdict(match.group(1)) == SafeMasterKey()
@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl")
def test_rotation_command_prints_a_key_the_boot_check_accepts_and_saves_nothing(tmp_path: Path):
completed = subprocess.run(
["bash", "-c", PRINT_NEW_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True
)
assert re.fullmatch(r"sk-[0-9a-f]{64}\n", completed.stdout) is not None
assert _verdict(completed.stdout.strip()) == SafeMasterKey()
assert list(tmp_path.iterdir()) == []
def test_announced_fix_is_the_last_thing_a_crashing_process_prints():
crash_after_announcing = (
"from litellm.proxy.auth.master_key_boot_check import announce_on_stderr_at_exit\n"
"announce_on_stderr_at_exit('THE FIX')\n"
"print('buffered stdout')\n"
"raise RuntimeError('lifespan failed')\n"
)
completed = subprocess.run([sys.executable, "-c", crash_after_announcing], capture_output=True, text=True)
assert completed.returncode != 0
assert "RuntimeError: lifespan failed" in completed.stderr
assert completed.stderr.endswith("THE FIX")
assert completed.stdout == "buffered stdout\n"

View file

@ -0,0 +1,155 @@
import json
import os
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from textwrap import dedent
from types import SimpleNamespace
from typing import Final
from unittest.mock import ANY
import litellm
import litellm.proxy.proxy_server
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.videos.utils import encode_video_id_with_provider
class _RoutingRequest:
def __init__(self, headers=None, query_params=None):
self.headers = headers or {}
self.query_params = query_params or {}
self.state = SimpleNamespace()
# ---------------------------------------------------------------------------
# _run_centralized_common_checks — centralized authz gate
# ---------------------------------------------------------------------------
_RECORDING_DDTRACE = dedent(
"""
import functools
import inspect
class _Span:
def __enter__(self):
return self
def __exit__(self, *exc):
return None
class _Tracer:
def __init__(self):
self.spans = []
def wrap(self, name=None, **kwargs):
def decorator(f):
span_name = name or f"{f.__module__}.{f.__name__}"
if inspect.iscoroutinefunction(f):
@functools.wraps(f)
async def async_wrapped(*args, **kw):
self.spans.append(span_name)
return await f(*args, **kw)
return async_wrapped
@functools.wraps(f)
def wrapped(*args, **kw):
self.spans.append(span_name)
return f(*args, **kw)
return wrapped
return decorator
def trace(self, name, **kwargs):
return _Span()
def current_span(self):
return None
def current_root_span(self):
return None
tracer = _Tracer()
"""
)
_DDTRACE_AUTH_PROBE = dedent(
"""
import asyncio
import json
import ddtrace
from fastapi import Request
from starlette.datastructures import URL
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
proxy_server.master_key = "sk-probe"
async def auth(api_key):
request = Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/chat/completions"})
request._url = URL(url="/chat/completions")
try:
await user_api_key_auth(
request=request,
api_key=api_key,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
custom_litellm_key_header=None,
)
return "accepted"
except ProxyException:
return "rejected"
async def main():
outcomes = [await auth("Bearer sk-probe"), await auth("Bearer sk-wrong")]
print(json.dumps({"outcomes": outcomes, "spans": ddtrace.tracer.spans}))
asyncio.run(main())
"""
)
def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(tmp_path: Path):
stub_root = tmp_path / "site"
(stub_root / "ddtrace").mkdir(parents=True)
(stub_root / "ddtrace" / "__init__.py").write_text(_RECORDING_DDTRACE)
probe = tmp_path / "probe.py"
probe.write_text(_DDTRACE_AUTH_PROBE)
repo_root = Path(litellm.__file__).resolve().parent.parent
env = {
**os.environ,
"USE_DDTRACE": "true",
"PYTHONPATH": os.pathsep.join(
[str(stub_root), str(repo_root)] + [p for p in (os.environ.get("PYTHONPATH"),) if p]
),
}
result = subprocess.run(
[sys.executable, str(probe)], env=env, cwd=repo_root, capture_output=True, text=True, check=False
)
assert result.returncode == 0, result.stderr[-4000:]
report = json.loads(result.stdout.strip().splitlines()[-1])
assert report["outcomes"] == ["accepted", "rejected"]
auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span]
_VIDEO_ID: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2")
_EXPIRED_BEFORE_ANY_TEST_RUN: Final = datetime(2000, 1, 1, tzinfo=timezone.utc)

View file

@ -0,0 +1,119 @@
"""
Tests for EvictedClientCloser.
An evicted client must stay open long enough for a request that already holds it
to finish, and must then actually be closed, otherwise its connection pool is
retained until a generational collection runs. A client the caller supplied is
never closed, because litellm does not own its lifecycle.
"""
import asyncio
import httpx
import pytest
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
class FakeClock:
"""Hand-advanced monotonic clock, so grace windows need no real waiting."""
def __init__(self) -> None:
self.now = 1000.0
def __call__(self) -> float:
return self.now
def advance(self, seconds: float) -> None:
self.now += seconds
def make_closer(clock: FakeClock, grace_seconds: float = 60.0) -> EvictedClientCloser:
return EvictedClientCloser(grace_seconds=grace_seconds, clock=clock)
async def _trickling_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
"""Serves a chunked body slowly, so a request stays on the wire long enough to observe."""
await reader.read(4096)
writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n")
await writer.drain()
for _ in range(6):
writer.write(b"5\r\nhello\r\n")
await writer.drain()
await asyncio.sleep(0.1)
writer.write(b"0\r\n\r\n")
await writer.drain()
@pytest.mark.asyncio
async def test_a_client_serving_a_request_is_not_closed_when_its_grace_window_ends():
"""The grace window on its own cannot promise that a request has finished.
``litellm.request_timeout`` defaults to 6000 seconds and a streaming response
is bounded only by how long the upstream keeps sending, so a client past its
deadline is closed only once its own pool reports nothing in flight.
"""
server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
clock = FakeClock()
closer = make_closer(clock)
client = httpx.AsyncClient()
closer.mark_owned(client)
closer.schedule(client)
async def read_the_stream() -> int:
received = 0
async with client.stream("GET", f"http://127.0.0.1:{port}/") as response:
async for chunk in response.aiter_bytes():
received += len(chunk)
return received
streaming = asyncio.create_task(read_the_stream())
await asyncio.sleep(0.25) # the request is on the wire
clock.advance(3600.0) # and its grace window is long gone
closer.reap()
await asyncio.sleep(0.05)
assert client.is_closed is False, "closed a client that was serving a request"
assert await streaming > 0, "the in-flight request did not survive the reap"
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.is_closed is True, "an idle client past its grace window must be closed"
assert closer.pending_count == 0
server.close()
@pytest.mark.asyncio
async def test_the_aiohttp_backed_handler_is_not_closed_mid_request():
"""The default async path is aiohttp-backed, whose pool accounts for its own leases."""
server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
clock = FakeClock()
closer = make_closer(clock)
handler = AsyncHTTPHandler()
held_client = handler.client
closer.mark_owned(handler)
closer.schedule(handler)
request = asyncio.create_task(handler.get(f"http://127.0.0.1:{port}/"))
await asyncio.sleep(0.25)
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert held_client.is_closed is False, "closed a handler that was serving a request"
assert (await request).status_code == 200
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert held_client.is_closed is True
assert handler.client.is_closed is False, "a held handler must self-heal after its evicted client is closed"
server.close()

View file

@ -0,0 +1,32 @@
import os
import subprocess
import sys
import textwrap
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
def test_importing_caching_does_not_require_redis():
# redis is an optional dependency (extra_proxy), so the base SDK can be
# installed without it. Selecting valkey-semantic needs redis, but merely
# importing litellm.caching.caching must not, or `import litellm` breaks for
# every base-SDK user. This runs in a subprocess with redis blocked so the
# check is not polluted by redis already being imported in this session.
code = textwrap.dedent("""
import sys
for name in ("redis", "redis.asyncio", "redis.commands",
"redis.commands.search"):
sys.modules[name] = None
import litellm.caching.caching # must not import redis at module top
from litellm.types.caching import LiteLLMCacheType
assert LiteLLMCacheType.VALKEY_SEMANTIC == "valkey-semantic"
print("ok")
""")
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
env={**os.environ, "PYTHONPATH": _REPO_ROOT},
)
assert result.returncode == 0, result.stderr
assert "ok" in result.stdout

View file

@ -0,0 +1,272 @@
import os
import pathlib
import shlex
import stat
import subprocess
import sys
import time
import pytest
from click.testing import CliRunner
from litellm.proxy.client.cli.commands.claude_settings import (
ApiKeyHelper,
KeepModel,
StartOn,
StaticToken,
UnpinModel,
install_statusline_script,
installed_executable,
lite_api_key_helper,
statusline_command,
)
class TestLiteApiKeyHelper:
@pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX")
def test_sh_runs_the_lite_on_path_against_this_proxy(self, tmp_path):
lite = tmp_path / "my tools" / "lite"
lite.parent.mkdir()
lite.write_text('#!/bin/sh\nprintf "%s\\n" "$0" "$@"\n')
lite.chmod(0o700)
helper = lite_api_key_helper("https://proxy.example.com/", platform="linux", which=lambda name: str(lite))
assert helper is not None
ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True, check=True)
assert ran.stdout.splitlines() == [str(lite), "--base-url", "https://proxy.example.com", "auth", "print-token"]
@pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX")
def test_a_lite_found_through_a_relative_path_entry_keeps_running_from_any_directory(self, tmp_path, monkeypatch):
for repo, name in ((tmp_path / "login-dir", "mine"), (tmp_path / "other-repo", "planted")):
lite = repo / "bin" / "lite"
lite.parent.mkdir(parents=True)
lite.write_text(f"#!/bin/sh\necho {name}\n")
lite.chmod(0o700)
monkeypatch.chdir(tmp_path / "login-dir")
helper = lite_api_key_helper("https://proxy.example.com", platform="linux", which=lambda name: "bin/lite")
assert helper is not None
ran = subprocess.run(
["sh", "-c", helper.command], cwd=tmp_path / "other-repo", capture_output=True, text=True, check=True
)
assert ran.stdout == "mine\n"
@pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX")
class TestInstalledExecutable:
@staticmethod
def _script(path, body):
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(f"#!/bin/sh\n{body}\n")
path.chmod(0o700)
def test_the_helper_runs_the_lite_installed_beside_this_python_over_a_shim_earlier_on_path(
self, tmp_path, monkeypatch
):
self._script(tmp_path / "shims" / "lite", "echo 'pyenv: lite: command not found' >&2; exit 127")
self._script(tmp_path / "venv" / "bin" / "lite", "echo real")
monkeypatch.setenv("PATH", f"{tmp_path / 'shims'}{os.pathsep}{os.environ['PATH']}")
helper = lite_api_key_helper(
"https://proxy.example.com",
platform="linux",
which=lambda name: installed_executable(name, scripts_dir=str(tmp_path / "venv" / "bin")),
)
assert helper is not None
ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True)
assert (ran.returncode, ran.stdout) == (0, "real\n")
PROXY = "http://127.0.0.1:4000"
ANTHROPIC = "https://api.anthropic.com"
RELOGIN = StaticToken("sk-fresh-login")
HELPER = ApiKeyHelper("/usr/local/bin/lite --base-url http://127.0.0.1:4000 auth print-token")
ORIGINAL = {
"theme": "dark",
"permissions": {"allow": ["Bash"]},
"env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "sk-ant-mine", "ANTHROPIC_BASE_URL": ANTHROPIC},
"apiKeyHelper": "/usr/local/bin/lite auth print-token",
"model": "claude-opus-5",
}
def _set(path, value):
"""A user edit: set (or with `_ABSENT`, remove) the key at a dotted path in the settings file."""
def edit(settings):
section, _, key = path.rpartition(".")
container = settings.setdefault(section, {}) if section else settings
if value is _ABSENT:
container.pop(key, None)
else:
container[key] = value
return settings
return edit
_ABSENT = object()
# Each row: initial file, steps (configure kwargs dicts or edit callables) between the first configure
# and unconfigure, the expected file afterwards, and the expected outcome fields. Sequences that used
# to be one test each; the receipt's rules are what make them all come out right.
UNDO_SCENARIOS = {
"plain round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}),
"helper round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}, {"credential": HELPER, "model": KeepModel()}),
"helper with no file before": (None, [], None, {"file_removed": True}, {"credential": HELPER}),
"re-logins switching between the token and the helper keep the first snapshot": (
ORIGINAL,
[{"credential": HELPER, "model": KeepModel()}, {"credential": RELOGIN, "model": KeepModel()}],
ORIGINAL,
{"kept": ()},
{"credential": HELPER, "model": KeepModel()},
),
"no file before": (None, [], None, {"file_removed": True}),
"no env before": ({"theme": "dark"}, [], {"theme": "dark"}, {}),
"null env before": ({"theme": "dark", "env": None}, [], {"theme": "dark", "env": None}, {}),
"empty env before": ({"theme": "dark", "env": {}}, [], {"theme": "dark", "env": {}}, {}),
"user edits stay and are named": (
ORIGINAL,
[_set("env.ENABLE_TOOL_SEARCH", "false"), _set("model", "claude-sonnet-4-6")],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "claude-sonnet-4-6"},
{"kept": {"env.ENABLE_TOOL_SEARCH", "model"}, "withheld": ()},
),
"user filled an env configure created": (None, [_set("env.MY_VAR", "mine")], {"env": {"MY_VAR": "mine"}}, {}),
"user deleted the file": (None, [lambda s: None], None, {"file_removed": True, "restored": (), "kept": ()}),
"user removed our key: neither restored nor kept": (
ORIGINAL,
[_set("env.ANTHROPIC_AUTH_TOKEN", _ABSENT)],
ORIGINAL,
{"not_restored": {"env.ANTHROPIC_AUTH_TOKEN"}, "kept": ()},
),
"restored names only what changed": (
{"model": "claude-opus-5"},
[],
{"model": "claude-opus-5"},
{
"restored": {
"env.ANTHROPIC_BASE_URL",
"env.ANTHROPIC_AUTH_TOKEN",
"env.ENABLE_TOOL_SEARCH",
"env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",
"statusLine",
},
"kept": (),
},
{"credential": RELOGIN, "model": KeepModel()},
),
"repeat across credential kinds keeps the first snapshot": (
ORIGINAL,
[
{"credential": RELOGIN, "model": UnpinModel()},
{"credential": StaticToken("sk-rotated"), "model": StartOn("claude-sonnet-4-6")},
],
ORIGINAL,
{},
),
"repeat without a model lets go of our pin, user had none": ({}, [{"model": UnpinModel()}], {}, {}),
"repeat without a model lets go of our pin, user had one": (
{"model": "claude-opus-5"},
[{"model": UnpinModel()}],
{"model": "claude-opus-5"},
{},
),
"re-login keeps our pin": (None, [{"credential": RELOGIN, "model": KeepModel()}], None, {"file_removed": True}),
"edit between configures survives an unpin repeat": (
ORIGINAL,
[
_set("model", "my-favourite"),
_set("env.ENABLE_TOOL_SEARCH", "false"),
{"credential": RELOGIN, "model": UnpinModel()},
],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"},
{"kept": {"env.ENABLE_TOOL_SEARCH", "model"}},
),
"edit between configures survives a re-login": (
ORIGINAL,
[
_set("model", "my-favourite"),
_set("env.ENABLE_TOOL_SEARCH", "false"),
{"credential": RELOGIN, "model": KeepModel()},
],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"},
{"kept": {"env.ENABLE_TOOL_SEARCH", "model"}},
),
"edit between configures: a same-model repeat displaces it, so it is what comes back": (
ORIGINAL,
[_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": RELOGIN}],
{**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"},
{"kept": {"env.ENABLE_TOOL_SEARCH"}, "restored_includes": {"model"}},
),
"base URL changed since: credentials withheld, receipt kept": (
ORIGINAL,
[_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")],
{**ORIGINAL, "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, "apiKeyHelper": _ABSENT},
{
"withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC), ("apiKeyHelper", ANTHROPIC)},
"kept": {"env.ANTHROPIC_BASE_URL"},
"receipt_kept": True,
},
),
"base URL changed and back: judged against the URL the restored file holds": (
ORIGINAL,
[_set("env.ANTHROPIC_BASE_URL", ANTHROPIC)],
ORIGINAL,
{"withheld": ()},
),
"credential captured beside no URL goes back only beside no URL": (
{"env": {"ANTHROPIC_API_KEY": "sk-default-endpoint"}},
[_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")],
{"env": {"ANTHROPIC_BASE_URL": "http://other-proxy:4000"}},
{
"withheld": {("env.ANTHROPIC_API_KEY", "no ANTHROPIC_BASE_URL (Anthropic's default endpoint)")},
"receipt_kept": True,
},
),
"restored document empty while a credential is withheld: file goes, receipt stays": (
None,
[
_set("env.ANTHROPIC_API_KEY", "sk-user"),
{"credential": RELOGIN, "model": KeepModel()},
_set("env.ANTHROPIC_BASE_URL", _ABSENT),
],
None,
{"withheld": {("env.ANTHROPIC_API_KEY", PROXY)}, "file_removed": True, "receipt_kept": True},
{"credential": RELOGIN, "model": KeepModel()},
),
"a credential the user changed is kept, never also withheld": (
ORIGINAL,
[_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000"), _set("apiKeyHelper", "/opt/mine/helper")],
{
**ORIGINAL,
"env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"},
"apiKeyHelper": "/opt/mine/helper",
},
{
"withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC)},
"kept": {"env.ANTHROPIC_BASE_URL", "apiKeyHelper"},
"receipt_kept": True,
},
),
}
class TestStatusLine:
"""Every configure registers the status line, and only while the slot is empty or already ours."""
COMMAND = "/opt/lite/bin/python /Users/me/.litellm/statusline.py"
def test_the_installed_script_is_the_bundled_one_and_the_command_runs_this_interpreter(self, tmp_path):
from litellm.proxy.client.cli.commands import statusline_script
script = tmp_path / "lite" / "statusline.py"
command = install_statusline_script(script)
assert script.read_bytes().split(b"\n", 1)[1] == pathlib.Path(statusline_script.__file__).read_bytes()
assert shlex.split(command) == [sys.executable, str(script)]
assert command == statusline_command(script)
assert stat.S_IMODE(script.stat().st_mode) == 0o600
assert stat.S_IMODE(script.parent.stat().st_mode) == 0o700
assert install_statusline_script(script) == command

View file

@ -0,0 +1,20 @@
import os
import subprocess
import sys
from typing import Final
def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks():
script: Final = (
"import litellm; "
"litellm.callbacks = []; "
"import litellm.proxy.hooks; "
"assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'"
)
result: Final = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"

View file

@ -0,0 +1,83 @@
import signal
import subprocess
import sys
from pathlib import Path
from types import MappingProxyType
from typing import Final
import anyio
import pytest
from litellm.proxy.rust_forward.platform import Platform, current_platform
from litellm.proxy.rust_forward.platform.windows import WindowsPlatform
from litellm.proxy.rust_forward.supervisor import supervise
from litellm.proxy.rust_forward.supervisor.children import ChildSpec
from litellm.proxy.rust_forward.supervisor.lifecycle import Child, ChildExited
from litellm.proxy.rust_forward.supervisor.runtime import (
SpawnFailed,
exit_like,
open_children,
spawn_subprocess,
)
from tests.unit.proxy.rust_forward.supervisor._support import python_server
PYTHON_CHILD: Final = """
import json, os, pathlib, socket, time
time.sleep(0.5)
pathlib.Path(os.environ["MARKERS"], "listening").touch()
private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"]
listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"]))
listener.listen()
time.sleep(10)
"""
RUST_CHILD: Final = """
import os, pathlib
markers = pathlib.Path(os.environ["MARKERS"])
(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists()))
"""
def _fake_gateway(scripts: Path) -> Platform:
if sys.platform != "win32":
gateway: Final = scripts / "litellm-gateway"
gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}")
gateway.chmod(0o755)
return current_platform()
script: Final = scripts / "rust_child.py"
script.write_text(RUST_CHILD)
(scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n')
return WindowsPlatform(gateway_name="litellm-gateway.cmd")
def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("MARKERS", str(tmp_path))
command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD)
supervise(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path))
assert (tmp_path / "rust_saw_python_listening").read_text() == "True"
POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="the child traps SIGTERM")
RUST_MISSING: Final = SpawnFailed(FileNotFoundError("litellm-gateway"))
def _children(argv: tuple[str, ...]) -> MappingProxyType[Child, ChildSpec]:
spec: Final = ChildSpec(argv, MappingProxyType({}), current_platform().spawn_options(()))
return MappingProxyType({Child.PYTHON: spec, Child.RUST: spec})
async def test_exiting_children_are_reported_as_events() -> None:
async with open_children(
current_platform(), spawn_subprocess, _children((sys.executable, "-c", "exit(7)")), 1
) as children:
assert await children.spawn(Child.RUST) is None
with anyio.fail_after(5):
assert await children.next_event() == ChildExited(Child.RUST, 7)
@POSIX_ONLY
@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGKILL])
def test_the_supervisor_dies_by_the_signal_that_killed_python(signum: signal.Signals) -> None:
script: Final = (
"from litellm.proxy.rust_forward.supervisor.runtime import exit_like\n"
f"raise SystemExit(exit_like({-signum}, frozenset()))"
)
assert subprocess.run((sys.executable, "-I", "-c", script), check=False, timeout=60).returncode == -signum

View file

@ -0,0 +1,88 @@
import subprocess
import sys
import textwrap
from types import MappingProxyType
from typing import Final
from mcp.types import Tool as MCPTool
from litellm.types.mcp_server.mcp_toolset import MCPToolset
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
def test_completion_with_function_tools_works_without_fastapi_installed():
script = textwrap.dedent(
"""
import sys
class _FastapiBlocker:
def find_spec(self, fullname, path=None, target=None):
if fullname == "fastapi" or fullname.startswith("fastapi."):
raise ModuleNotFoundError("No module named 'fastapi'")
return None
sys.meta_path.insert(0, _FastapiBlocker())
import litellm
response = litellm.completion(
model="openai/gpt-5.5",
messages=[{"role": "user", "content": "What is the weather in SF?"}],
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
],
mock_response="sunny",
)
assert response.choices[0].message.content == "sunny"
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
timeout=120,
)
assert result.returncode == 0, result.stderr
TOOLSET_ROWS: Final = (
MCPToolset(
toolset_id="ts-granted", toolset_name="team-toolset", tools=[{"server_id": "srv-1", "tool_name": "add"}]
),
MCPToolset(
toolset_id="ts-other", toolset_name="other-toolset", tools=[{"server_id": "srv-2", "tool_name": "lookup"}]
),
MCPToolset(toolset_id="ts-shouty", toolset_name="CALC", tools=[{"server_id": "srv-2", "tool_name": "lookup"}]),
MCPToolset(toolset_id="ts-group", toolset_name="reference", tools=[{"server_id": "srv-1", "tool_name": "add"}]),
MCPToolset(
toolset_id="ts-prefixed",
toolset_name="prefixed-toolset",
tools=[{"server_id": "srv-1", "tool_name": "add"}, {"server_id": "srv-1", "tool_name": "calc-subtract"}],
),
MCPToolset(
toolset_id="ts-subtract", toolset_name="subtract-toolset", tools=[{"server_id": "srv-1", "tool_name": "subtract"}]
),
)
GATEWAY_SERVERS: Final = (
MCPServer(server_id="srv-1", name="calc", alias="calc", transport=MCPTransport.http, access_groups=["wiki"]),
MCPServer(server_id="srv-2", name="docs", alias="docs", transport=MCPTransport.http, access_groups=["reference"]),
)
UPSTREAM_TOOLS: Final = MappingProxyType(
{
"calc": (
MCPTool(name="add", inputSchema={"type": "object"}),
MCPTool(name="subtract", inputSchema={"type": "object"}),
),
"docs": (MCPTool(name="lookup", inputSchema={"type": "object"}),),
}
)

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,44 @@
"""
Tests for OpenAPI to MCP generator, focusing on security and edge cases.
This test suite ensures that:
1. Parameter names with invalid Python identifiers are handled safely
2. No exec() is used (security)
3. All edge cases (hyphens, dots, keywords, special chars) work correctly
4. Path traversal attacks are prevented
5. Path parameters are properly URL encoded
"""
import pytest
if __name__ == "__main__":
pytest.main([__file__, "-v"])
# ---------------------------------------------------------------------------
# Tool name sanitization for OpenAPI -> MCP
# Repro: GitHub's REST OpenAPI uses tag-namespaced operationIds like
# "actions/download-job-logs-for-workflow-run". Without sanitization the
# generated MCP tool name contains '/', which Anthropic/OpenAI/Bedrock all
# reject (^[a-zA-Z0-9_-]+$). This block guards the registration + preview
# paths against that.
# ---------------------------------------------------------------------------
def test_openapi_generator_import_does_not_require_mcp_sdk() -> None:
import subprocess
import sys
script = """
import builtins
original_import = builtins.__import__
def without_mcp(name, *args, **kwargs):
if name == 'mcp' or name.startswith('mcp.'):
raise ModuleNotFoundError('MCP SDK unavailable')
return original_import(name, *args, **kwargs)
builtins.__import__ = without_mcp
import litellm.proxy.mcp_server.openapi_to_mcp_generator
"""
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True)
assert result.returncode == 0, result.stderr

View file

@ -0,0 +1,65 @@
import json
import os
import subprocess
import sys
import textwrap
from typing import Optional, Tuple
import pytest
import litellm
@pytest.fixture(autouse=True)
def _no_openai_api_base_override(monkeypatch):
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
monkeypatch.setattr(litellm, "api_base", None)
class TestAnthropicPromptCachingEnvVars:
"""Both settings are read from the environment at import, so an admin can enable
auto-caching without a config file. Each case re-imports litellm in a subprocess
so the env is read fresh without contaminating this process's module graph.
"""
@staticmethod
def _import_litellm_with_env(env_override: dict) -> Tuple[bool, Optional[str]]:
env = os.environ.copy()
env.pop("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", None)
env.pop("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL", None)
env.update(env_override)
script = textwrap.dedent(
"""
import json, litellm
print(json.dumps([litellm.enable_anthropic_prompt_caching, litellm.anthropic_prompt_caching_ttl]))
"""
)
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300)
assert result.returncode == 0, result.stderr
enabled, ttl = json.loads(result.stdout.strip().splitlines()[-1])
return enabled, ttl
def test_unset_env_leaves_auto_caching_off(self):
assert self._import_litellm_with_env({}) == (False, None)
@pytest.mark.parametrize("value", ["true", "True", "TRUE"])
def test_env_enables_auto_caching_case_insensitively(self, value):
enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value})
assert enabled is True
@pytest.mark.parametrize("value", ["false", "0", "yes", ""])
def test_env_only_enables_on_true(self, value):
enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value})
assert enabled is False
@pytest.mark.parametrize("value", ["5m", "1h"])
def test_ttl_env_is_applied(self, value):
_, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value})
assert ttl == value
@pytest.mark.parametrize("value", ["10m", "1H", "3600", "ephemeral"])
def test_unsupported_ttl_env_falls_back_to_provider_default(self, value):
"""An unparseable TTL must fall back to Anthropic's 5m default, never reach the provider verbatim."""
_, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value})
assert ttl is None

View file

@ -0,0 +1,41 @@
import json
import subprocess
import sys
from typing import Final
_MASKED_HOSTED_WEB_CALL_PROBE: Final = """
import json
from litellm.llms.anthropic.common_utils import write_request_hosted_web_call_arguments
message = {
"role": "assistant",
"provider_specific_fields": {
"web_search_calls": [
{
"id": "ws_srvtoolu_1",
"type": "web_search_call",
"status": "completed",
"action": {"type": "search", "query": "deploy token 12345", "queries": ["deploy token 12345"]},
}
]
},
}
write_request_hosted_web_call_arguments(message, {"srvtoolu_1": json.dumps({"query": "deploy token [REDACTED]"})})
print(json.dumps(message["provider_specific_fields"]["web_search_calls"]))
"""
def test_a_masked_replayed_hosted_web_call_serializes_before_its_openai_model_was_ever_built() -> None:
result: Final = subprocess.run(
(sys.executable, "-I", "-c", _MASKED_HOSTED_WEB_CALL_PROBE),
capture_output=True,
text=True,
check=True,
timeout=120,
)
hosted: Final = json.loads(result.stdout.strip().splitlines()[-1])
assert [call["action"] for call in hosted] == [
{"type": "search", "query": "deploy token [REDACTED]", "queries": ["deploy token [REDACTED]"]}
]

View file

@ -0,0 +1,204 @@
import os
from collections.abc import Mapping
from typing import Final, Optional
from unittest.mock import MagicMock
import openai
import requests
import litellm
from litellm.llms.azure.common_utils import BaseAzureLLM
def test_default_max_retries_env_var_reaches_azure_sdk_client():
import subprocess
import sys
code = (
"from litellm.llms.azure.common_utils import BaseAzureLLM\n"
"client = BaseAzureLLM().get_azure_openai_client("
"api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01',"
" client=None, _is_async=True, litellm_params={}, model='gpt-4')\n"
"print(client.max_retries)"
)
completed = subprocess.run(
[sys.executable, "-c", code],
env={**os.environ, "DEFAULT_MAX_RETRIES": "0"},
capture_output=True,
text=True,
check=True,
)
assert completed.stdout.strip() == "0"
# Test parameters for different API functions with Azure models
AZURE_API_FUNCTION_PARAMS = [
# (function_name, is_async, args)
(
"completion",
False,
{
"model": "azure/gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 10,
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
(
"completion",
True,
{
"model": "azure/gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 10,
"stream": True,
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
(
"embedding",
False,
{
"model": "azure/text-embedding-ada-002",
"input": "Hello world",
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
(
"embedding",
True,
{
"model": "azure/text-embedding-ada-002",
"input": "Hello world",
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
(
"speech",
False,
{
"model": "azure/tts",
"input": "Hello, this is a test of text to speech",
"voice": "alloy",
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
(
"speech",
True,
{
"model": "azure/tts-1",
"input": "Hello, this is a test of text to speech",
"voice": "alloy",
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
(
"transcription",
False,
{
"model": "azure/whisper-1",
"file": MagicMock(),
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
(
"transcription",
True,
{
"model": "azure/whisper-1",
"file": MagicMock(),
"api_key": "test-api-key",
"api_base": "https://test.openai.azure.com",
"api_version": "2023-05-15",
},
),
]
_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500"
def _entra_error(aadsts_code: str, description: str) -> Mapping[str, object]:
return {
"error": "invalid_client",
"error_description": f"{aadsts_code}: {description} Trace ID: {_ENTRA_TRACE_ID}",
"error_codes": [int(aadsts_code.removeprefix("AADSTS"))],
}
_REAL_CREDENTIAL_NON_REJECTIONS: Final = {
"unreachable_through_a_credential_chain": ("default_azure_credential_chain", None, ""),
"html_503_through_a_credential_chain": (
"default_azure_credential_chain",
503,
"<html><body>Service Unavailable</body></html>",
),
"json_503_through_a_credential_chain": (
"default_azure_credential_chain",
503,
_entra_error("AADSTS90033", "A transient error has occurred."),
),
"throttled_client_secret": (
"client_secret",
429,
_entra_error("AADSTS50196", "The server terminated an operation because it encountered a client request loop."),
),
"tenant_throttled_client_secret": (
"client_secret",
429,
_entra_error("AADSTS90055", "There are too many incoming requests."),
),
"request_budget_exceeded_client_secret": (
"client_secret",
400,
_entra_error("AADSTS90024", "A transient error has occurred. Try again."),
),
"retryable_error_client_secret": (
"client_secret",
400,
_entra_error("AADSTS50033", "A transient error has occurred."),
),
"external_server_retryable_client_secret": (
"client_secret",
400,
_entra_error("AADSTS90006", "The service is temporarily unavailable."),
),
"graph_retryable_client_secret": (
"client_secret",
400,
_entra_error("AADSTS90090", "The service is temporarily unavailable."),
),
"token_issuance_error_client_secret": (
"client_secret",
400,
_entra_error("AADSTS50000", "There was an error issuing a token or an issue with our sign-in service."),
),
"throttled_without_a_code_through_a_credential_chain": (
"default_azure_credential_chain",
429,
{"error": "temporarily_unavailable", "error_description": "Too many requests."},
),
"json_503_through_workload_identity": (
"workload_identity",
503,
_entra_error("AADSTS90033", "A transient error has occurred."),
),
"throttled_workload_identity_naming_a_rejection_code": (
"workload_identity",
429,
_entra_error("AADSTS7000215", "Invalid client secret provided."),
),
}

View file

@ -0,0 +1,39 @@
from __future__ import annotations
import os
import subprocess
import sys
from collections.abc import Generator
from typing import Final
import pytest
from litellm.rust_bridge import configuration
@pytest.fixture(autouse=True)
def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest discovers fixtures dynamically
monkeypatch: pytest.MonkeyPatch,
) -> Generator[None]:
configuration.reset_rust_configuration()
monkeypatch.delenv("LITELLM_RUST", raising=False)
yield
configuration.reset_rust_configuration()
@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False")))
def test_environment_controls_startup(value: str, expected: str) -> None:
environment: Final = {**os.environ, "LITELLM_RUST": value}
result: Final = subprocess.run(
(
sys.executable,
"-c",
"from litellm.rust_bridge.configuration import rust_enabled; print(rust_enabled())",
),
check=True,
capture_output=True,
text=True,
env=environment,
)
assert result.stdout.strip() == expected

View file

@ -0,0 +1,146 @@
import asyncio
import base64
import uuid
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass, field
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from threading import Event, Thread
from typing import Final, Literal
import pytest
import litellm
from litellm.llms.anthropic.messages.context_management.inference_view import CountContextEdits
from litellm.llms.vertex_ai.count_tokens.context import VertexGeminiCountContextError, gemini_count_contents
_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None)
_PNG_BYTES: Final = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
_LOOP_PROOF_TIMEOUT_SECONDS: Final = 5
@dataclass(frozen=True, slots=True)
class _LoopProbe:
request_arrived: Event = field(default_factory=Event)
loop_ran: Event = field(default_factory=Event)
loop_ran_while_download_pending: Event = field(default_factory=Event)
def _png_handler_waiting_for_loop(probe: _LoopProbe) -> type[BaseHTTPRequestHandler]:
class _PngHandler(BaseHTTPRequestHandler):
def do_GET(self) -> None:
probe.request_arrived.set()
if probe.loop_ran.wait(_LOOP_PROOF_TIMEOUT_SECONDS):
probe.loop_ran_while_download_pending.set()
self.send_response(200)
self.send_header("Content-Type", "image/png")
self.send_header("Content-Length", str(len(_PNG_BYTES)))
self.end_headers()
self.wfile.write(_PNG_BYTES)
def log_message(self, format: str, *args: object) -> None:
return None
return _PngHandler
@contextmanager
def _png_server_waiting_for_loop(probe: _LoopProbe) -> Generator[str, None, None]:
with ThreadingHTTPServer(("127.0.0.1", 0), _png_handler_waiting_for_loop(probe)) as server:
thread: Final = Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield f"http://127.0.0.1:{server.server_address[1]}"
finally:
server.shutdown()
async def _run_loop_once_download_starts(probe: _LoopProbe) -> None:
assert await asyncio.to_thread(probe.request_arrived.wait, _LOOP_PROOF_TIMEOUT_SECONDS)
probe.loop_ran.set()
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
async def test_remote_image_downloads_without_blocking_the_event_loop(
custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "user_url_validation", False)
probe: Final = _LoopProbe()
with _png_server_waiting_for_loop(probe) as base_url:
image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png"
_, contents = await asyncio.gather(
_run_loop_once_download_starts(probe),
gemini_count_contents(
(
{
"role": "user",
"content": [
{"type": "text", "text": "Describe it"},
{"type": "image", "source": {"type": "url", "url": image_url}},
],
},
),
"anthropic",
"gemini-2.5-flash",
custom_llm_provider,
{},
),
)
assert contents == [
{
"role": "user",
"parts": [
{"text": "Describe it"},
{"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}},
],
}
]
assert probe.loop_ran_while_download_pending.is_set()
class _MissingImageHandler(BaseHTTPRequestHandler):
def do_GET(self) -> None:
self.send_response(404)
self.send_header("Content-Length", "0")
self.end_headers()
def log_message(self, format: str, *args: object) -> None:
return None
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
async def test_image_that_cannot_be_fetched_returns_an_error_value(
custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "user_url_validation", False)
with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server:
Thread(target=server.serve_forever, daemon=True).start()
try:
contents = await gemini_count_contents(
(
{
"role": "user",
"content": [
{
"type": "image",
"source": {
"type": "url",
"url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png",
},
}
],
},
),
"anthropic",
"gemini-2.5-flash",
custom_llm_provider,
{},
)
finally:
server.shutdown()
assert isinstance(contents, VertexGeminiCountContextError)

View file

@ -0,0 +1,50 @@
import hashlib
import shutil
import subprocess
import pytest
from litellm import sandbox
DOCKER_IMAGE = "alpine:3.20"
def _docker_usable() -> bool:
if shutil.which("docker") is None:
return False
try:
return (
subprocess.run(
["docker", "info"], capture_output=True, timeout=20
).returncode
== 0
)
except (OSError, subprocess.SubprocessError):
return False
@pytest.mark.skipif(not _docker_usable(), reason="docker daemon not available")
async def test_real_docker_roundtrip():
box = sandbox.docker(DOCKER_IMAGE, workdir="/workspace")
try:
result = await box.run(["echo", "hello"], timeout=120)
assert result.stdout.strip() == "hello"
assert result.exit_code == 0
await box.write("seed.txt", b"seed")
await box.write("sub/out.txt", b"from host")
assert await box.read("sub/out.txt") == b"from host"
await box.write("node_modules/skip.js", b"x")
assert await box.which("sh") is not None
assert await box.which("definitely-not-a-binary-xyz") is None
tmp = await box.tempdir()
assert tmp.startswith("/")
snap = await box.snapshot()
assert snap == {
"seed.txt": hashlib.sha256(b"seed").hexdigest(),
"sub/out.txt": hashlib.sha256(b"from host").hexdigest(),
}
finally:
await box.close()

View file

@ -0,0 +1,296 @@
import asyncio
import importlib
import json
import os
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
import httpx
import pytest
import litellm
from litellm.exceptions import Timeout as LitellmTimeout
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.model_cost_table import publish
from litellm.llms.base_llm.auth.types import TokenExchangeSpec
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
MaskedHTTPStatusError,
get_httpx_client,
)
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
async def _read_http_request(reader: asyncio.StreamReader) -> None:
raw = b""
while b"\r\n\r\n" not in raw:
chunk = await reader.read(1024)
if not chunk:
return
raw += chunk
head, _, body = raw.partition(b"\r\n\r\n")
content_length = next(
(int(line.split(b":", 1)[1]) for line in head.split(b"\r\n") if line.lower().startswith(b"content-length")),
0,
)
while len(body) < content_length:
body += await reader.read(content_length - len(body))
@pytest.mark.asyncio
async def test_init_held_async_handler_survives_external_client_close():
handler = AsyncHTTPHandler(timeout=42.5)
held_client = handler.client
await held_client.aclose()
assert held_client.is_closed
async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
await _read_http_request(reader)
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
await writer.drain()
writer.close()
server = await asyncio.start_server(respond, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
try:
response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []})
finally:
server.close()
await server.wait_closed()
assert response.status_code == 200
assert handler.client is not held_client
assert handler.client.timeout == httpx.Timeout(42.5)
await handler.close()
@pytest.mark.asyncio
async def test_init_held_async_handler_survives_evicted_client_close():
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
cache = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0))
handler = AsyncHTTPHandler(timeout=42.5)
held_client = handler.client
cache.set_cache("init-held-handler", handler, litellm_owned_client=True, ttl=0)
await asyncio.sleep(0.02)
assert cache.get_cache("init-held-handler") is None
await asyncio.sleep(0.05)
assert held_client.is_closed
async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
await _read_http_request(reader)
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
await writer.drain()
writer.close()
server = await asyncio.start_server(respond, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
try:
response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []})
finally:
server.close()
await server.wait_closed()
assert response.status_code == 200
assert handler.client is not held_client
assert handler.client.timeout == httpx.Timeout(42.5)
await handler.close()
def test_init_held_sync_handler_recreates_closed_client():
from http.server import BaseHTTPRequestHandler, HTTPServer
class OkRequestHandler(BaseHTTPRequestHandler):
def do_GET(self):
self.send_response(200)
self.send_header("Content-Length", "2")
self.end_headers()
self.wfile.write(b"ok")
def log_message(self, format, *args):
pass
handler = HTTPHandler(timeout=7)
held_client = handler.client
held_client.close()
server = HTTPServer(("127.0.0.1", 0), OkRequestHandler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
response = handler.get(f"http://127.0.0.1:{server.server_port}/")
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
assert response.status_code == 200
assert handler.client is not held_client
assert handler.client.timeout == httpx.Timeout(7)
handler.close()
_WIF_SPEC: Final = TokenExchangeSpec(
token_url="https://api.anthropic.com/v1/oauth/token",
assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN",
assertion_field="assertion",
static_body={"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", "federation_rule_id": "fdrl_1"},
body_encoding="json",
request_headers={},
cache_key_identity=("fdrl_1", "org-1", "", ""),
timeout_seconds=2.0,
)
@pytest.fixture()
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
yield
record_vcr_outcome(request, vcr)
@pytest.fixture(scope="function")
def isolate_litellm_state():
"""
Per-function isolation fixture.
Resets litellm globals to their true defaults before each test and
restores them afterward, so tests don't leak side effects.
Works safely under pytest-xdist parallel execution.
"""
original_state = {}
for attr in (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
):
if hasattr(litellm, attr):
val = getattr(litellm, attr)
original_state[attr] = val.copy() if val else []
for attr in ("pre_call_rules", "post_call_rules"):
if hasattr(litellm, attr):
val = getattr(litellm, attr)
original_state[attr] = val.copy() if val else []
for attr in _SCALAR_DEFAULTS:
if hasattr(litellm, attr):
original_state[attr] = getattr(litellm, attr)
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
for attr in (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
"pre_call_rules",
"post_call_rules",
):
if hasattr(litellm, attr):
setattr(litellm, attr, [])
for attr, default_val in _SCALAR_DEFAULTS.items():
if hasattr(litellm, attr):
setattr(litellm, attr, default_val)
yield
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
for attr, original_value in original_state.items():
if hasattr(litellm, attr):
setattr(litellm, attr, original_value)
publish(litellm.model_cost)
_SCALAR_DEFAULTS = {
"num_retries": getattr(litellm, "num_retries", None),
"num_retries_per_request": getattr(litellm, "num_retries_per_request", None),
"request_timeout": getattr(litellm, "request_timeout", None),
"set_verbose": getattr(litellm, "set_verbose", False),
"cache": getattr(litellm, "cache", None),
"allowed_fails": getattr(litellm, "allowed_fails", 3),
"default_fallbacks": getattr(litellm, "default_fallbacks", None),
"enable_azure_ad_token_refresh": getattr(litellm, "enable_azure_ad_token_refresh", None),
"tag_budget_config": getattr(litellm, "tag_budget_config", None),
"model_cost": getattr(litellm, "model_cost", None),
"token_counter": getattr(litellm, "token_counter", None),
"disable_aiohttp_transport": getattr(litellm, "disable_aiohttp_transport", False),
"force_ipv4": getattr(litellm, "force_ipv4", False),
"drop_params": getattr(litellm, "drop_params", None),
"modify_params": getattr(litellm, "modify_params", False),
"api_base": getattr(litellm, "api_base", None),
"api_key": getattr(litellm, "api_key", None),
}
@pytest.fixture(scope="module")
def setup_and_teardown():
"""
Module-scoped setup. Reloads litellm only in single-process mode
(skipped under xdist to avoid cross-worker interference).
"""
import litellm
worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
if worker_id is None:
importlib.reload(litellm)
try:
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
import litellm.proxy.proxy_server
importlib.reload(litellm.proxy.proxy_server)
except Exception:
pass
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
yield
_SERVER_DELAY_S = 5
_PER_REQUEST_TIMEOUT_S = 1.0
_CLIENT_DEFAULT_TIMEOUT_S = 60.0
class _SlowHandler(BaseHTTPRequestHandler):
def do_POST(self):
time.sleep(_SERVER_DELAY_S)
try:
self.send_response(200)
self.end_headers()
self.wfile.write(b"{}")
except OSError:
pass
def log_message(self, *args):
pass
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_post_delay_exceeds_per_request_timeout_raises():
server = ThreadingHTTPServer(("127.0.0.1", 0), _SlowHandler)
threading.Thread(target=server.serve_forever, daemon=True).start()
host, port = server.server_address
handler = get_httpx_client(params={"timeout": _CLIENT_DEFAULT_TIMEOUT_S})
try:
with pytest.raises(LitellmTimeout):
handler.post(
f"http://{host}:{port}/delay",
headers={"content-type": "application/json"},
data=json.dumps({"model": "claude", "messages": []}),
timeout=_PER_REQUEST_TIMEOUT_S,
)
except MaskedHTTPStatusError as e:
pytest.skip(f"httpbin.org unavailable: {e}")
finally:
handler.close()
server.shutdown()
server.server_close()
_TIMEOUT_CLIENT_LIMITS: Final = httpx.Timeout(42.0, connect=7.0)

View file

@ -0,0 +1,37 @@
import subprocess
import sys
import textwrap
from typing import Final
from litellm.llms.base_llm.auth.jwt_signing import MISSING_SIGNING_DEPENDENCIES_MESSAGE
class TestBaseSdkImport:
"""A base ``pip install litellm`` has neither PyJWT nor cryptography (both are proxy extras),
and ``litellm/__init__`` reaches this module through the Anthropic provider, so a
module-level import of either would break ``import litellm`` for every base SDK user."""
def test_module_imports_with_pyjwt_and_cryptography_absent(self):
script: Final = textwrap.dedent(
"""
import sys
class Blocker:
def find_spec(self, name, path=None, target=None):
if name.split(".")[0] in {"jwt", "cryptography"}:
raise ModuleNotFoundError(f"No module named {name!r}")
sys.meta_path.insert(0, Blocker())
import litellm
from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json
try:
jwks_document_json("not a key")
except ImportError as e:
print(e)
"""
)
result: Final = subprocess.run(
[sys.executable, "-I", "-c", script], capture_output=True, text=True, check=False
)
assert result.returncode == 0, result.stderr[-2000:]
assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE

View file

@ -0,0 +1,23 @@
import os
import subprocess
import sys
from typing import Final
import litellm
import litellm.utils
def test_import_litellm_does_not_load_fastapi_or_bpe_table():
result: Final = subprocess.run(
[
sys.executable,
"-c",
"import sys, litellm; print(','.join(m for m in ('fastapi','starlette','litellm.proxy.proxy_cli','litellm.litellm_core_utils.default_encoding','openai.types.realtime') if m in sys.modules))",
],
check=True,
capture_output=True,
text=True,
env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.stdout.strip() == ""

View file

@ -0,0 +1,384 @@
import json
import os
import subprocess
import sys
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import pytest
import litellm
from litellm.cli.main import cli
if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib
_HELP_ARGV: Final = (
("--help",),
("logout", "--help"),
("login", "--help"),
("models", "list", "--help"),
("claude", "--help"),
("opencode", "--help"),
("codex", "--help"),
("pi", "--help"),
("keys", "list", "--help"),
)
_COMMAND_ARGV: Final = (("logout",), ("whoami",))
_RUN_EACH_ARGV: Final = """
import contextlib, io, json, sys
from {module} import cli
runs = {{}}
for argv in json.loads(sys.argv[1]):
sys.argv = ["lite", *argv]
output = io.StringIO()
with contextlib.redirect_stdout(output):
try:
cli()
except SystemExit as stopped:
code = stopped.code
else:
code = 0
runs[" ".join(argv)] = {{"exit": code, "output": output.getvalue()}}
print(json.dumps({{"runs": runs, "sdk_initialized": "completion" in vars(sys.modules["litellm"])}}))
"""
_USE_SDK_AFTER_HELP: Final = """
import builtins, importlib, json, sys
interpreter_imports = (builtins.__import__, importlib.import_module)
from {module} import cli
sys.argv = ["lite", "--help"]
try:
cli()
except SystemExit:
pass
before = "completion" in vars(sys.modules["litellm"])
{first_use}
from litellm.router import Router
import litellm
import litellm.responses.main
from litellm import acompletion
print(json.dumps({{
"before": before,
"router": Router.__name__,
"completion": callable(litellm.completion),
"acompletion": acompletion is litellm.acompletion,
"responses_is_function": callable(litellm.responses) and not isinstance(litellm.responses, type(litellm)),
"namespace": sorted(vars(litellm)),
"module_type": type(litellm).__name__,
"imports_restored": (builtins.__import__, importlib.import_module) == interpreter_imports,
}}))
"""
_FIRST_USE_FROM_THREADS: Final = """
import json, sys, threading
from litellm_cli.main import cli
sys.argv = ["lite", "--help"]
try:
cli()
except SystemExit:
pass
import litellm
start = threading.Barrier(4)
seen = []
def use_sdk():
start.wait()
try:
seen.append(callable(litellm.completion))
except AttributeError as error:
seen.append(str(error))
threads = [threading.Thread(target=use_sdk) for _ in range(4)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
print(json.dumps({"seen": seen}))
"""
_STAR_IMPORT_AFTER_HELP: Final = """
import json, sys
from litellm_cli.main import cli
sys.argv = ["lite", "--help"]
try:
cli()
except SystemExit:
pass
before = "completion" in vars(sys.modules["litellm"])
from litellm import *
print(json.dumps({"before": before, "completion": callable(globals().get("completion")), "acompletion": callable(globals().get("acompletion"))}))
"""
_SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING: Final = {
"sdk_sync.py": """
import threading
initializing = threading.Event()
submodule_loading = threading.Event()
""",
"litellm/__init__.py": """
import sdk_sync
sdk_sync.initializing.set()
sdk_sync.submodule_loading.wait(2)
from litellm import submodule
ready = True
""",
"litellm/submodule.py": """
import litellm
import sdk_sync
sdk_sync.submodule_loading.set()
ready_when_loaded = getattr(litellm, "ready", False)
""",
"litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n",
"litellm/proxy/__init__.py": "",
"litellm/proxy/client/__init__.py": "",
"litellm/proxy/client/cli/__init__.py": "",
"litellm/proxy/client/cli/main.py": """
import json, threading
import litellm
import sdk_sync
def cli():
def initialize_sdk():
litellm.ready
def import_submodule_once_initialization_starts():
sdk_sync.initializing.wait()
import litellm.submodule
threads = [threading.Thread(target=initialize_sdk, daemon=True),
threading.Thread(target=import_submodule_once_initialization_starts, daemon=True)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(10)
print(json.dumps({"finished": [not thread.is_alive() for thread in threads]}))
litellm_proxy_cli = cli
""",
}
_SDK_WHOSE_INIT_FAILS: Final = {
"litellm/__init__.py": 'raise ImportError("a dependency of the SDK is broken")\n',
"litellm/utils.py": "",
"litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n",
"litellm/proxy/__init__.py": "",
"litellm/proxy/client/__init__.py": "",
"litellm/proxy/client/cli/__init__.py": "",
"litellm/proxy/client/cli/main.py": """
import json
import litellm
def error_from(use):
try:
use()
except Exception as error:
return f"{type(error).__name__}: {error}"
return None
def import_sdk_code():
import litellm.utils
def cli():
print(json.dumps([error_from(lambda: litellm.completion), error_from(lambda: litellm.completion), error_from(import_sdk_code)]))
litellm_proxy_cli = cli
""",
}
_RUN_LITE_AGAINST_SDK_AT: Final = """
import sys
sys.path.insert(0, sys.argv[1])
from litellm_cli.main import cli
cli()
"""
_PYPROJECT: Final = Path(__file__).resolve().parents[3] / "pyproject.toml"
_RUN_ENTRY_POINT_HELP: Final = """
import contextlib, io, json, sys
from {module} import {function}
imported_sdk = "litellm" in sys.modules
sys.argv = ["litellm", "--help"]
output = io.StringIO()
with contextlib.redirect_stdout(output):
try:
{function}()
except SystemExit as stopped:
code = stopped.code
else:
code = 0
print(json.dumps({{"imported_sdk": imported_sdk, "exit": code, "help": output.getvalue()}}))
"""
_WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG: Final = """
import json, logging, os
os.environ["LITELLM_WORKER_LOG_CONFIG"] = json.dumps({{"version": 1, "root": {{"level": "INFO"}}}})
from litellm_cli.main import proxy_worker_app
proxy_worker_app()
print(json.dumps({{name: logging.getLogger(name).disabled for name in {names!r}}}))
"""
_COMMAND_MODULES_LOADED_BY: Final = """
import json, sys
from litellm_cli.main import cli
sys.argv = ["lite", *sys.argv[1:]]
try:
cli()
except SystemExit:
pass
prefix = "litellm.proxy.client.cli.commands."
print(json.dumps({"modules": sorted(name[len(prefix):] for name in sys.modules if name.startswith(prefix))}))
"""
def _run(source: str, argv: tuple[str, ...], home: Path) -> dict[str, object]:
checkout: Final = str(_PYPROJECT.parent)
bootstrap: Final = (
"import pathlib,sys\n"
f"sys.path.insert(0, {checkout!r})\n"
f"sys.path.insert(0, str(pathlib.Path({checkout!r}) / '.github/packaging'))\n"
"from cli_package import editable_support\n"
f"exec(editable_support(pathlib.Path({checkout!r}) / 'litellm/cli')[0][1])\n"
)
result: Final = subprocess.run(
[sys.executable, "-I", "-c", bootstrap + source, *argv],
capture_output=True,
text=True,
timeout=120,
env={**os.environ, "HOME": str(home), "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.returncode == 0, result.stderr
return json.loads(result.stdout.rstrip("\n").rpartition("\n")[2])
@pytest.fixture(scope="module")
def cli_runs(tmp_path_factory: pytest.TempPathFactory) -> Mapping[str, dict[str, object]]:
argv: Final = (json.dumps(_HELP_ARGV + _COMMAND_ARGV),)
return {
"thin": _run(_RUN_EACH_ARGV.format(module="litellm_cli.main"), argv, tmp_path_factory.mktemp("thin")),
"full": _run(
_RUN_EACH_ARGV.format(module="litellm.proxy.client.cli.main"), argv, tmp_path_factory.mktemp("full")
),
}
def test_litellm_command_imports_the_proxy_only_when_it_runs(tmp_path: Path):
module, _, function = tomllib.loads(_PYPROJECT.read_text())["project"]["scripts"]["litellm"].partition(":")
run: Final = _run(_RUN_ENTRY_POINT_HELP.format(module=module, function=function), (), tmp_path)
assert run["imported_sdk"] is False
assert run["exit"] == 0
assert "--num_workers" in str(run["help"])
def test_a_worker_log_config_silences_the_sdk_loggers_but_not_the_ones_the_proxy_app_creates(tmp_path: Path):
disabled: Final = _run(
_WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG.format(names=("LiteLLM Proxy", "litellm.proxy.health_check")),
(),
tmp_path,
)
assert disabled == {"LiteLLM Proxy": True, "litellm.proxy.health_check": False}
def test_lite_runs_commands_without_initializing_the_sdk(cli_runs: Mapping[str, dict[str, object]]):
assert cli_runs["thin"]["sdk_initialized"] is False
assert cli_runs["full"]["sdk_initialized"] is True
@pytest.mark.parametrize("argv", _HELP_ARGV + _COMMAND_ARGV, ids=" ".join)
def test_lite_prints_what_the_full_sdk_cli_prints(argv: tuple[str, ...], cli_runs: Mapping[str, dict[str, object]]):
thin: Final = cli_runs["thin"]["runs"]
full: Final = cli_runs["full"]["runs"]
assert isinstance(thin, dict) and isinstance(full, dict)
assert thin[" ".join(argv)] == full[" ".join(argv)]
assert thin[" ".join(argv)]["output"]
@pytest.mark.parametrize(
("argv", "loaded", "not_loaded"),
[
(("logout", "--help"), {"auth"}, {"agents", "autoroute.commands", "models", "keys"}),
(("models", "list", "--help"), {"models"}, {"agents", "autoroute.commands", "keys"}),
(("claude", "--help"), {"agents"}, {"autoroute.commands", "models", "keys"}),
],
)
def test_lite_imports_only_the_command_it_runs(
argv: tuple[str, ...], loaded: set[str], not_loaded: set[str], tmp_path: Path
):
modules: Final = _run(_COMMAND_MODULES_LOADED_BY, argv, tmp_path)["modules"]
assert isinstance(modules, list)
assert loaded <= set(modules)
assert not not_loaded & set(modules)
@pytest.fixture(scope="module")
def sdk_used_after_full_cli(tmp_path_factory: pytest.TempPathFactory) -> dict[str, object]:
return _run(
_USE_SDK_AFTER_HELP.format(module="litellm.proxy.client.cli.main", first_use="pass"),
(),
tmp_path_factory.mktemp("full"),
)
@pytest.mark.parametrize(
"first_use",
[
"import litellm; litellm.completion",
"import litellm; litellm.types.utils.ModelResponse",
"from litellm import acompletion",
"from litellm.router import Router",
"from litellm.types import utils",
"import litellm.responses.main",
'importlib.import_module("litellm.router")',
'importlib.import_module(".utils", "litellm")',
'exec("from .types import utils", {"__name__": "litellm.lazy_import", "__package__": "litellm"})',
],
)
def test_sdk_initializes_in_place_on_first_use_after_lite_imports(
first_use: str, sdk_used_after_full_cli: dict[str, object], tmp_path: Path
):
thin: Final = _run(_USE_SDK_AFTER_HELP.format(module="litellm_cli.main", first_use=first_use), (), tmp_path)
assert thin["before"] is False
assert sdk_used_after_full_cli["before"] is True
assert {**thin, "before": True} == sdk_used_after_full_cli
assert (thin["router"], thin["completion"], thin["acompletion"], thin["responses_is_function"]) == (
"Router",
True,
True,
True,
)
def test_star_import_after_lite_imports_gets_the_whole_sdk(tmp_path: Path):
assert _run(_STAR_IMPORT_AFTER_HELP, (), tmp_path) == {"before": False, "completion": True, "acompletion": True}
def test_threads_that_use_the_sdk_at_once_all_wait_for_one_initialization(tmp_path: Path):
used: Final = _run(_FIRST_USE_FROM_THREADS, (), tmp_path)
assert used == {"seen": [True, True, True, True]}
def _write_sdk(root: Path, files: Mapping[str, str]) -> Path:
for relative_path, source in files.items():
(root / relative_path).parent.mkdir(parents=True, exist_ok=True)
(root / relative_path).write_text(source)
return root
def test_importing_sdk_code_while_another_thread_initializes_the_sdk_does_not_hang(tmp_path: Path):
sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING)
assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == {"finished": [True, True]}
def test_every_use_after_a_failed_sdk_init_raises_the_init_error(tmp_path: Path):
sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_INIT_FAILS)
assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == ["ImportError: a dependency of the SDK is broken"] * 3

View file

@ -0,0 +1,323 @@
import json
import threading
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import count
from typing import Final
import httpx
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
MESSAGES: Final = ({"role": "user", "content": "Prove it"},)
HOP_1: Final = {
"candidates": [
{
"content": {"role": "model", "parts": [{"text": "hello"}]},
"finishReason": "CONTINUATION",
"continuationToken": "dG9rLTE=",
}
],
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 30, "totalTokenCount": 40},
}
HOP_2: Final = {
"candidates": [{"content": {"role": "model", "parts": [{"text": " world"}]}, "finishReason": "STOP"}],
"usageMetadata": {
"promptTokenCount": 10,
"cachedContentTokenCount": 9,
"candidatesTokenCount": 20,
"totalTokenCount": 30,
},
}
STREAM_HOP_1_END: Final = {
"candidates": [{**HOP_1["candidates"][0], "content": {"role": "model", "parts": [{"text": "lo"}]}}],
"usageMetadata": HOP_1["usageMetadata"],
}
def _replay(hops: Sequence[httpx.Response], sent: list[dict[str, object]]) -> Callable[[httpx.Request], httpx.Response]:
def handler(request: httpx.Request) -> httpx.Response:
sent.append(json.loads(request.content))
return hops[len(sent) - 1]
return handler
def _sse(*chunks: Mapping[str, object]) -> httpx.Response:
return httpx.Response(
200,
content="".join(f"data: {json.dumps(chunk)}\r\n\r\n" for chunk in chunks).encode(),
headers={"content-type": "text/event-stream"},
)
NO_CITATIONS_OR_URL_CONTEXT: Final = (
pytest.param({}, id="missing"),
pytest.param({"citationMetadata": None, "urlContextMetadata": None}, id="null"),
pytest.param({"citationMetadata": {}, "urlContextMetadata": {}}, id="empty"),
)
STREAM_FIRST_CHUNK: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hel"}]}}]}
def _stream_hops() -> tuple[httpx.Response, httpx.Response]:
return (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(HOP_2))
def _stream_text(chunks: list) -> str:
return "".join(c.choices[0].delta.content or "" for c in chunks if c.choices)
def _stream_finish_reasons(chunks: list) -> list[str]:
return [c.choices[0].finish_reason for c in chunks if c.choices and c.choices[0].finish_reason]
def _delta_fields(chunks: list) -> list[Mapping[str, object]]:
return [
choice.delta.provider_specific_fields
for chunk in chunks
for choice in chunk.choices # comprehension-ok: flattening stream chunks into their choices
if choice.delta.provider_specific_fields
]
def _client_tokens(chunks: list) -> list[object]:
return [fields["continuation_token"] for fields in _delta_fields(chunks) if "continuation_token" in fields]
@dataclass(frozen=True, slots=True)
class _TokenEndpoint:
credentials: str
down: threading.Event
def _service_account_json(token_uri: str) -> str:
private_key: Final = (
rsa.generate_private_key(public_exponent=65537, key_size=2048)
.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
.decode()
)
return json.dumps(
{
"type": "service_account",
"project_id": "p",
"private_key_id": "k",
"private_key": private_key,
"client_email": "sa@p.iam.gserviceaccount.com",
"token_uri": token_uri,
}
)
@pytest.fixture
def token_endpoint() -> Iterator[_TokenEndpoint]:
down: Final = threading.Event()
minted: Final = count(1)
class Handler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
self.rfile.read(int(self.headers["Content-Length"]))
refused: Final = down.is_set()
payload: Final = json.dumps(
{"error": "invalid_grant"} if refused else {"access_token": f"token-{next(minted)}", "expires_in": 0}
).encode()
self.send_response(400 if refused else 200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, *_args: object) -> None:
return None
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
threading.Thread(target=server.serve_forever, daemon=True).start()
yield _TokenEndpoint(_service_account_json(f"http://127.0.0.1:{server.server_port}/token"), down)
server.shutdown()
server.server_close()
def _google_taking_each_access_token_once(
hops: Sequence[httpx.Response], used: list[str], on_answer: Callable[[], object]
) -> Callable[[httpx.Request], httpx.Response]:
def handler(request: httpx.Request) -> httpx.Response:
token: Final = request.headers["authorization"]
if token in used:
return httpx.Response(401, json={"error": {"code": 401, "status": "UNAUTHENTICATED"}})
used.append(token)
on_answer()
return hops[len(used) - 1]
return handler
async def _vertex_answer(
is_async: bool,
stream: bool,
handler: Callable[[httpx.Request], httpx.Response],
credentials: str,
extra_headers: Mapping[str, str] | None = None,
) -> tuple[str, list[str], list[object]]:
kwargs: Final = {
"model": "vertex_ai/gemini-3.8-flash",
"extra_headers": extra_headers,
"messages": list(MESSAGES),
"max_tokens": 100_000,
"stream": stream,
"vertex_project": "p",
"vertex_location": "global",
"vertex_credentials": credentials,
}
transport: Final = httpx.MockTransport(handler)
response: Final = (
await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport))
if is_async
else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport)))
)
if stream:
chunks: Final = [chunk async for chunk in response] if is_async else list(response)
return _stream_text(chunks), _stream_finish_reasons(chunks), _client_tokens(chunks)
choice: Final = response.choices[0]
fields: Final = choice.message.provider_specific_fields or {}
return choice.message.content, [choice.finish_reason], [fields["continuation_token"]] if fields else []
def _two_hops(stream: bool) -> tuple[httpx.Response, httpx.Response]:
return _stream_hops() if stream else (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2))
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.asyncio
async def test_vertex_long_decode_finishes_when_each_continuation_outlives_the_access_token_before_it(
is_async: bool, stream: bool, token_endpoint: _TokenEndpoint
):
used: Final[list[str]] = []
google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, lambda: None)
answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials)
assert answer == ("hello world", ["stop"], [])
assert len(used) == 2
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.asyncio
async def test_vertex_long_decode_returns_the_answer_so_far_with_its_token_when_no_new_access_token_can_be_had(
is_async: bool, stream: bool, token_endpoint: _TokenEndpoint
):
used: Final[list[str]] = []
google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, token_endpoint.down.set)
answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials)
assert answer == ("hello", ["length"], ["dG9rLTE="])
assert len(used) == 1
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.asyncio
async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_every_continuation(
is_async: bool, stream: bool, token_endpoint: _TokenEndpoint
):
hops: Final = _two_hops(stream)
seen: Final[list[str]] = []
def google(request: httpx.Request) -> httpx.Response:
seen.append(request.headers["authorization"])
return hops[len(seen) - 1]
answer: Final = await _vertex_answer(
is_async, stream, google, token_endpoint.credentials, {"Authorization": "Bearer mine"}
)
assert answer == ("hello world", ["stop"], [])
assert seen == ["Bearer mine", "Bearer mine"]
BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS"
ARGON: Final = "gemini/gemini-4-argon"
OUTPUT_LIMIT_KEYS: Final = ("max_output_tokens", "maxOutputTokens")
def _output_limits(request: Mapping[str, object]) -> dict[str, object]:
generation_config: Final = request.get("generationConfig", {})
assert isinstance(generation_config, dict)
return {key: value for key, value in generation_config.items() if key in OUTPUT_LIMIT_KEYS}
def _cost_map_max_output_tokens(model: str) -> int:
provider, name = model.split("/", 1)
max_output_tokens: Final = litellm.get_model_info(model=name, custom_llm_provider=provider)["max_output_tokens"]
assert isinstance(max_output_tokens, int)
return max_output_tokens
def _gemini_chat_models_answering_no_text() -> tuple[str, ...]:
rows: Final = (
(name, row)
for name, row in litellm.model_cost.items()
if row.get("mode") == "chat" and "text" not in row.get("supported_output_modalities", ("text",))
)
return tuple(
name if name.startswith("gemini/") else f"vertex_ai/{name}"
for name, row in rows
if row.get("litellm_provider") in ("gemini", "vertex_ai-language-models")
)
@pytest.mark.asyncio
async def test_env_budget_skips_every_chat_model_the_cost_map_lists_as_answering_no_text(
monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint
):
monkeypatch.setenv(BUDGET_ENV, "50")
models: Final = _gemini_chat_models_answering_no_text()
sent: Final[list[dict[str, object]]] = []
replies: Final = tuple(httpx.Response(200, json=HOP_2) for _ in models)
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent)))
for model in models:
await litellm.acompletion(
model=model,
messages=list(MESSAGES),
api_key="fake-gemini-key",
vertex_project="p",
vertex_location="us-central1",
vertex_credentials=token_endpoint.credentials,
client=client,
)
assert models
assert {model: _output_limits(request) for model, request in zip(models, sent, strict=True)} == dict.fromkeys(
models, {}
)
@pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "vertex_ai/gemini-2.5-flash", ARGON])
@pytest.mark.asyncio
async def test_env_budget_above_a_models_output_limit_sends_that_limit_instead(
monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint, model: str
):
budget: Final = 2 * _cost_map_max_output_tokens(model)
monkeypatch.setenv(BUDGET_ENV, str(budget))
sent: Final[list[dict[str, object]]] = []
await litellm.acompletion(
model=model,
messages=list(MESSAGES),
api_key="fake-gemini-key",
vertex_project="p",
vertex_location="us-central1",
vertex_credentials=token_endpoint.credentials,
client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))),
)
assert _output_limits(sent[0]) == {"max_output_tokens": _cost_map_max_output_tokens(model)}

View file

@ -0,0 +1,193 @@
"""
Router behavior the router-code differential fuzz run found changed once the Rust core became the
only core: count settings below zero, an allowed_fails_policy write that is no policy, a model
response object as a deployment's mock_response, and endpoints read or patched on the Router class.
"""
import asyncio
import json
import socket
import threading
from collections.abc import Iterator, Mapping
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
import pytest
import litellm
from litellm import Router
MESSAGES: Final = [{"role": "user", "content": "hi"}]
class _Upstream(BaseHTTPRequestHandler):
"""An OpenAI upstream whose path's first segment names how it answers: ``cut`` sends two
Responses stream events then cuts the chunked body, ``filtered`` answers a chat completion the
provider blocked, ``down`` answers 502, and anything else a plain chat completion."""
protocol_version = "HTTP/1.1"
def log_message(self, format: str, *args: object) -> None: # noqa: A002 # the base class's name
del format, args
def _json(self, status: int, body: Mapping[str, object]) -> None:
data: Final = json.dumps(body).encode()
self.send_response(status)
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(data)))
self.end_headers()
_ = self.wfile.write(data)
def _chunk(self, event: Mapping[str, object]) -> None:
data: Final = f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
_ = self.wfile.write(f"{len(data):x}\r\n".encode() + data + b"\r\n")
self.wfile.flush()
def _cut(self) -> None:
self.send_response(200)
self.send_header("content-type", "text/event-stream")
self.send_header("transfer-encoding", "chunked")
self.end_headers()
response: Final = {"id": "r1", "object": "response", "created_at": 1, "status": "in_progress", "output": []}
self._chunk({"type": "response.created", "sequence_number": 0, "response": {**response, "model": "gpt-4o"}})
self._chunk(
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "m1",
"output_index": 0,
"content_index": 0,
"delta": "hello ",
}
)
self.connection.shutdown(socket.SHUT_RDWR)
self.close_connection = True
def do_POST(self) -> None:
_ = self.rfile.read(int(self.headers.get("content-length") or 0))
mode: Final = self.path.split("/")[1]
if mode == "cut":
self._cut()
return
if mode == "down":
self._json(502, {"error": {"message": "provider error", "type": "api_error"}})
return
finish: Final = "content_filter" if mode == "filtered" else "stop"
message: Final = {"role": "assistant", "content": mode}
choice: Final = {"index": 0, "finish_reason": finish, "message": message}
usage: Final = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
self._json(200, {"id": "c1", "object": "chat.completion", "created": 1, "choices": [choice], "usage": usage})
@pytest.fixture
def upstream() -> Iterator[str]:
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _Upstream)
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
yield f"http://127.0.0.1:{server.server_address[1]}"
server.shutdown()
thread.join()
def _served(name: str, deployment_id: str, model: str, api_base: str) -> dict[str, object]:
return {
"model_name": name,
"litellm_params": {"model": model, "api_key": "k", "api_base": api_base},
"model_info": {"id": deployment_id},
}
def _request_logging(call_type: str, *, stream: bool) -> object:
from litellm.litellm_core_utils.litellm_logging import Logging
return Logging(
model="g",
messages=MESSAGES,
stream=stream,
call_type=call_type,
start_time=datetime.now(), # noqa: DTZ005 # litellm's own naive start time
litellm_call_id="call-1",
function_id="f",
)
async def _recorded() -> None:
"""Let the background recordings of a stream's outcome run."""
for _ in range(50):
await asyncio.sleep(0.02)
@pytest.mark.asyncio
async def test_a_responses_stream_cut_after_it_opened_counts_one_fail_with_a_shared_logging_object(
upstream: str,
) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/cut/v1")], num_retries=0, disable_cooldowns=True
)
stream: Final = await router.aresponses(
model="g", input="hi", stream=True, litellm_logging_obj=_request_logging("aresponses", stream=True)
)
with pytest.raises(Exception): # noqa: B017, PT011 # the transport's read error, whatever class maps it
async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True
pass
await _recorded()
assert router.health().store.get_cache(key="d1:fails") == 1
@pytest.mark.asyncio
async def test_a_blocked_response_raised_for_a_fallback_counts_no_success_with_a_shared_logging_object(
upstream: str,
) -> None:
router: Final = Router(
model_list=[
_served("a", "a1", "openai/gpt-4o", f"{upstream}/filtered/v1"),
_served("safe", "s1", "openai/gpt-4o", f"{upstream}/ok/v1"),
],
content_policy_fallbacks=[{"a": ["safe"]}],
num_retries=0,
disable_cooldowns=True,
)
response: Final = await router.acompletion(
model="a", messages=MESSAGES, litellm_logging_obj=_request_logging("acompletion", stream=False)
)
await _recorded()
assert response.choices[0].message.content == "ok"
assert router.health().store.get_cache(key="a1:successes") is None
assert router.health().store.get_cache(key="s1:successes") == 1
@pytest.mark.asyncio
async def test_a_text_completion_stream_that_fails_before_its_first_chunk_counts_nothing(upstream: str) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-3.5-turbo-instruct", f"{upstream}/down/v1")],
num_retries=0,
disable_cooldowns=True,
)
stream: Final = await router.atext_completion(model="g", prompt="hi", stream=True)
with pytest.raises(litellm.BadGatewayError):
async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True
pass
await _recorded()
assert router.health().store.get_cache(key="d1:successes") is None
assert router.health().store.get_cache(key="d1:fails") is None
@pytest.mark.asyncio
async def test_a_request_refused_before_the_provider_counts_against_no_deployment(upstream: str) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/ok/v1")], num_retries=0, disable_cooldowns=True
)
with pytest.raises(litellm.UnsupportedParamsError):
await router.acompletion(model="g", messages=MESSAGES, reasoning_effort="low")
await _recorded()
assert router.health().store.get_cache(key="d1:fails") is None
assert not [key for key in router.health().store.in_memory_cache.cache_dict if ":rpm" in key]

View file

@ -0,0 +1,28 @@
import subprocess
import sys
import pytest
import litellm.router_host as router_host_module
from litellm.router_host import router_host
def test_router_uses_the_proxy_host_without_importing_the_proxy_first():
# A fresh interpreter, so no earlier test has imported litellm.proxy.
script = (
"import sys\n"
"from litellm.router import Router\n"
"from litellm.router_host import router_host\n"
"assert 'litellm.proxy' not in sys.modules, 'importing the router imported the proxy'\n"
"print(type(router_host()).__name__)\n"
)
result = subprocess.run([sys.executable, "-P", "-c", script], capture_output=True, text=True, check=True)
assert result.stdout.strip() == "ProxyRouterHost"
@pytest.fixture(autouse=True)
def nothing_installed():
saved = router_host_module._installed[0]
router_host_module._installed[0] = None
yield
router_host_module._installed[0] = saved

View file

@ -0,0 +1,367 @@
import itertools
import json
import os
import subprocess
import sys
from collections.abc import Mapping
from pathlib import Path
from types import MappingProxyType
from typing import Final
import pytest
from pydantic import TypeAdapter
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
from litellm.utils import ModelResponseListIterator
from litellm import completion
bedrock_chunks = [
ModelResponseStream(
id="chatcmpl-d249def8-a78b-464c-87b5-3a6f43565292",
created=1742056047,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
provider_specific_fields=None,
content="I'm Claude",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=None,
),
ModelResponseStream(
id="chatcmpl-fe559823-b383-4249-ab87-52f6ad9d08c2",
created=1742056047,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
provider_specific_fields=None,
content=", an AI",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=None,
),
ModelResponseStream(
id="chatcmpl-c1c6cc2f-75b9-4a24-88b9-4e5aacd0268b",
created=1742056047,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(
provider_specific_fields=None,
content="",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=None,
),
]
FINISH_ONLY_LOGPROBS_STREAM_SCRIPT: Final = """
import json
import sys
sys.path.insert(0, sys.argv[1])
import httpx
from openai import OpenAI
import litellm
upstream_choices = (
{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None},
{
"index": 0,
"delta": {},
"logprobs": {"content": [{"token": "<|return|>", "logprob": -0.5, "top_logprobs": []}]},
"finish_reason": "stop",
},
)
body = "".join(
"data: "
+ json.dumps({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [choice]})
+ "\\n\\n"
for choice in upstream_choices
) + "data: [DONE]\\n\\n"
client = OpenAI(
api_key="unused",
base_url="https://provider.example/v1",
http_client=httpx.Client(
transport=httpx.MockTransport(
lambda request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"})
)
),
)
stream = litellm.completion(
model="openai/m", messages=[{"role": "user", "content": "hi"}], stream=True, logprobs=True, client=client
)
print(json.dumps([json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in stream]))
"""
def test_stream_with_logprobs_only_on_finish_chunk_serializes_in_a_fresh_process() -> None:
result: Final = subprocess.run(
[sys.executable, "-I", "-c", FINISH_ONLY_LOGPROBS_STREAM_SCRIPT, str(Path(litellm.__file__).parent.parent)],
capture_output=True,
text=True,
timeout=120,
env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.returncode == 0, result.stdout + result.stderr
wire_choices: Final = tuple(
itertools.chain.from_iterable(chunk["choices"] for chunk in json.loads(result.stdout.strip().splitlines()[-1]))
)
finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")]
assert [choice["finish_reason"] for choice in finish_choices] == ["stop"]
assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == ["<|return|>"]
_REPETITION_TEST_CASES = [
# Basic cases
pytest.param(
["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
True,
id="all_identical_raises",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1),
False,
id="below_threshold_no_raise",
),
pytest.param(
[None] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="none_content_no_raise",
),
pytest.param(
[""] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="empty_content_no_raise",
),
# Short content (len <= 2) should not raise
pytest.param(
["##"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="short_content_2chars_no_raise",
),
pytest.param(
["{"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="short_content_1char_no_raise",
),
pytest.param(
["ab"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="short_content_2chars_ab_no_raise",
),
# All different chunks
pytest.param(
["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="all_different_no_raise",
),
# One chunk different at various positions
pytest.param(
["different_first"] + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1),
False,
id="first_chunk_different_no_raise",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1) + ["different_last"],
False,
id="last_chunk_different_no_raise",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1)
+ ["different_mid"]
+ ["same"]
* (
litellm.REPEATED_STREAMING_CHUNK_LIMIT
- litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2
+ 1
),
False,
id="middle_chunk_different_no_raise",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 2) + ["diff", "diff"],
False,
id="last_two_different_no_raise",
),
pytest.param(
["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT
+ ["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT
+ ["diff"],
True,
id="in_between_same_and_diff_raise",
),
]
# Azure streaming chunks that reproduce issue #24221:
# Azure sends an initial chunk with prompt_filter_results and choices=[],
# then a chunk with role='assistant' and content='', then content chunks.
# With stream_options.include_usage=True, the empty-choices chunk was
# forwarded with an inflated default choice, consuming the sent_first_chunk
# flag and causing strip_role_from_delta to strip the role from the real
# first chunk.
_AZURE_CHUNKS_WITH_PROMPT_FILTER = [
# Chunk 1: prompt_filter_results, no choices (Azure-specific)
ModelResponseStream(
id="chatcmpl-abc123",
created=1742056047,
model=None,
object="chat.completion.chunk",
choices=[],
usage=None,
),
# Chunk 2: first real chunk with role='assistant' and empty content
ModelResponseStream(
id="chatcmpl-abc123",
created=1742056047,
model=None,
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(content="", role="assistant"),
)
],
usage=None,
),
# Chunk 3: content
ModelResponseStream(
id="chatcmpl-abc123",
created=1742056047,
model=None,
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(content="Hello!"),
)
],
usage=None,
),
# Chunk 4: finish_reason
ModelResponseStream(
id="chatcmpl-abc123",
created=1742056047,
model=None,
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(),
)
],
usage=None,
),
# Chunk 5: final usage chunk, no choices
ModelResponseStream(
id="chatcmpl-abc123",
created=1742056047,
model=None,
object="chat.completion.chunk",
choices=[],
usage=Usage(
completion_tokens=10,
prompt_tokens=20,
total_tokens=30,
),
),
]
_AZURE_CHOICE_EXTRAS: Final = MappingProxyType(
{"content_filter_results": {"hate": {"filtered": False, "severity": "safe"}}}
)
_OPENAI_AND_AZURE_CHOICE_EXTRAS: Final = (
pytest.param("openai", MappingProxyType({}), id="openai"),
pytest.param("azure", _AZURE_CHOICE_EXTRAS, id="azure"),
)
_TOKEN_IDS: Final = TypeAdapter(list[int] | None)
_CHOICE_INDEX: Final = TypeAdapter(int)
_DELTA: Final = TypeAdapter(Mapping[str, object])
def _streaming_wrapper_for_unit_test(
model: str,
custom_llm_provider: str,
chunks: list[ModelResponseStream],
) -> CustomStreamWrapper:
return CustomStreamWrapper(
completion_stream=ModelResponseListIterator(model_responses=chunks),
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=Logging(
model=model,
messages=[{"role": "user", "content": "unit test"}],
stream=True,
call_type="completion",
start_time=0.0,
litellm_call_id="unit-test-call",
function_id="unit-test-function",
),
)
def _unit_test_streaming_chunk(
content: str | None,
finish_reason: str | None = None,
role: str | None = "assistant",
) -> ModelResponseStream:
return ModelResponseStream(
id="unit-test-stream",
created=1,
model="unit-test-model",
choices=[
StreamingChoices(
index=0,
delta=Delta(content=content, role=role),
finish_reason=finish_reason,
)
],
)

View file

@ -0,0 +1,207 @@
#### What this tests ####
# This tests litellm.token_counter.token_counter() function
import json
import os
import subprocess
import sys
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import tiktoken
from tokenizers import Regex, Tokenizer, models, pre_tokenizers
from unittest.mock import AsyncMock
import litellm
from litellm import decode, encode
import litellm.constants
from litellm.types.utils import SelectTokenizerResponse
## Choose which token_counter the test will use.
# test_token_counter_normal_plus_function_calling()
# test_tokenizers()
def test_encoding_and_decoding(tmp_path: Path):
sample_text = "Hellö World, this is my input string!"
# openai encoding + decoding
openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text)
openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens)
assert openai_text == sample_text
# claude encoding + decoding
claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text)
claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens)
assert claude_text == sample_text
# cohere encoding + decoding
cohere_tokens = encode(model="command-nightly", text=sample_text)
cohere_text = decode(model="command-nightly", tokens=cohere_tokens)
assert cohere_text == sample_text
# llama2 encoding + decoding
words = sample_text.split()
result = _run_in_memory_hub(
HUB_ROUND_TRIP_SCRIPT,
{
"hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(
pre_tokenizers.WhitespaceSplit(),
vocab={"[UNK]": 0, **{word: i + 1 for i, word in enumerate(words)}},
)
},
sample_text,
tmp_path,
)
assert result["decoded"] == sample_text
assert result["requested"] == ["hf-internal-testing/llama-tokenizer"]
assert len(result["tokens"]) == len(words)
assert len(result["tokens"]) != len(encode(model="gpt-3.5-turbo", text=sample_text))
# test_encoding_and_decoding()
# test_gpt_vision_token_counting()
from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, claude_json_str
# Clear the cache at module load to ensure clean state
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache()
HUB_SETUP_SCRIPT: Final = """
import json
import sys
sys.path.insert(0, sys.argv[1])
import httpx
import huggingface_hub
import litellm
served = json.loads(sys.argv[2])
text = sys.argv[3]
requested = []
def handle(request):
repo = request.url.path.lstrip("/").split("/resolve/")[0]
if repo not in served or not request.url.path.endswith("/tokenizer.json"):
return httpx.Response(404)
requested.append(repo)
payload = served[repo].encode()
headers = {"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}
return httpx.Response(200, headers=headers, content=payload if request.method == "GET" else b"")
huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle)))
"""
HUB_TOKENIZER_SCRIPT: Final = HUB_SETUP_SCRIPT + """
litellm.cohere_models = {"command-r-v1"}
litellm.anthropic_models = {"claude-2"}
custom = litellm.create_pretrained_tokenizer("Xenova/llama-3-tokenizer")
print(json.dumps({
"llama2": litellm.token_counter(model="meta-llama/Llama-2-7b-chat", text=text),
"llama3": litellm.token_counter(model="meta-llama/llama-3-70b-instruct", text=text),
"cohere": litellm.token_counter(model="command-r-v1", text=text),
"anthropic": litellm.token_counter(model="claude-2", text=text),
"custom": litellm.token_counter(custom_tokenizer=custom, text=text),
"requested": sorted(set(requested)),
}))
"""
HUB_ROUND_TRIP_SCRIPT: Final = HUB_SETUP_SCRIPT + """
tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text)
print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))}))
"""
HUB_RETRY_DURING_EXIT_SCRIPT: Final = HUB_SETUP_SCRIPT + """
import threading
import time
from litellm.utils import _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer
retry_reached_hub = threading.Event()
def answer_once_the_process_is_exiting(request):
retry_reached_hub.set()
while threading.main_thread().is_alive():
time.sleep(0.01)
time.sleep(0.5)
return handle(request)
now = [0.0]
outcomes = _HuggingFaceTokenizerLoadOutcomes(clock=lambda: now[0])
model = "meta-llama/llama-3-70b-instruct"
_return_huggingface_tokenizer(model, outcomes)
huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(answer_once_the_process_is_exiting)))
now[0] = 3600.0
_return_huggingface_tokenizer(model, outcomes)
print(json.dumps({"retry_reached_hub": retry_reached_hub.wait(timeout=60)}))
"""
def _word_level_tokenizer_json(
pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None
) -> str:
tokenizer: Final = Tokenizer(
models.WordLevel(vocab=dict(vocab) if vocab is not None else {"[UNK]": 0}, unk_token="[UNK]")
)
tokenizer.pre_tokenizer = pre_tokenizer
return tokenizer.to_str()
def _run_in_memory_hub(script: str, served: dict[str, str], text: str, tmp_path: Path) -> dict:
result: Final = subprocess.run(
[
sys.executable,
"-I",
"-c",
script,
str(Path(litellm.__file__).parent.parent),
json.dumps(served),
text,
],
capture_output=True,
text=True,
timeout=60,
env={
**os.environ,
"HF_HOME": str(tmp_path / "home"),
"HF_HUB_CACHE": str(tmp_path / "cache"),
"HF_ENDPOINT": "http://127.0.0.1:9",
"HF_HUB_OFFLINE": "0",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
},
)
assert result.returncode == 0, result.stdout + result.stderr
return json.loads(result.stdout.strip().splitlines()[-1])
def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_tokenizer(tmp_path: Path) -> None:
sample: Final = "Tokenizers disagree: anthropic, tiktoken; llama-2 & llama-3!"
served: Final = {
"hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()),
"Xenova/llama-3-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Split(Regex("."), "isolated")),
"Xenova/c4ai-command-r-v01-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Whitespace()),
}
expected: Final = {repo: len(Tokenizer.from_str(payload).encode(sample).ids) for repo, payload in served.items()}
anthropic_count: Final = len(Tokenizer.from_str(claude_json_str).encode(sample).ids)
tiktoken_count: Final = litellm.token_counter(model="gpt-3.5-turbo", text=sample)
assert len({*expected.values(), anthropic_count, tiktoken_count}) == len(expected) + 2
counts: Final = _run_in_memory_hub(HUB_TOKENIZER_SCRIPT, served, sample, tmp_path)
assert counts == {
"llama2": expected["hf-internal-testing/llama-tokenizer"],
"llama3": expected["Xenova/llama-3-tokenizer"],
"cohere": expected["Xenova/c4ai-command-r-v01-tokenizer"],
"anthropic": anthropic_count,
"custom": expected["Xenova/llama-3-tokenizer"],
"requested": sorted(served),
}
def test_process_exits_cleanly_while_a_background_tokenizer_retry_is_downloading(tmp_path: Path) -> None:
assert _run_in_memory_hub(HUB_RETRY_DURING_EXIT_SCRIPT, {}, "", tmp_path) == {"retry_reached_hub": True}

View file

@ -0,0 +1,119 @@
import os
import subprocess
import sys
from pathlib import Path
from typing import Final
import pytest
import litellm
from tests.unit.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
@pytest.mark.parametrize("offline", ("0", "1"))
def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None:
script: Final = """
import json
import sys
from pathlib import Path
sys.path.insert(0, sys.argv[1])
import httpx
import huggingface_hub
from huggingface_hub.errors import LocalEntryNotFoundError
import litellm
payload = sys.argv[2].encode()
offline = sys.argv[3] == "1"
observed = []
def handle(request):
assert not offline, "offline loading issued a request"
if request.url.path.endswith("/tokenizer.json"):
observed.append(request.headers.get("authorization"))
if request.headers.get("authorization") != "Bearer audit-fixture-token":
return httpx.Response(401)
return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"")
if not offline:
huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle)))
try:
tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"]
except LocalEntryNotFoundError:
assert offline
assert observed == []
else:
assert not offline
assert "Bearer audit-fixture-token" in observed
assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World"
assert tuple(Path(sys.argv[4]).rglob("tokenizer.json"))
print("compatible")
"""
result: Final = subprocess.run(
[
sys.executable,
"-I",
"-c",
script,
str(Path(litellm.__file__).parent.parent),
TOKENIZER_JSON,
offline,
str(tmp_path / "cache"),
],
capture_output=True,
text=True,
timeout=30,
env={
**os.environ,
"HF_HOME": str(tmp_path / "home"),
"HF_HUB_CACHE": str(tmp_path / "cache"),
"HF_ENDPOINT": "http://127.0.0.1:9",
"HF_TOKEN": "audit-fixture-token",
"HF_HUB_OFFLINE": offline,
"HF_HUB_DISABLE_IMPLICIT_TOKEN": "0",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
},
)
assert result.returncode == 0, result.stdout + result.stderr
assert result.stdout.strip() == "compatible"
@pytest.mark.parametrize("rust", (None, "0", "1"))
def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None:
script: Final = """
import sys
sys.path.insert(0, sys.argv[1])
def reject_network(event, args):
if event == "socket.connect":
raise AssertionError("tokenizer attempted a network connection")
sys.addaudithook(reject_network)
import litellm
from litellm.rust_bridge.tokenizer import get_encoding
import tiktoken
from tokenizers import Tokenizer
assert isinstance(litellm.encoding, tiktoken.Encoding)
for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"):
encoding = get_encoding(name)
text = "offline café 漢字 🙂" + " " * 64
assert encoding.decode(encoding.encode(text)) == text
ids = litellm.encode(text="hello world")
assert litellm.decode(tokens=ids) == "hello world"
assert litellm.token_counter(model=None, text="hello world") == len(ids)
custom = litellm.create_tokenizer(sys.argv[2])
assert isinstance(custom["tokenizer"], Tokenizer)
custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]")
assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World"
print("compatible")
"""
result: Final = subprocess.run(
[sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON],
capture_output=True,
text=True,
timeout=30,
cwd=tmp_path,
env={
**{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"},
**({"LITELLM_RUST": rust} if rust is not None else {}),
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
"TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"),
},
)
assert result.returncode == 0, result.stdout + result.stderr
assert result.stdout.strip() == "compatible"
assert not (tmp_path / "unused-tokenizer-cache").exists()

View file

@ -18,7 +18,6 @@ import pytest
from redis.asyncio import ConnectionPool, Redis
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
@ -111,19 +110,6 @@ async def test_redis_client_is_closed_only_after_its_subscription_releases_the_c
assert closer.pending_count == 0
async def _trickling_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
"""Serves a chunked body slowly, so a request stays on the wire long enough to observe."""
await reader.read(4096)
writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n")
await writer.drain()
for _ in range(6):
writer.write(b"5\r\nhello\r\n")
await writer.drain()
await asyncio.sleep(0.1)
writer.write(b"0\r\n\r\n")
await writer.drain()
@pytest.mark.asyncio
async def test_owned_client_is_closed_once_the_grace_window_elapses():
clock = FakeClock()
@ -416,79 +402,6 @@ async def test_a_client_evicted_on_a_collected_loop_is_dropped():
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_a_client_serving_a_request_is_not_closed_when_its_grace_window_ends():
"""The grace window on its own cannot promise that a request has finished.
``litellm.request_timeout`` defaults to 6000 seconds and a streaming response
is bounded only by how long the upstream keeps sending, so a client past its
deadline is closed only once its own pool reports nothing in flight.
"""
server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
clock = FakeClock()
closer = make_closer(clock)
client = httpx.AsyncClient()
closer.mark_owned(client)
closer.schedule(client)
async def read_the_stream() -> int:
received = 0
async with client.stream("GET", f"http://127.0.0.1:{port}/") as response:
async for chunk in response.aiter_bytes():
received += len(chunk)
return received
streaming = asyncio.create_task(read_the_stream())
await asyncio.sleep(0.25) # the request is on the wire
clock.advance(3600.0) # and its grace window is long gone
closer.reap()
await asyncio.sleep(0.05)
assert client.is_closed is False, "closed a client that was serving a request"
assert await streaming > 0, "the in-flight request did not survive the reap"
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.is_closed is True, "an idle client past its grace window must be closed"
assert closer.pending_count == 0
server.close()
@pytest.mark.asyncio
async def test_the_aiohttp_backed_handler_is_not_closed_mid_request():
"""The default async path is aiohttp-backed, whose pool accounts for its own leases."""
server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
clock = FakeClock()
closer = make_closer(clock)
handler = AsyncHTTPHandler()
held_client = handler.client
closer.mark_owned(handler)
closer.schedule(handler)
request = asyncio.create_task(handler.get(f"http://127.0.0.1:{port}/"))
await asyncio.sleep(0.25)
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert held_client.is_closed is False, "closed a handler that was serving a request"
assert (await request).status_code == 200
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert held_client.is_closed is True
assert handler.client.is_closed is False, "a held handler must self-heal after its evicted client is closed"
server.close()
def test_the_pending_queue_cannot_grow_past_its_bound():
"""A caller that churns the client cache must not be able to grow this queue."""
clock = FakeClock()

View file

@ -1,9 +1,6 @@
import hashlib
import os
import struct
import subprocess
import sys
import textwrap
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
@ -508,29 +505,3 @@ async def test_index_info_uses_valkey_ft_info():
assert result == {"num_docs": 3}
async_client.ft.assert_called_once_with("test_index")
def test_importing_caching_does_not_require_redis():
# redis is an optional dependency (extra_proxy), so the base SDK can be
# installed without it. Selecting valkey-semantic needs redis, but merely
# importing litellm.caching.caching must not, or `import litellm` breaks for
# every base-SDK user. This runs in a subprocess with redis blocked so the
# check is not polluted by redis already being imported in this session.
code = textwrap.dedent("""
import sys
for name in ("redis", "redis.asyncio", "redis.commands",
"redis.commands.search"):
sys.modules[name] = None
import litellm.caching.caching # must not import redis at module top
from litellm.types.caching import LiteLLMCacheType
assert LiteLLMCacheType.VALKEY_SEMANTIC == "valkey-semantic"
print("ok")
""")
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
env={**os.environ, "PYTHONPATH": _REPO_ROOT},
)
assert result.returncode == 0, result.stderr
assert "ok" in result.stdout

View file

@ -1,12 +1,17 @@
# conftest.py
import asyncio
import importlib
import os
from typing import Final
import pytest
import litellm
import litellm.proxy.proxy_server
from tests._module_state import collect_flushable_caches, flush_caches, restore_mutable_state, snapshot_mutable_state
_LITELLM_STATE: Final = snapshot_mutable_state(litellm)
_PROXY_SERVER_STATE: Final = snapshot_mutable_state(litellm.proxy.proxy_server)
_FLUSHABLE_CACHES: Final = collect_flushable_caches((litellm, litellm.proxy.proxy_server))
@pytest.fixture(scope="session")
@ -19,38 +24,17 @@ def event_loop():
loop.close()
@pytest.fixture(scope="function", autouse=True)
def setup_and_teardown():
"""
This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained.
"""
curr_dir = os.getcwd() # Get the current working directory
from litellm import Router
importlib.reload(litellm)
try:
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
importlib.reload(litellm.proxy.proxy_server)
except Exception as e:
print(f"Error reloading litellm.proxy.proxy_server: {e}")
litellm.in_memory_llm_clients_cache.flush_cache()
import asyncio
restore_mutable_state(litellm, _LITELLM_STATE)
restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE)
flush_caches(_FLUSHABLE_CACHES)
loop = asyncio.get_event_loop_policy().new_event_loop()
asyncio.set_event_loop(loop)
print(litellm)
# from litellm import Router, completion, aembedding, acompletion, embedding
yield
# Teardown code (executes after the yield point)
loop.close() # Close the loop created earlier
asyncio.set_event_loop(None) # Remove the reference to the loop
loop.close()
asyncio.set_event_loop(None)
def pytest_collection_modifyitems(config, items):

View file

@ -1,7 +1,5 @@
import asyncio
import hashlib
import shutil
import subprocess
from typing import Optional
import pytest
@ -10,8 +8,6 @@ from litellm import sandbox
from litellm.harness.errors import SandboxError
from litellm.harness.sandbox import DockerSandbox, Sandbox
from litellm.harness.sandbox.docker import parse_sha256sum
DOCKER_IMAGE = "alpine:3.20"
CID = "cid123"
@ -198,44 +194,3 @@ async def test_missing_docker_binary(monkeypatch):
def test_parse_sha256sum_ignores_junk():
assert parse_sha256sum("garbage\n\n") == {}
def _docker_usable() -> bool:
if shutil.which("docker") is None:
return False
try:
return (
subprocess.run(
["docker", "info"], capture_output=True, timeout=20
).returncode
== 0
)
except (OSError, subprocess.SubprocessError):
return False
@pytest.mark.skipif(not _docker_usable(), reason="docker daemon not available")
async def test_real_docker_roundtrip():
box = sandbox.docker(DOCKER_IMAGE, workdir="/workspace")
try:
result = await box.run(["echo", "hello"], timeout=120)
assert result.stdout.strip() == "hello"
assert result.exit_code == 0
await box.write("seed.txt", b"seed")
await box.write("sub/out.txt", b"from host")
assert await box.read("sub/out.txt") == b"from host"
await box.write("node_modules/skip.js", b"x")
assert await box.which("sh") is not None
assert await box.which("definitely-not-a-binary-xyz") is None
tmp = await box.tempdir()
assert tmp.startswith("/")
snap = await box.snapshot()
assert snap == {
"seed.txt": hashlib.sha256(b"seed").hexdigest(),
"sub/out.txt": hashlib.sha256(b"from host").hexdigest(),
}
finally:
await box.close()

View file

@ -4,12 +4,10 @@ import json
import logging
import os
import pathlib
import subprocess
import sys
import textwrap
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, List, Optional, Tuple, cast
from typing import Final, List, cast
from unittest.mock import MagicMock, patch
import pytest
@ -3501,54 +3499,6 @@ class TestConfiguredInjectionPointsSurviveClientMarks:
assert sys2 == sys1
class TestAnthropicPromptCachingEnvVars:
"""Both settings are read from the environment at import, so an admin can enable
auto-caching without a config file. Each case re-imports litellm in a subprocess
so the env is read fresh without contaminating this process's module graph.
"""
@staticmethod
def _import_litellm_with_env(env_override: dict) -> Tuple[bool, Optional[str]]:
env = os.environ.copy()
env.pop("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", None)
env.pop("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL", None)
env.update(env_override)
script = textwrap.dedent(
"""
import json, litellm
print(json.dumps([litellm.enable_anthropic_prompt_caching, litellm.anthropic_prompt_caching_ttl]))
"""
)
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300)
assert result.returncode == 0, result.stderr
enabled, ttl = json.loads(result.stdout.strip().splitlines()[-1])
return enabled, ttl
def test_unset_env_leaves_auto_caching_off(self):
assert self._import_litellm_with_env({}) == (False, None)
@pytest.mark.parametrize("value", ["true", "True", "TRUE"])
def test_env_enables_auto_caching_case_insensitively(self, value):
enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value})
assert enabled is True
@pytest.mark.parametrize("value", ["false", "0", "yes", ""])
def test_env_only_enables_on_true(self, value):
enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value})
assert enabled is False
@pytest.mark.parametrize("value", ["5m", "1h"])
def test_ttl_env_is_applied(self, value):
_, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value})
assert ttl == value
@pytest.mark.parametrize("value", ["10m", "1H", "3600", "ephemeral"])
def test_unsupported_ttl_env_falls_back_to_provider_default(self, value):
"""An unparseable TTL must fall back to Anthropic's 5m default, never reach the provider verbatim."""
_, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value})
assert ttl is None
def _contains_key(value, key) -> bool:
if isinstance(value, dict):
return key in value or any(_contains_key(v, key) for v in value.values())

View file

@ -215,36 +215,6 @@ async def test_hook_searches_through_the_injected_router_with_the_request_metada
assert messages[0]["content"] == "Context:\n\ncontext from vs-router\n\n"
@pytest.mark.asyncio
@pytest.mark.parametrize("sdk_verbose", [False, True])
async def test_hook_falls_back_to_the_sdk_when_the_runtime_has_no_router(
registry_with: RegisterStores,
warnings: list[logging.LogRecord],
monkeypatch: pytest.MonkeyPatch,
sdk_verbose: bool,
) -> None:
monkeypatch.setattr(litellm, "set_verbose", sdk_verbose)
registry_with("vs-sdk", custom_llm_provider="lit6752-not-a-provider")
_, messages, _ = await _run_hook(
VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)),
["vs-sdk"],
FakeLoggingObj({"user_api_key_team_id": "team-a"}),
)
assert messages == [{"role": "user", "content": "what is litellm?"}]
search_warnings: Final = tuple(
record.getMessage()
for record in warnings
if record.getMessage().startswith("Vector store search failed for vector_store_id=")
)
assert len(search_warnings) == 1
assert search_warnings[0].startswith(
"Vector store search failed for vector_store_id=vs-sdk, continuing without its context: "
)
assert "is not a valid LlmProviders" in search_warnings[0]
@pytest.mark.asyncio
async def test_every_healthy_vector_store_contributes_its_own_context(registry_with: RegisterStores) -> None:
"""Regression (LIT-6752): each store appended its context to the original messages, so only the last one survived."""

View file

@ -2,16 +2,13 @@ import asyncio
import datetime
import itertools
import json
import os
import queue
import subprocess
import sys
import threading
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from types import CodeType, FrameType, MappingProxyType, TracebackType
from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch
@ -2071,66 +2068,6 @@ async def test_stream_logprobs_without_top_logprobs_reach_caller_as_litellm_logp
]
FINISH_ONLY_LOGPROBS_STREAM_SCRIPT: Final = """
import json
import sys
sys.path.insert(0, sys.argv[1])
import httpx
from openai import OpenAI
import litellm
upstream_choices = (
{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None},
{
"index": 0,
"delta": {},
"logprobs": {"content": [{"token": "<|return|>", "logprob": -0.5, "top_logprobs": []}]},
"finish_reason": "stop",
},
)
body = "".join(
"data: "
+ json.dumps({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [choice]})
+ "\\n\\n"
for choice in upstream_choices
) + "data: [DONE]\\n\\n"
client = OpenAI(
api_key="unused",
base_url="https://provider.example/v1",
http_client=httpx.Client(
transport=httpx.MockTransport(
lambda request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"})
)
),
)
stream = litellm.completion(
model="openai/m", messages=[{"role": "user", "content": "hi"}], stream=True, logprobs=True, client=client
)
print(json.dumps([json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in stream]))
"""
def test_stream_with_logprobs_only_on_finish_chunk_serializes_in_a_fresh_process() -> None:
result: Final = subprocess.run(
[sys.executable, "-I", "-c", FINISH_ONLY_LOGPROBS_STREAM_SCRIPT, str(Path(litellm.__file__).parent.parent)],
capture_output=True,
text=True,
timeout=120,
env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.returncode == 0, result.stdout + result.stderr
wire_choices: Final = tuple(
itertools.chain.from_iterable(chunk["choices"] for chunk in json.loads(result.stdout.strip().splitlines()[-1]))
)
finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")]
assert [choice["finish_reason"] for choice in finish_choices] == ["stop"]
assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == ["<|return|>"]
def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper:
"""A CustomStreamWrapper over the real OpenAI-compatible line iterator,
fed an HTTP 200 SSE body that carries an in-body error payload the way
@ -2398,7 +2335,6 @@ async def test_mid_stream_error_leaves_out_an_api_base_it_cannot_resolve():
assert "API Base:" not in debug_info
@pytest.mark.asyncio
@pytest.mark.parametrize(
"routed_api_base",
@ -7169,7 +7105,6 @@ def test_chunk_creator_skips_full_dump_of_model_response_stream_chunks(
assert getattr(emitted_chunks[1], "custom_chunk_attribute") == "custom value"
def test_openai_chunk_creator_keeps_multiple_choice_logprobs_and_extras(
logging_obj: Logging,
) -> None:
@ -9765,7 +9700,6 @@ async def test_logged_response_keeps_finish_reason_from_last_content_chunk(
assert logged.choices[0].message.content == "Hello, this reply is cut off"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_no_synthetic_finish_reason_logged_when_provider_sent_none(sync_mode: bool, logging_obj: Logging):

View file

@ -6,15 +6,12 @@ import copy
import importlib
import io
import json
import os
import subprocess
import sys
import threading
import time
from collections.abc import Callable, Mapping
from concurrent.futures import Future, ThreadPoolExecutor, wait
from itertools import accumulate, chain
from pathlib import Path
from queue import SimpleQueue
from typing import Final, cast
from unittest.mock import MagicMock, sentinel
@ -23,7 +20,7 @@ import anyio.to_thread
import pytest
import tiktoken
from PIL import Image
from tokenizers import Regex, Tokenizer, models, pre_tokenizers
from tokenizers import Tokenizer
from unittest.mock import AsyncMock, patch
@ -922,48 +919,6 @@ class NeedsToleranceUpdateError(Exception):
# test_tokenizers()
def test_encoding_and_decoding(tmp_path: Path):
sample_text = "Hellö World, this is my input string!"
# openai encoding + decoding
openai_tokens = encode(model="gpt-3.5-turbo", text=sample_text)
openai_text = decode(model="gpt-3.5-turbo", tokens=openai_tokens)
assert openai_text == sample_text
# claude encoding + decoding
claude_tokens = encode(model="claude-3-5-haiku-20241022", text=sample_text)
claude_text = decode(model="claude-3-5-haiku-20241022", tokens=claude_tokens)
assert claude_text == sample_text
# cohere encoding + decoding
cohere_tokens = encode(model="command-nightly", text=sample_text)
cohere_text = decode(model="command-nightly", tokens=cohere_tokens)
assert cohere_text == sample_text
# llama2 encoding + decoding
words = sample_text.split()
result = _run_in_memory_hub(
HUB_ROUND_TRIP_SCRIPT,
{
"hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(
pre_tokenizers.WhitespaceSplit(),
vocab={"[UNK]": 0, **{word: i + 1 for i, word in enumerate(words)}},
)
},
sample_text,
tmp_path,
)
assert result["decoded"] == sample_text
assert result["requested"] == ["hf-internal-testing/llama-tokenizer"]
assert len(result["tokens"]) == len(words)
assert len(result["tokens"]) != len(encode(model="gpt-3.5-turbo", text=sample_text))
# test_encoding_and_decoding()
@ -1162,7 +1117,6 @@ from litellm.utils import (
_HuggingFaceTokenizerLoadOutcomes,
_return_huggingface_tokenizer,
_select_tokenizer_helper,
claude_json_str,
encoding,
)
@ -3564,129 +3518,6 @@ def test_image_estimate_bounds_header_decoding_for_large_inline_payloads() -> No
assert len(decode_header.call_args.args[0]) <= 65_536
HUB_SETUP_SCRIPT: Final = """
import json
import sys
sys.path.insert(0, sys.argv[1])
import httpx
import huggingface_hub
import litellm
served = json.loads(sys.argv[2])
text = sys.argv[3]
requested = []
def handle(request):
repo = request.url.path.lstrip("/").split("/resolve/")[0]
if repo not in served or not request.url.path.endswith("/tokenizer.json"):
return httpx.Response(404)
requested.append(repo)
payload = served[repo].encode()
headers = {"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}
return httpx.Response(200, headers=headers, content=payload if request.method == "GET" else b"")
huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle)))
"""
HUB_TOKENIZER_SCRIPT: Final = HUB_SETUP_SCRIPT + """
litellm.cohere_models = {"command-r-v1"}
litellm.anthropic_models = {"claude-2"}
custom = litellm.create_pretrained_tokenizer("Xenova/llama-3-tokenizer")
print(json.dumps({
"llama2": litellm.token_counter(model="meta-llama/Llama-2-7b-chat", text=text),
"llama3": litellm.token_counter(model="meta-llama/llama-3-70b-instruct", text=text),
"cohere": litellm.token_counter(model="command-r-v1", text=text),
"anthropic": litellm.token_counter(model="claude-2", text=text),
"custom": litellm.token_counter(custom_tokenizer=custom, text=text),
"requested": sorted(set(requested)),
}))
"""
HUB_ROUND_TRIP_SCRIPT: Final = HUB_SETUP_SCRIPT + """
tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text)
print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))}))
"""
HUB_RETRY_DURING_EXIT_SCRIPT: Final = HUB_SETUP_SCRIPT + """
import threading
import time
from litellm.utils import _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer
retry_reached_hub = threading.Event()
def answer_once_the_process_is_exiting(request):
retry_reached_hub.set()
while threading.main_thread().is_alive():
time.sleep(0.01)
time.sleep(0.5)
return handle(request)
now = [0.0]
outcomes = _HuggingFaceTokenizerLoadOutcomes(clock=lambda: now[0])
model = "meta-llama/llama-3-70b-instruct"
_return_huggingface_tokenizer(model, outcomes)
huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(answer_once_the_process_is_exiting)))
now[0] = 3600.0
_return_huggingface_tokenizer(model, outcomes)
print(json.dumps({"retry_reached_hub": retry_reached_hub.wait(timeout=60)}))
"""
def _word_level_tokenizer_json(
pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None
) -> str:
tokenizer: Final = Tokenizer(
models.WordLevel(vocab=dict(vocab) if vocab is not None else {"[UNK]": 0}, unk_token="[UNK]")
)
tokenizer.pre_tokenizer = pre_tokenizer
return tokenizer.to_str()
def _run_in_memory_hub(script: str, served: dict[str, str], text: str, tmp_path: Path) -> dict:
result: Final = subprocess.run(
[
sys.executable,
"-I",
"-c",
script,
str(Path(litellm.__file__).parent.parent),
json.dumps(served),
text,
],
capture_output=True,
text=True,
timeout=60,
env={
**os.environ,
"HF_HOME": str(tmp_path / "home"),
"HF_HUB_CACHE": str(tmp_path / "cache"),
"HF_ENDPOINT": "http://127.0.0.1:9",
"HF_HUB_OFFLINE": "0",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
},
)
assert result.returncode == 0, result.stdout + result.stderr
return json.loads(result.stdout.strip().splitlines()[-1])
def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_tokenizer(tmp_path: Path) -> None:
sample: Final = "Tokenizers disagree: anthropic, tiktoken; llama-2 & llama-3!"
served: Final = {
"hf-internal-testing/llama-tokenizer": _word_level_tokenizer_json(pre_tokenizers.WhitespaceSplit()),
"Xenova/llama-3-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Split(Regex("."), "isolated")),
"Xenova/c4ai-command-r-v01-tokenizer": _word_level_tokenizer_json(pre_tokenizers.Whitespace()),
}
expected: Final = {repo: len(Tokenizer.from_str(payload).encode(sample).ids) for repo, payload in served.items()}
anthropic_count: Final = len(Tokenizer.from_str(claude_json_str).encode(sample).ids)
tiktoken_count: Final = litellm.token_counter(model="gpt-3.5-turbo", text=sample)
assert len({*expected.values(), anthropic_count, tiktoken_count}) == len(expected) + 2
counts: Final = _run_in_memory_hub(HUB_TOKENIZER_SCRIPT, served, sample, tmp_path)
assert counts == {
"llama2": expected["hf-internal-testing/llama-tokenizer"],
"llama3": expected["Xenova/llama-3-tokenizer"],
"cohere": expected["Xenova/c4ai-command-r-v01-tokenizer"],
"anthropic": anthropic_count,
"custom": expected["Xenova/llama-3-tokenizer"],
"requested": sorted(served),
}
def test_custom_openai_tokenizer_without_a_model_counts_with_the_default_encoding() -> None:
sample: Final = "A custom tokenizer passed with no model still counts."
@ -3695,10 +3526,6 @@ def test_custom_openai_tokenizer_without_a_model_counts_with_the_default_encodin
)
def test_process_exits_cleanly_while_a_background_tokenizer_retry_is_downloading(tmp_path: Path) -> None:
assert _run_in_memory_hub(HUB_RETRY_DURING_EXIT_SCRIPT, {}, "", tmp_path) == {"retry_reached_hub": True}
def test_empty_custom_tokenizer_uses_model_tokenizer() -> None:
text_value: Final = "A tokenizer fallback should preserve the model encoding."
custom_count: Final = _get_exact_count_function("gpt-3.5-turbo", {})(text_value)

View file

@ -1,7 +1,5 @@
import copy
import os
import pickle
import subprocess
import sys
from pathlib import Path
from typing import Final, Literal
@ -176,115 +174,6 @@ def test_public_tokenizer_objects_survive_pickle_and_deepcopy(tmp_path: Path) ->
assert pickle.loads(pickle.dumps(litellm.encoding)).encode("hello") == litellm.encoding.encode("hello")
@pytest.mark.parametrize("offline", ("0", "1"))
def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None:
script: Final = """
import json
import sys
from pathlib import Path
sys.path.insert(0, sys.argv[1])
import httpx
import huggingface_hub
from huggingface_hub.errors import LocalEntryNotFoundError
import litellm
payload = sys.argv[2].encode()
offline = sys.argv[3] == "1"
observed = []
def handle(request):
assert not offline, "offline loading issued a request"
if request.url.path.endswith("/tokenizer.json"):
observed.append(request.headers.get("authorization"))
if request.headers.get("authorization") != "Bearer audit-fixture-token":
return httpx.Response(401)
return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"")
if not offline:
huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle)))
try:
tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"]
except LocalEntryNotFoundError:
assert offline
assert observed == []
else:
assert not offline
assert "Bearer audit-fixture-token" in observed
assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World"
assert tuple(Path(sys.argv[4]).rglob("tokenizer.json"))
print("compatible")
"""
result: Final = subprocess.run(
[
sys.executable,
"-I",
"-c",
script,
str(Path(litellm.__file__).parent.parent),
TOKENIZER_JSON,
offline,
str(tmp_path / "cache"),
],
capture_output=True,
text=True,
timeout=30,
env={
**os.environ,
"HF_HOME": str(tmp_path / "home"),
"HF_HUB_CACHE": str(tmp_path / "cache"),
"HF_ENDPOINT": "http://127.0.0.1:9",
"HF_TOKEN": "audit-fixture-token",
"HF_HUB_OFFLINE": offline,
"HF_HUB_DISABLE_IMPLICIT_TOKEN": "0",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
},
)
assert result.returncode == 0, result.stdout + result.stderr
assert result.stdout.strip() == "compatible"
@pytest.mark.parametrize("rust", (None, "0", "1"))
def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None:
script: Final = """
import sys
sys.path.insert(0, sys.argv[1])
def reject_network(event, args):
if event == "socket.connect":
raise AssertionError("tokenizer attempted a network connection")
sys.addaudithook(reject_network)
import litellm
from litellm.rust_bridge.tokenizer import get_encoding
import tiktoken
from tokenizers import Tokenizer
assert isinstance(litellm.encoding, tiktoken.Encoding)
for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"):
encoding = get_encoding(name)
text = "offline café 漢字 🙂" + " " * 64
assert encoding.decode(encoding.encode(text)) == text
ids = litellm.encode(text="hello world")
assert litellm.decode(tokens=ids) == "hello world"
assert litellm.token_counter(model=None, text="hello world") == len(ids)
custom = litellm.create_tokenizer(sys.argv[2])
assert isinstance(custom["tokenizer"], Tokenizer)
custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]")
assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World"
print("compatible")
"""
result: Final = subprocess.run(
[sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON],
capture_output=True,
text=True,
timeout=30,
cwd=tmp_path,
env={
**{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"},
**({"LITELLM_RUST": rust} if rust is not None else {}),
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
"TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"),
},
)
assert result.returncode == 0, result.stdout + result.stderr
assert result.stdout.strip() == "compatible"
assert not (tmp_path / "unused-tokenizer-cache").exists()
@pytest.mark.parametrize("is_pretokenized", (False, True))
def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: bool) -> None:
reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)

View file

@ -14,7 +14,6 @@ import itertools
import json
import logging
import os
import subprocess
import sys
import threading
from collections.abc import Mapping
@ -4910,40 +4909,3 @@ def test_a_register_model_opt_out_for_one_claude_model_leaves_every_other_routed
litellm.register_model({f"anthropic/{opted_out}": {"supports_prompt_caching": False}})
assert supports_anthropic_cache_control(untouched, "anthropic", {"id": "plain-id"}) is True
_MASKED_HOSTED_WEB_CALL_PROBE: Final = """
import json
from litellm.llms.anthropic.common_utils import write_request_hosted_web_call_arguments
message = {
"role": "assistant",
"provider_specific_fields": {
"web_search_calls": [
{
"id": "ws_srvtoolu_1",
"type": "web_search_call",
"status": "completed",
"action": {"type": "search", "query": "deploy token 12345", "queries": ["deploy token 12345"]},
}
]
},
}
write_request_hosted_web_call_arguments(message, {"srvtoolu_1": json.dumps({"query": "deploy token [REDACTED]"})})
print(json.dumps(message["provider_specific_fields"]["web_search_calls"]))
"""
def test_a_masked_replayed_hosted_web_call_serializes_before_its_openai_model_was_ever_built() -> None:
result: Final = subprocess.run(
(sys.executable, "-I", "-c", _MASKED_HOSTED_WEB_CALL_PROBE),
capture_output=True,
text=True,
check=True,
timeout=120,
)
hosted: Final = json.loads(result.stdout.strip().splitlines()[-1])
assert [call["action"] for call in hosted] == [
{"type": "search", "query": "deploy token [REDACTED]", "queries": ["deploy token [REDACTED]"]}
]

View file

@ -437,28 +437,6 @@ def test_initialize_honors_explicit_max_retries(setup_mocks, configured, expecte
assert result["max_retries"] == expected
def test_default_max_retries_env_var_reaches_azure_sdk_client():
import subprocess
import sys
code = (
"from litellm.llms.azure.common_utils import BaseAzureLLM\n"
"client = BaseAzureLLM().get_azure_openai_client("
"api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01',"
" client=None, _is_async=True, litellm_params={}, model='gpt-4')\n"
"print(client.max_retries)"
)
completed = subprocess.run(
[sys.executable, "-c", code],
env={**os.environ, "DEFAULT_MAX_RETRIES": "0"},
capture_output=True,
text=True,
check=True,
)
assert completed.stdout.strip() == "0"
@pytest.mark.parametrize(
"call_type",
[
@ -3268,4 +3246,3 @@ def test_non_chat_calls_on_a_foundry_project_v1_base_reach_azure(
respx_mock.post().mock(return_value=httpx.Response(200, json=upstream_body))
assert call() == "ok"

View file

@ -1,9 +1,7 @@
import base64
import hashlib
import json
import subprocess
import sys
import textwrap
import time
from typing import Final
@ -13,7 +11,6 @@ from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec, rsa
from litellm.llms.base_llm.auth.jwt_signing import (
MISSING_SIGNING_DEPENDENCIES_MESSAGE,
build_jwk,
build_jwks,
jwks_document_json,
@ -183,33 +180,7 @@ class TestBaseSdkImport:
and ``litellm/__init__`` reaches this module through the Anthropic provider, so a
module-level import of either would break ``import litellm`` for every base SDK user."""
def test_module_imports_with_pyjwt_and_cryptography_absent(self):
script: Final = textwrap.dedent(
"""
import sys
class Blocker:
def find_spec(self, name, path=None, target=None):
if name.split(".")[0] in {"jwt", "cryptography"}:
raise ModuleNotFoundError(f"No module named {name!r}")
sys.meta_path.insert(0, Blocker())
import litellm
from litellm.llms.base_llm.auth.jwt_signing import jwks_document_json
try:
jwks_document_json("not a key")
except ImportError as e:
print(e)
"""
)
result: Final = subprocess.run(
[sys.executable, "-I", "-c", script], capture_output=True, text=True, check=False
)
assert result.returncode == 0, result.stderr[-2000:]
assert result.stdout.strip() == MISSING_SIGNING_DEPENDENCIES_MESSAGE
def test_signing_reports_the_missing_extra(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setitem(sys.modules, "jwt", None)
with pytest.raises(ImportError, match="litellm\\[proxy\\]"):
sign_es256_jwt(pem_of(fixed_private_key()), {"sub": "x"})

View file

@ -1,19 +1,16 @@
import asyncio
import gc
import importlib
import json
import logging
import os
import pathlib
import ssl
import threading
import time
import weakref
from collections.abc import AsyncIterator, Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from contextvars import Context
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from unittest.mock import MagicMock, patch
@ -29,8 +26,6 @@ from yarl import URL
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.exceptions import Timeout as LitellmTimeout
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.model_cost_table import publish
from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine, default_token_exchange_engine
from litellm.llms.base_llm.auth.types import MintedToken, TokenExchangeSpec
from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
@ -46,7 +41,6 @@ from litellm.llms.custom_httpx.http_handler import (
realtime_ssl_for_url,
)
from litellm.types.llms.custom_http import VerifyTypes
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@pytest.mark.asyncio
@ -879,116 +873,6 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout:
assert client.timeout.read == 300.0
async def _read_http_request(reader: asyncio.StreamReader) -> None:
raw = b""
while b"\r\n\r\n" not in raw:
chunk = await reader.read(1024)
if not chunk:
return
raw += chunk
head, _, body = raw.partition(b"\r\n\r\n")
content_length = next(
(int(line.split(b":", 1)[1]) for line in head.split(b"\r\n") if line.lower().startswith(b"content-length")),
0,
)
while len(body) < content_length:
body += await reader.read(content_length - len(body))
@pytest.mark.asyncio
async def test_init_held_async_handler_survives_external_client_close():
handler = AsyncHTTPHandler(timeout=42.5)
held_client = handler.client
await held_client.aclose()
assert held_client.is_closed
async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
await _read_http_request(reader)
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
await writer.drain()
writer.close()
server = await asyncio.start_server(respond, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
try:
response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []})
finally:
server.close()
await server.wait_closed()
assert response.status_code == 200
assert handler.client is not held_client
assert handler.client.timeout == httpx.Timeout(42.5)
await handler.close()
@pytest.mark.asyncio
async def test_init_held_async_handler_survives_evicted_client_close():
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
cache = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0))
handler = AsyncHTTPHandler(timeout=42.5)
held_client = handler.client
cache.set_cache("init-held-handler", handler, litellm_owned_client=True, ttl=0)
await asyncio.sleep(0.02)
assert cache.get_cache("init-held-handler") is None
await asyncio.sleep(0.05)
assert held_client.is_closed
async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
await _read_http_request(reader)
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
await writer.drain()
writer.close()
server = await asyncio.start_server(respond, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
try:
response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []})
finally:
server.close()
await server.wait_closed()
assert response.status_code == 200
assert handler.client is not held_client
assert handler.client.timeout == httpx.Timeout(42.5)
await handler.close()
def test_init_held_sync_handler_recreates_closed_client():
from http.server import BaseHTTPRequestHandler, HTTPServer
class OkRequestHandler(BaseHTTPRequestHandler):
def do_GET(self):
self.send_response(200)
self.send_header("Content-Length", "2")
self.end_headers()
self.wfile.write(b"ok")
def log_message(self, format, *args):
pass
handler = HTTPHandler(timeout=7)
held_client = handler.client
held_client.close()
server = HTTPServer(("127.0.0.1", 0), OkRequestHandler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
response = handler.get(f"http://127.0.0.1:{server.server_port}/")
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
assert response.status_code == 200
assert handler.client is not held_client
assert handler.client.timeout == httpx.Timeout(7)
handler.close()
def test_caller_supplied_sync_client_is_not_replaced_when_closed():
supplied = httpx.Client()
handler = HTTPHandler(client=supplied)
@ -2245,66 +2129,6 @@ async def test_http2_disabled_by_default(monkeypatch: pytest.MonkeyPatch):
assert AsyncHTTPHandler.should_use_aiohttp_transport() is True
@pytest.fixture()
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
yield
record_vcr_outcome(request, vcr)
@pytest.fixture(scope="function")
def isolate_litellm_state():
"""
Per-function isolation fixture.
Resets litellm globals to their true defaults before each test and
restores them afterward, so tests don't leak side effects.
Works safely under pytest-xdist parallel execution.
"""
original_state = {}
for attr in (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
):
if hasattr(litellm, attr):
val = getattr(litellm, attr)
original_state[attr] = val.copy() if val else []
for attr in ("pre_call_rules", "post_call_rules"):
if hasattr(litellm, attr):
val = getattr(litellm, attr)
original_state[attr] = val.copy() if val else []
for attr in _SCALAR_DEFAULTS:
if hasattr(litellm, attr):
original_state[attr] = getattr(litellm, attr)
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
for attr in (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
"pre_call_rules",
"post_call_rules",
):
if hasattr(litellm, attr):
setattr(litellm, attr, [])
for attr, default_val in _SCALAR_DEFAULTS.items():
if hasattr(litellm, attr):
setattr(litellm, attr, default_val)
yield
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
for attr, original_value in original_state.items():
if hasattr(litellm, attr):
setattr(litellm, attr, original_value)
publish(litellm.model_cost)
_SCALAR_DEFAULTS = {
"num_retries": getattr(litellm, "num_retries", None),
"num_retries_per_request": getattr(litellm, "num_retries_per_request", None),
@ -2326,73 +2150,6 @@ _SCALAR_DEFAULTS = {
}
@pytest.fixture(scope="module")
def setup_and_teardown():
"""
Module-scoped setup. Reloads litellm only in single-process mode
(skipped under xdist to avoid cross-worker interference).
"""
import litellm
worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
if worker_id is None:
importlib.reload(litellm)
try:
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
import litellm.proxy.proxy_server
importlib.reload(litellm.proxy.proxy_server)
except Exception:
pass
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
yield
_SERVER_DELAY_S = 5
_PER_REQUEST_TIMEOUT_S = 1.0
_CLIENT_DEFAULT_TIMEOUT_S = 60.0
class _SlowHandler(BaseHTTPRequestHandler):
def do_POST(self):
time.sleep(_SERVER_DELAY_S)
try:
self.send_response(200)
self.end_headers()
self.wfile.write(b"{}")
except OSError:
pass
def log_message(self, *args):
pass
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_post_delay_exceeds_per_request_timeout_raises():
server = ThreadingHTTPServer(("127.0.0.1", 0), _SlowHandler)
threading.Thread(target=server.serve_forever, daemon=True).start()
host, port = server.server_address
handler = get_httpx_client(params={"timeout": _CLIENT_DEFAULT_TIMEOUT_S})
try:
with pytest.raises(LitellmTimeout):
handler.post(
f"http://{host}:{port}/delay",
headers={"content-type": "application/json"},
data=json.dumps({"model": "claude", "messages": []}),
timeout=_PER_REQUEST_TIMEOUT_S,
)
except MaskedHTTPStatusError as e:
pytest.skip(f"httpbin.org unavailable: {e}")
finally:
handler.close()
server.shutdown()
server.server_close()
def test_realtime_ssl_for_url_sends_no_tls_argument_for_a_plain_ws_endpoint() -> None:
assert (
realtime_ssl_for_url("ws://127.0.0.1:8080/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent")

View file

@ -606,26 +606,6 @@ def cost_recorder(monkeypatch: pytest.MonkeyPatch) -> CostRecorder:
return recorder
@pytest.mark.asyncio
@pytest.mark.parametrize("upstream", ["chat/completions", "messages", "responses", "gemini"])
@pytest.mark.parametrize("ingress", ["chat", "responses"])
async def test_reported_list_cost_is_billed_instead_of_rows(
go_catalog: None, cost_recorder: CostRecorder, upstream: str, ingress: str
) -> None:
model: Final = f"fixture-{upstream.split('/')[0]}"
with respx.mock(using="httpx") as transport:
transport.post().mock(
return_value=httpx.Response(200, json=go_priced_reply(upstream, model, GO_REPORTED_USAGE))
)
kwargs: Final = {"model": f"opencode_go/{model}", "api_key": "fake-go-key"}
if ingress == "chat":
response: Final = await litellm.acompletion(messages=[{"role": "user", "content": "hello"}], **kwargs)
else:
response = await litellm.aresponses(input="hello", **kwargs)
assert response._hidden_params["additional_headers"][RESPONSE_COST_HEADER] == pytest.approx(0.00412)
assert await logged_cost(cost_recorder) == pytest.approx(0.00412)
@pytest.mark.asyncio
@pytest.mark.parametrize("usage_extra", [{}, {"cost": 0}, {"opencode": {"billing_source": "go"}}])
async def test_rows_price_the_request_when_no_list_cost_is_reported(

View file

@ -1,11 +1,4 @@
import asyncio
import base64
import uuid
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass, field
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from threading import Event, Thread
from typing import Final, Literal
from unittest.mock import patch
@ -298,131 +291,6 @@ async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: L
_PNG_BYTES: Final = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
_LOOP_PROOF_TIMEOUT_SECONDS: Final = 5
@dataclass(frozen=True, slots=True)
class _LoopProbe:
request_arrived: Event = field(default_factory=Event)
loop_ran: Event = field(default_factory=Event)
loop_ran_while_download_pending: Event = field(default_factory=Event)
def _png_handler_waiting_for_loop(probe: _LoopProbe) -> type[BaseHTTPRequestHandler]:
class _PngHandler(BaseHTTPRequestHandler):
def do_GET(self) -> None:
probe.request_arrived.set()
if probe.loop_ran.wait(_LOOP_PROOF_TIMEOUT_SECONDS):
probe.loop_ran_while_download_pending.set()
self.send_response(200)
self.send_header("Content-Type", "image/png")
self.send_header("Content-Length", str(len(_PNG_BYTES)))
self.end_headers()
self.wfile.write(_PNG_BYTES)
def log_message(self, format: str, *args: object) -> None:
return None
return _PngHandler
@contextmanager
def _png_server_waiting_for_loop(probe: _LoopProbe) -> Generator[str, None, None]:
with ThreadingHTTPServer(("127.0.0.1", 0), _png_handler_waiting_for_loop(probe)) as server:
thread: Final = Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield f"http://127.0.0.1:{server.server_address[1]}"
finally:
server.shutdown()
async def _run_loop_once_download_starts(probe: _LoopProbe) -> None:
assert await asyncio.to_thread(probe.request_arrived.wait, _LOOP_PROOF_TIMEOUT_SECONDS)
probe.loop_ran.set()
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
async def test_remote_image_downloads_without_blocking_the_event_loop(
custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "user_url_validation", False)
probe: Final = _LoopProbe()
with _png_server_waiting_for_loop(probe) as base_url:
image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png"
_, contents = await asyncio.gather(
_run_loop_once_download_starts(probe),
gemini_count_contents(
(
{
"role": "user",
"content": [
{"type": "text", "text": "Describe it"},
{"type": "image", "source": {"type": "url", "url": image_url}},
],
},
),
"anthropic",
"gemini-2.5-flash",
custom_llm_provider,
{},
),
)
assert contents == [
{
"role": "user",
"parts": [
{"text": "Describe it"},
{"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}},
],
}
]
assert probe.loop_ran_while_download_pending.is_set()
class _MissingImageHandler(BaseHTTPRequestHandler):
def do_GET(self) -> None:
self.send_response(404)
self.send_header("Content-Length", "0")
self.end_headers()
def log_message(self, format: str, *args: object) -> None:
return None
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
async def test_image_that_cannot_be_fetched_returns_an_error_value(
custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "user_url_validation", False)
with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server:
Thread(target=server.serve_forever, daemon=True).start()
try:
contents = await gemini_count_contents(
(
{
"role": "user",
"content": [
{
"type": "image",
"source": {
"type": "url",
"url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png",
},
}
],
},
),
"anthropic",
"gemini-2.5-flash",
custom_llm_provider,
{},
)
finally:
server.shutdown()
assert isinstance(contents, VertexGeminiCountContextError)
async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None:

View file

@ -1,16 +1,12 @@
import json
import sys
import threading
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import chain, count, repeat
from itertools import chain, repeat
from typing import Final
import httpx
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@ -1185,158 +1181,6 @@ async def test_streaming_continuation_request_is_sent_with_the_callers_timeout(i
assert read_timeouts == [7.5, 7.5]
@dataclass(frozen=True, slots=True)
class _TokenEndpoint:
credentials: str
down: threading.Event
def _service_account_json(token_uri: str) -> str:
private_key: Final = (
rsa.generate_private_key(public_exponent=65537, key_size=2048)
.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
.decode()
)
return json.dumps(
{
"type": "service_account",
"project_id": "p",
"private_key_id": "k",
"private_key": private_key,
"client_email": "sa@p.iam.gserviceaccount.com",
"token_uri": token_uri,
}
)
@pytest.fixture
def token_endpoint() -> Iterator[_TokenEndpoint]:
down: Final = threading.Event()
minted: Final = count(1)
class Handler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
self.rfile.read(int(self.headers["Content-Length"]))
refused: Final = down.is_set()
payload: Final = json.dumps(
{"error": "invalid_grant"} if refused else {"access_token": f"token-{next(minted)}", "expires_in": 0}
).encode()
self.send_response(400 if refused else 200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, *_args: object) -> None:
return None
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
threading.Thread(target=server.serve_forever, daemon=True).start()
yield _TokenEndpoint(_service_account_json(f"http://127.0.0.1:{server.server_port}/token"), down)
server.shutdown()
server.server_close()
def _google_taking_each_access_token_once(
hops: Sequence[httpx.Response], used: list[str], on_answer: Callable[[], object]
) -> Callable[[httpx.Request], httpx.Response]:
def handler(request: httpx.Request) -> httpx.Response:
token: Final = request.headers["authorization"]
if token in used:
return httpx.Response(401, json={"error": {"code": 401, "status": "UNAUTHENTICATED"}})
used.append(token)
on_answer()
return hops[len(used) - 1]
return handler
async def _vertex_answer(
is_async: bool,
stream: bool,
handler: Callable[[httpx.Request], httpx.Response],
credentials: str,
extra_headers: Mapping[str, str] | None = None,
) -> tuple[str, list[str], list[object]]:
kwargs: Final = {
"model": "vertex_ai/gemini-3.8-flash",
"extra_headers": extra_headers,
"messages": list(MESSAGES),
"max_tokens": 100_000,
"stream": stream,
"vertex_project": "p",
"vertex_location": "global",
"vertex_credentials": credentials,
}
transport: Final = httpx.MockTransport(handler)
response: Final = (
await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport))
if is_async
else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport)))
)
if stream:
chunks: Final = [chunk async for chunk in response] if is_async else list(response)
return _stream_text(chunks), _stream_finish_reasons(chunks), _client_tokens(chunks)
choice: Final = response.choices[0]
fields: Final = choice.message.provider_specific_fields or {}
return choice.message.content, [choice.finish_reason], [fields["continuation_token"]] if fields else []
def _two_hops(stream: bool) -> tuple[httpx.Response, httpx.Response]:
return _stream_hops() if stream else (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2))
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.asyncio
async def test_vertex_long_decode_finishes_when_each_continuation_outlives_the_access_token_before_it(
is_async: bool, stream: bool, token_endpoint: _TokenEndpoint
):
used: Final[list[str]] = []
google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, lambda: None)
answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials)
assert answer == ("hello world", ["stop"], [])
assert len(used) == 2
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.asyncio
async def test_vertex_long_decode_returns_the_answer_so_far_with_its_token_when_no_new_access_token_can_be_had(
is_async: bool, stream: bool, token_endpoint: _TokenEndpoint
):
used: Final[list[str]] = []
google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, token_endpoint.down.set)
answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials)
assert answer == ("hello", ["length"], ["dG9rLTE="])
assert len(used) == 1
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.asyncio
async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_every_continuation(
is_async: bool, stream: bool, token_endpoint: _TokenEndpoint
):
hops: Final = _two_hops(stream)
seen: Final[list[str]] = []
def google(request: httpx.Request) -> httpx.Response:
seen.append(request.headers["authorization"])
return hops[len(seen) - 1]
answer: Final = await _vertex_answer(
is_async, stream, google, token_endpoint.credentials, {"Authorization": "Bearer mine"}
)
assert answer == ("hello world", ["stop"], [])
assert seen == ["Bearer mine", "Bearer mine"]
BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS"
ARGON: Final = "gemini/gemini-4-argon"
UNMAPPED_GEMINI: Final = "gemini/gemini-99-unreleased"
@ -1555,46 +1399,6 @@ async def test_env_budget_leaves_a_snake_case_generation_config_from_extra_body_
assert _output_limits(sent[0]) == {}
def _gemini_chat_models_answering_no_text() -> tuple[str, ...]:
rows: Final = (
(name, row)
for name, row in litellm.model_cost.items()
if row.get("mode") == "chat" and "text" not in row.get("supported_output_modalities", ("text",))
)
return tuple(
name if name.startswith("gemini/") else f"vertex_ai/{name}"
for name, row in rows
if row.get("litellm_provider") in ("gemini", "vertex_ai-language-models")
)
@pytest.mark.asyncio
async def test_env_budget_skips_every_chat_model_the_cost_map_lists_as_answering_no_text(
monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint
):
monkeypatch.setenv(BUDGET_ENV, "50")
models: Final = _gemini_chat_models_answering_no_text()
sent: Final[list[dict[str, object]]] = []
replies: Final = tuple(httpx.Response(200, json=HOP_2) for _ in models)
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent)))
for model in models:
await litellm.acompletion(
model=model,
messages=list(MESSAGES),
api_key="fake-gemini-key",
vertex_project="p",
vertex_location="us-central1",
vertex_credentials=token_endpoint.credentials,
client=client,
)
assert models
assert {model: _output_limits(request) for model, request in zip(models, sent, strict=True)} == dict.fromkeys(
models, {}
)
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.asyncio
async def test_env_budget_leaves_a_continuation_token_resume_with_no_limit_alone(
@ -1619,28 +1423,6 @@ async def test_env_budget_leaves_a_continuation_token_resume_with_no_limit_alone
assert _output_limits(sent[0]) == {}
@pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "vertex_ai/gemini-2.5-flash", ARGON])
@pytest.mark.asyncio
async def test_env_budget_above_a_models_output_limit_sends_that_limit_instead(
monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint, model: str
):
budget: Final = 2 * _cost_map_max_output_tokens(model)
monkeypatch.setenv(BUDGET_ENV, str(budget))
sent: Final[list[dict[str, object]]] = []
await litellm.acompletion(
model=model,
messages=list(MESSAGES),
api_key="fake-gemini-key",
vertex_project="p",
vertex_location="us-central1",
vertex_credentials=token_endpoint.credentials,
client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))),
)
assert _output_limits(sent[0]) == {"max_output_tokens": _cost_map_max_output_tokens(model)}
@pytest.mark.asyncio
async def test_env_budget_never_goes_above_the_cost_map_limit_when_a_deployment_overrides_it(
monkeypatch: pytest.MonkeyPatch,

File diff suppressed because it is too large Load diff

View file

@ -1,10 +1,6 @@
import asyncio
import re
import shutil
import subprocess
import sys
from collections.abc import Mapping
from pathlib import Path
from typing import Final
import pytest
@ -26,7 +22,6 @@ from litellm.proxy.auth.master_key_boot_check import (
UnsafeMasterKeyError,
UnsafeMasterKeyReason,
UnsafeMasterKeyRefused,
announce_on_stderr_at_exit,
enforce_master_key_boot_verdict,
master_key_boot_verdict,
render_refusal,
@ -361,24 +356,6 @@ def test_migration_with_an_exported_key_replaces_it_in_place_and_numbers_every_s
assert text.index("os.environ/") < text.index(MIGRATE_FROM_MASTER_KEY_ENV_VAR + "=") < text.index("Start the proxy")
@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl")
@pytest.mark.parametrize("from_master_key", [PUBLICLY_KNOWN_KEY, "", " " + PUBLICLY_KNOWN_KEY])
def test_printed_migration_commands_save_both_keys_to_the_env_file(tmp_path: Path, from_master_key: str):
from dotenv import dotenv_values
text = render_refusal(
_refusal(migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1))
)
commands = [line.strip() for line in text.splitlines() if line.strip().startswith("echo ")]
subprocess.run(["bash", "-c", "\n".join(commands)], cwd=tmp_path, capture_output=True, text=True, check=True)
saved = dotenv_values(tmp_path / ".env")
assert saved[MIGRATE_FROM_MASTER_KEY_ENV_VAR] == from_master_key
assert _verdict(saved[MASTER_KEY_ENV_VAR]) == SafeMasterKey()
assert len(commands) == 2
@pytest.mark.parametrize("migration", [_MIGRATION, None])
def test_override_hint_is_the_last_paragraph(migration: StoredSecretsMigration | None):
text = render_refusal(_refusal(migration=migration))
@ -388,30 +365,6 @@ def test_override_hint_is_the_last_paragraph(migration: StoredSecretsMigration |
assert f"general_settings.{WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING}" in last_paragraph
@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl")
def test_printed_command_saves_a_key_the_boot_check_accepts(tmp_path: Path):
completed = subprocess.run(
["bash", "-c", GENERATE_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True
)
saved = (tmp_path / ".env").read_text()
match = re.fullmatch(rf"{MASTER_KEY_ENV_VAR}=(sk-[0-9a-f]{{64}})\n", saved)
assert match is not None
assert completed.stdout == saved
assert _verdict(match.group(1)) == SafeMasterKey()
@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl")
def test_rotation_command_prints_a_key_the_boot_check_accepts_and_saves_nothing(tmp_path: Path):
completed = subprocess.run(
["bash", "-c", PRINT_NEW_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True
)
assert re.fullmatch(r"sk-[0-9a-f]{64}\n", completed.stdout) is not None
assert _verdict(completed.stdout.strip()) == SafeMasterKey()
assert list(tmp_path.iterdir()) == []
def test_refusal_announces_the_fix_and_aborts_the_boot():
announced: list[str] = []
refusal = _refusal(source=ConfigFileSource(config_file_path="/app/config.yaml"))
@ -429,19 +382,3 @@ def test_safe_and_overridden_keys_boot_without_announcing(verdict: MasterKeyBoot
enforce_master_key_boot_verdict(verdict, announce=announced.append)
assert announced == []
def test_announced_fix_is_the_last_thing_a_crashing_process_prints():
crash_after_announcing = (
"from litellm.proxy.auth.master_key_boot_check import announce_on_stderr_at_exit\n"
"announce_on_stderr_at_exit('THE FIX')\n"
"print('buffered stdout')\n"
"raise RuntimeError('lifespan failed')\n"
)
completed = subprocess.run([sys.executable, "-c", crash_after_announcing], capture_output=True, text=True)
assert completed.returncode != 0
assert "RuntimeError: lifespan failed" in completed.stderr
assert completed.stderr.endswith("THE FIX")
assert completed.stdout == "buffered stdout\n"

View file

@ -3,16 +3,12 @@ import contextvars
import json
import logging
import os
import subprocess
import sys
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager, nullcontext
from datetime import datetime, timedelta, timezone
from functools import partial
from io import BytesIO
from pathlib import Path
from textwrap import dedent
from types import MappingProxyType, SimpleNamespace
from typing import Final, Literal
from unittest.mock import ANY, AsyncMock, MagicMock, patch
@ -12170,129 +12166,6 @@ def test_the_request_span_parents_datastore_spans_of_the_request(monkeypatch, al
assert contextvars.copy_context().run(ambient_parent_after_auth) is request_span
_RECORDING_DDTRACE = dedent(
"""
import functools
import inspect
class _Span:
def __enter__(self):
return self
def __exit__(self, *exc):
return None
class _Tracer:
def __init__(self):
self.spans = []
def wrap(self, name=None, **kwargs):
def decorator(f):
span_name = name or f"{f.__module__}.{f.__name__}"
if inspect.iscoroutinefunction(f):
@functools.wraps(f)
async def async_wrapped(*args, **kw):
self.spans.append(span_name)
return await f(*args, **kw)
return async_wrapped
@functools.wraps(f)
def wrapped(*args, **kw):
self.spans.append(span_name)
return f(*args, **kw)
return wrapped
return decorator
def trace(self, name, **kwargs):
return _Span()
def current_span(self):
return None
def current_root_span(self):
return None
tracer = _Tracer()
"""
)
_DDTRACE_AUTH_PROBE = dedent(
"""
import asyncio
import json
import ddtrace
from fastapi import Request
from starlette.datastructures import URL
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
proxy_server.master_key = "sk-probe"
async def auth(api_key):
request = Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/chat/completions"})
request._url = URL(url="/chat/completions")
try:
await user_api_key_auth(
request=request,
api_key=api_key,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
custom_litellm_key_header=None,
)
return "accepted"
except ProxyException:
return "rejected"
async def main():
outcomes = [await auth("Bearer sk-probe"), await auth("Bearer sk-wrong")]
print(json.dumps({"outcomes": outcomes, "spans": ddtrace.tracer.spans}))
asyncio.run(main())
"""
)
def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(tmp_path: Path):
stub_root = tmp_path / "site"
(stub_root / "ddtrace").mkdir(parents=True)
(stub_root / "ddtrace" / "__init__.py").write_text(_RECORDING_DDTRACE)
probe = tmp_path / "probe.py"
probe.write_text(_DDTRACE_AUTH_PROBE)
repo_root = Path(litellm.__file__).resolve().parent.parent
env = {
**os.environ,
"USE_DDTRACE": "true",
"PYTHONPATH": os.pathsep.join(
[str(stub_root), str(repo_root)] + [p for p in (os.environ.get("PYTHONPATH"),) if p]
),
}
result = subprocess.run(
[sys.executable, str(probe)], env=env, cwd=repo_root, capture_output=True, text=True, check=False
)
assert result.returncode == 0, result.stderr[-4000:]
report = json.loads(result.stdout.strip().splitlines()[-1])
assert report["outcomes"] == ["accepted", "rejected"]
auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span]
@pytest.mark.asyncio
@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard-return", "proxy-admin-return"])
async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_admin):

View file

@ -1,9 +1,6 @@
import json
import os
import pathlib
import shlex
import stat
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError
@ -56,7 +53,6 @@ def _owners(*backup_paths):
CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings"
AUTH_MODULE = "litellm.proxy.client.cli.commands.auth"
@pytest.fixture
def paths(tmp_path):
return tmp_path / "claude" / "settings.json", tmp_path / "backup.json"
@ -174,35 +170,6 @@ class TestConfigureClaudeSettings:
class TestLiteApiKeyHelper:
@pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX")
def test_sh_runs_the_lite_on_path_against_this_proxy(self, tmp_path):
lite = tmp_path / "my tools" / "lite"
lite.parent.mkdir()
lite.write_text('#!/bin/sh\nprintf "%s\\n" "$0" "$@"\n')
lite.chmod(0o700)
helper = lite_api_key_helper("https://proxy.example.com/", platform="linux", which=lambda name: str(lite))
assert helper is not None
ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True, check=True)
assert ran.stdout.splitlines() == [str(lite), "--base-url", "https://proxy.example.com", "auth", "print-token"]
@pytest.mark.skipif(sys.platform.startswith("win"), reason="runs the command through sh, as Claude Code does on POSIX")
def test_a_lite_found_through_a_relative_path_entry_keeps_running_from_any_directory(self, tmp_path, monkeypatch):
for repo, name in ((tmp_path / "login-dir", "mine"), (tmp_path / "other-repo", "planted")):
lite = repo / "bin" / "lite"
lite.parent.mkdir(parents=True)
lite.write_text(f"#!/bin/sh\necho {name}\n")
lite.chmod(0o700)
monkeypatch.chdir(tmp_path / "login-dir")
helper = lite_api_key_helper("https://proxy.example.com", platform="linux", which=lambda name: "bin/lite")
assert helper is not None
ran = subprocess.run(
["sh", "-c", helper.command], cwd=tmp_path / "other-repo", capture_output=True, text=True, check=True
)
assert ran.stdout == "mine\n"
def test_quotes_every_token_for_cmd_on_windows(self):
helper = lite_api_key_helper(
@ -262,23 +229,6 @@ class TestInstalledExecutable:
path.write_text(f"#!/bin/sh\n{body}\n")
path.chmod(0o700)
def test_the_helper_runs_the_lite_installed_beside_this_python_over_a_shim_earlier_on_path(
self, tmp_path, monkeypatch
):
self._script(tmp_path / "shims" / "lite", "echo 'pyenv: lite: command not found' >&2; exit 127")
self._script(tmp_path / "venv" / "bin" / "lite", "echo real")
monkeypatch.setenv("PATH", f"{tmp_path / 'shims'}{os.pathsep}{os.environ['PATH']}")
helper = lite_api_key_helper(
"https://proxy.example.com",
platform="linux",
which=lambda name: installed_executable(name, scripts_dir=str(tmp_path / "venv" / "bin")),
)
assert helper is not None
ran = subprocess.run(["sh", "-c", helper.command], capture_output=True, text=True)
assert (ran.returncode, ran.stdout) == (0, "real\n")
def test_falls_back_to_path_when_this_python_has_no_lite(self, tmp_path, monkeypatch):
self._script(tmp_path / "pipx" / "lite", "echo pipx")
(tmp_path / "venv" / "bin").mkdir(parents=True)
@ -931,18 +881,6 @@ class TestStatusLine:
merged = merge_claude_settings({}, PROXY, StaticToken("tok"), status_line=self.COMMAND)
assert merged["statusLine"] == {"type": "command", "command": self.COMMAND}
def test_the_installed_script_is_the_bundled_one_and_the_command_runs_this_interpreter(self, tmp_path):
from litellm.proxy.client.cli.commands import statusline_script
script = tmp_path / "lite" / "statusline.py"
command = install_statusline_script(script)
assert script.read_bytes().split(b"\n", 1)[1] == pathlib.Path(statusline_script.__file__).read_bytes()
assert shlex.split(command) == [sys.executable, str(script)]
assert command == statusline_command(script)
assert stat.S_IMODE(script.stat().st_mode) == 0o600
assert stat.S_IMODE(script.parent.stat().st_mode) == 0o700
assert install_statusline_script(script) == command
def test_a_reinstall_replaces_the_script_in_one_step_and_a_refused_one_leaves_the_old_script_whole(self, tmp_path):
# Claude Code may be running the script at the moment `lite` reinstalls it; the file it has open
# must stay complete, and a reinstall that cannot land must not leave a truncated script behind.

View file

@ -2,15 +2,12 @@
import _hashlib
import asyncio
import copy
import hashlib
import importlib
import inspect
import logging
import os
import sys
import tempfile
import warnings
from collections.abc import Awaitable, Callable, Iterator
from contextlib import contextmanager
from types import ModuleType
@ -36,6 +33,7 @@ from litellm.integrations.otel.model.payloads import ServiceSpanData
from litellm.integrations.otel.model.spans import service_span_name
from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault
from tests._master_key import MASTER_KEY
from tests._module_state import collect_flushable_caches, flush_caches, restore_mutable_state, snapshot_mutable_state
class StubClientNotConnectedError(ClientNotConnectedError):
@ -56,104 +54,9 @@ def disconnected_prisma() -> DisconnectedPrisma:
return DisconnectedPrisma()
# Top-level assignments of these types are the ones importlib.reload(litellm)
# would have effectively reset. We snapshot them at conftest import time and
# deep-copy the snapshot back before every test.
_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes)
def _snapshot_mutable_state(module):
"""Capture a per-module snapshot of primitive and collection attributes."""
snapshot = {}
for attr in list(vars(module)):
if attr.startswith("_"):
continue
try:
value = getattr(module, attr)
except Exception as exc:
warnings.warn(
f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}",
stacklevel=2,
)
continue
if value is None or isinstance(value, _SNAPSHOT_TYPES):
try:
snapshot[attr] = _restored_value(value)
except Exception as exc:
warnings.warn(
f"conftest: could not snapshot {module.__name__}.{attr}: {exc}",
stacklevel=2,
)
return snapshot
_MUTABLE_CONTAINERS = (list, dict, set, bytearray)
def _holds_mutable_container(value) -> bool:
if isinstance(value, _MUTABLE_CONTAINERS):
return True
if isinstance(value, tuple):
return any(_holds_mutable_container(element) for element in value)
return False
def _restored_value(value):
return copy.deepcopy(value) if _holds_mutable_container(value) else value
def _restore_mutable_state(module, snapshot):
for attr, default in snapshot.items():
try:
setattr(module, attr, _restored_value(default))
except Exception as exc:
warnings.warn(
f"conftest: could not restore {module.__name__}.{attr}: {exc}",
stacklevel=2,
)
def _collect_flushable_caches():
"""Return (module, attr) pairs whose values expose flush_cache()."""
targets = []
for module in (litellm, litellm.proxy.proxy_server):
for attr in list(vars(module)):
if attr.startswith("_"):
continue
try:
value = getattr(module, attr)
except Exception:
continue
# Only instances — a class reference has an unbound flush_cache
# that can't be called without a self argument.
if inspect.isclass(value) or inspect.ismodule(value):
continue
if callable(getattr(value, "flush_cache", None)):
targets.append((module, attr))
return targets
def _flush_caches(targets):
for module, attr in targets:
try:
value = getattr(module, attr)
except Exception:
continue
flush = getattr(value, "flush_cache", None)
if callable(flush):
try:
flush()
except Exception as exc:
warnings.warn(
f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}",
stacklevel=2,
)
# Snapshot once at conftest import — these are the "clean" module states.
_LITELLM_STATE = _snapshot_mutable_state(litellm)
_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server)
_FLUSHABLE_CACHES = _collect_flushable_caches()
_LITELLM_STATE: Final = snapshot_mutable_state(litellm)
_PROXY_SERVER_STATE: Final = snapshot_mutable_state(litellm.proxy.proxy_server)
_FLUSHABLE_CACHES: Final = collect_flushable_caches((litellm, litellm.proxy.proxy_server))
_PROXY_APP_LIFESPAN: Final = litellm.proxy.proxy_server.app.router.lifespan_context
@ -182,9 +85,9 @@ def setup_and_teardown():
Use pytest's monkeypatch.setattr() or a local fixture for those
cases — don't rely on this autouse fixture to undo them.
"""
_restore_mutable_state(litellm, _LITELLM_STATE)
_restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE)
_flush_caches(_FLUSHABLE_CACHES)
restore_mutable_state(litellm, _LITELLM_STATE)
restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE)
flush_caches(_FLUSHABLE_CACHES)
loop = asyncio.get_event_loop_policy().new_event_loop()
asyncio.set_event_loop(loop)

View file

@ -1,7 +1,3 @@
import os
import subprocess
import sys
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
@ -69,19 +65,3 @@ async def test_execute_code_loop_dispatches_litellm_skill_tool():
mock_exec.assert_awaited_once()
assert mock_exec.await_args.args[0] == SKILL_TOOL_NAME
assert result is final_response
def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks():
script: Final = (
"import litellm; "
"litellm.callbacks = []; "
"import litellm.proxy.hooks; "
"assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'"
)
result: Final = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"

View file

@ -1,11 +1,15 @@
import asyncio
import importlib
import os
from typing import Final
import pytest
import litellm
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from tests._module_state import collect_flushable_caches, flush_caches, restore_mutable_state, snapshot_mutable_state
_LITELLM_STATE: Final = snapshot_mutable_state(litellm)
_FLUSHABLE_CACHES: Final = collect_flushable_caches((litellm,))
@pytest.fixture(scope="session")
@ -20,11 +24,8 @@ def event_loop():
@pytest.fixture(scope="function", autouse=True)
def setup_and_teardown():
"""
This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained.
"""
importlib.reload(litellm)
import asyncio
restore_mutable_state(litellm, _LITELLM_STATE)
flush_caches(_FLUSHABLE_CACHES)
loop = asyncio.get_event_loop_policy().new_event_loop()
asyncio.set_event_loop(loop)

View file

@ -1636,21 +1636,3 @@ class TestBoundedOpenAPISpecLoading:
else:
assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}}
assert destination.call_count == 1
def test_openapi_generator_import_does_not_require_mcp_sdk() -> None:
import subprocess
import sys
script = """
import builtins
original_import = builtins.__import__
def without_mcp(name, *args, **kwargs):
if name == 'mcp' or name.startswith('mcp.'):
raise ModuleNotFoundError('MCP SDK unavailable')
return original_import(name, *args, **kwargs)
builtins.__import__ = without_mcp
import litellm.proxy.mcp_server.openapi_to_mcp_generator
"""
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True)
assert result.returncode == 0, result.stderr

View file

@ -1,5 +1,4 @@
import signal
import subprocess
import sys
from dataclasses import dataclass, field
from pathlib import Path
@ -9,9 +8,7 @@ from typing import Final
import anyio
import pytest
from litellm.proxy.rust_forward.platform import ChildProcess, Platform, current_platform
from litellm.proxy.rust_forward.platform.windows import WindowsPlatform
from litellm.proxy.rust_forward.supervisor import supervise
from litellm.proxy.rust_forward.platform import ChildProcess, current_platform
from litellm.proxy.rust_forward.supervisor.children import ChildSpec
from litellm.proxy.rust_forward.supervisor.lifecycle import (
Action,
@ -24,42 +21,6 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import (
Spawn,
)
from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, drive, exit_like, open_children, spawn_subprocess
from tests.unit.proxy.rust_forward.supervisor._support import python_server
PYTHON_CHILD: Final = """
import json, os, pathlib, socket, time
time.sleep(0.5)
pathlib.Path(os.environ["MARKERS"], "listening").touch()
private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"]
listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"]))
listener.listen()
time.sleep(10)
"""
RUST_CHILD: Final = """
import os, pathlib
markers = pathlib.Path(os.environ["MARKERS"])
(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists()))
"""
def _fake_gateway(scripts: Path) -> Platform:
if sys.platform != "win32":
gateway: Final = scripts / "litellm-gateway"
gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}")
gateway.chmod(0o755)
return current_platform()
script: Final = scripts / "rust_child.py"
script.write_text(RUST_CHILD)
(scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n')
return WindowsPlatform(gateway_name="litellm-gateway.cmd")
def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("MARKERS", str(tmp_path))
command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD)
supervise(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path))
assert (tmp_path / "rust_saw_python_listening").read_text() == "True"
TERM: Final = signal.SIGTERM
@ -190,15 +151,6 @@ async def test_a_cancelled_supervisor_still_waits_for_the_child_it_stops() -> No
assert (while_child_runs, once_child_exits) == (((child,), 1, False), True)
async def test_exiting_children_are_reported_as_events() -> None:
async with open_children(
current_platform(), spawn_subprocess, _children((sys.executable, "-c", "exit(7)")), 1
) as children:
assert await children.spawn(Child.RUST) is None
with anyio.fail_after(5):
assert await children.next_event() == ChildExited(Child.RUST, 7)
async def test_a_child_that_cannot_start_is_returned_as_a_failure(tmp_path: Path) -> None:
async with open_children(
current_platform(), spawn_subprocess, _children((str(tmp_path / "missing"),)), 1
@ -223,13 +175,3 @@ def _dropped(_signum: signal.Signals) -> None:
)
def test_a_supervisor_the_kernel_will_not_signal_still_exits_as_python_alone_would(status: int, exits: int) -> None:
assert exit_like(status, STOPS, die_by=_dropped) == exits
@POSIX_ONLY
@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGKILL])
def test_the_supervisor_dies_by_the_signal_that_killed_python(signum: signal.Signals) -> None:
script: Final = (
"from litellm.proxy.rust_forward.supervisor.runtime import exit_like\n"
f"raise SystemExit(exit_like({-signum}, frozenset()))"
)
assert subprocess.run((sys.executable, "-I", "-c", script), check=False, timeout=60).returncode == -signum

View file

@ -151,37 +151,6 @@ def test_main_rejects_missing_multiproc_dir(monkeypatch: pytest.MonkeyPatch):
run.assert_not_called()
def test_start_metrics_server_process_returns_only_once_child_serves(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))
_write_worker_sample(pid=3001, value=4)
port: Final = _free_port()
with patch("atexit.register") as register:
process: Final = start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path))
try:
register.assert_called_once_with(process.terminate)
assert process.poll() is None
startup_metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=5.0)
assert startup_metrics.status_code == 200
assert startup_metrics.headers[PID_HEADER] == str(process.pid)
metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=10.0)
assert 'litellm_requests_metric_total{model="gpt-5"} 4.0' in metrics.text
finally:
process.kill()
process.wait(timeout=10)
def test_start_metrics_server_process_fails_when_port_is_taken(tmp_path: Path):
with socket.socket() as occupied:
occupied.bind(("127.0.0.1", 0))
occupied.listen()
port: Final = occupied.getsockname()[1]
with (
patch("atexit.register"),
pytest.raises(MetricsServerStartupError),
):
start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path))
class _ImpostorMetrics(BaseHTTPRequestHandler):
"""An unrelated service already on the port that answers /metrics with 200 and plausible metrics."""
@ -196,22 +165,6 @@ class _ImpostorMetrics(BaseHTTPRequestHandler):
return
def test_start_metrics_server_process_rejects_metrics_from_another_service_on_the_port(tmp_path: Path):
with ThreadingHTTPServer(("127.0.0.1", 0), _ImpostorMetrics) as impostor:
threading.Thread(target=impostor.serve_forever, daemon=True).start()
port: Final = impostor.server_address[1]
impostor_response: Final = httpx.get(f"http://127.0.0.1:{port}/metrics")
assert impostor_response.status_code == 200
assert "# HELP impostor_metric" in impostor_response.text
assert PID_HEADER not in impostor_response.headers
with (
patch("atexit.register"),
pytest.raises(MetricsServerStartupError),
):
start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path))
impostor.shutdown()
def test_metrics_server_process_serves_and_exits_with_parent(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))
_write_worker_sample(pid=2001, value=7)

View file

@ -3201,33 +3201,6 @@ async def test_user_config_router_enforces_its_limit_across_proxy_requests(
)
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_limit(
fake_redis_port: int, arrives_with_logging_obj: bool
) -> None:
tpm_for_three_requests: Final = await _tpm_for_three_requests()
def request() -> Coroutine[None, None, str]:
return _unsettled_user_config_request_outcome(
fake_redis_port,
{"routing_strategy": "usage-based-routing-v2"},
{"tpm": tpm_for_three_requests},
_proxy_logging_obj() if arrives_with_logging_obj else None,
)
async def settled(*requests: Coroutine[None, None, str]) -> list[str]:
outcomes: Final = await asyncio.gather(*requests)
await _settle_router_accounting()
return outcomes
await _wait_out_a_rate_limit_minute_rollover()
concurrent_outcomes: Final = await settled(request(), request())
later_outcomes: Final = [*(await settled(request())), *(await settled(request()))]
assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"]
@pytest.mark.asyncio
async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None:
def request() -> Coroutine[None, None, str]:

View file

@ -1,8 +1,6 @@
import asyncio
import importlib
import subprocess
import sys
import textwrap
import types
from types import MappingProxyType
from collections.abc import Mapping, Sequence
@ -1032,52 +1030,6 @@ async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monk
assert captured["metadata"]["tags"] == ["team-a", "prod"]
def test_completion_with_function_tools_works_without_fastapi_installed():
script = textwrap.dedent(
"""
import sys
class _FastapiBlocker:
def find_spec(self, fullname, path=None, target=None):
if fullname == "fastapi" or fullname.startswith("fastapi."):
raise ModuleNotFoundError("No module named 'fastapi'")
return None
sys.meta_path.insert(0, _FastapiBlocker())
import litellm
response = litellm.completion(
model="openai/gpt-5.5",
messages=[{"role": "user", "content": "What is the weather in SF?"}],
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
],
mock_response="sunny",
)
assert response.choices[0].message.content == "sunny"
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
timeout=120,
)
assert result.returncode == 0, result.stderr
def test_extract_tool_call_details_reads_anthropic_tool_use_input():
"""
Regression test (LIT-4517): an Anthropic tool_use block carries its arguments

View file

@ -1,8 +1,4 @@
from __future__ import annotations
import os
import subprocess
import sys
from collections.abc import Generator
from concurrent.futures import ThreadPoolExecutor
from typing import Final
@ -132,21 +128,3 @@ def test_process_override_and_reset_apply_to_existing_threads() -> None:
assert executor.submit(configuration.rust_enabled).result() is True
configuration.reset_rust_configuration()
assert executor.submit(configuration.rust_enabled).result() is False
@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False")))
def test_environment_controls_startup(value: str, expected: str) -> None:
environment: Final = {**os.environ, "LITELLM_RUST": value}
result: Final = subprocess.run(
(
sys.executable,
"-c",
"from litellm.rust_bridge.configuration import rust_enabled; print(rust_enabled())",
),
check=True,
capture_output=True,
text=True,
env=environment,
)
assert result.stdout.strip() == expected

View file

@ -1,7 +1,4 @@
import importlib
import os
import subprocess
import sys
from typing import Final
import pytest
@ -12,22 +9,6 @@ from litellm._lazy_imports_registry import LITELLM_ATTRIBUTES, UTILS_ATTRIBUTES
from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter
def test_import_litellm_does_not_load_fastapi_or_bpe_table():
result: Final = subprocess.run(
[
sys.executable,
"-c",
"import sys, litellm; print(','.join(m for m in ('fastapi','starlette','litellm.proxy.proxy_cli','litellm.litellm_core_utils.default_encoding','openai.types.realtime') if m in sys.modules))",
],
check=True,
capture_output=True,
text=True,
env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.stdout.strip() == ""
def _declared(location: tuple[str, str]) -> object:
module_path, attr_name = location
return getattr(importlib.import_module(module_path, package="litellm"), attr_name)

View file

@ -1,8 +1,4 @@
import json
import os
import subprocess
import sys
from collections.abc import Mapping
from pathlib import Path
from typing import Final
@ -16,373 +12,8 @@ if sys.version_info >= (3, 11):
else:
import tomli as tomllib
_HELP_ARGV: Final = (
("--help",),
("logout", "--help"),
("login", "--help"),
("models", "list", "--help"),
("claude", "--help"),
("opencode", "--help"),
("codex", "--help"),
("pi", "--help"),
("keys", "list", "--help"),
)
_COMMAND_ARGV: Final = (("logout",), ("whoami",))
_RUN_EACH_ARGV: Final = """
import contextlib, io, json, sys
from {module} import cli
runs = {{}}
for argv in json.loads(sys.argv[1]):
sys.argv = ["lite", *argv]
output = io.StringIO()
with contextlib.redirect_stdout(output):
try:
cli()
except SystemExit as stopped:
code = stopped.code
else:
code = 0
runs[" ".join(argv)] = {{"exit": code, "output": output.getvalue()}}
print(json.dumps({{"runs": runs, "sdk_initialized": "completion" in vars(sys.modules["litellm"])}}))
"""
_USE_SDK_AFTER_HELP: Final = """
import builtins, importlib, json, sys
interpreter_imports = (builtins.__import__, importlib.import_module)
from {module} import cli
sys.argv = ["lite", "--help"]
try:
cli()
except SystemExit:
pass
before = "completion" in vars(sys.modules["litellm"])
{first_use}
from litellm.router import Router
import litellm
import litellm.responses.main
from litellm import acompletion
print(json.dumps({{
"before": before,
"router": Router.__name__,
"completion": callable(litellm.completion),
"acompletion": acompletion is litellm.acompletion,
"responses_is_function": callable(litellm.responses) and not isinstance(litellm.responses, type(litellm)),
"namespace": sorted(vars(litellm)),
"module_type": type(litellm).__name__,
"imports_restored": (builtins.__import__, importlib.import_module) == interpreter_imports,
}}))
"""
_FIRST_USE_FROM_THREADS: Final = """
import json, sys, threading
from litellm_cli.main import cli
sys.argv = ["lite", "--help"]
try:
cli()
except SystemExit:
pass
import litellm
start = threading.Barrier(4)
seen = []
def use_sdk():
start.wait()
try:
seen.append(callable(litellm.completion))
except AttributeError as error:
seen.append(str(error))
threads = [threading.Thread(target=use_sdk) for _ in range(4)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
print(json.dumps({"seen": seen}))
"""
_STAR_IMPORT_AFTER_HELP: Final = """
import json, sys
from litellm_cli.main import cli
sys.argv = ["lite", "--help"]
try:
cli()
except SystemExit:
pass
before = "completion" in vars(sys.modules["litellm"])
from litellm import *
print(json.dumps({"before": before, "completion": callable(globals().get("completion")), "acompletion": callable(globals().get("acompletion"))}))
"""
_SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING: Final = {
"sdk_sync.py": """
import threading
initializing = threading.Event()
submodule_loading = threading.Event()
""",
"litellm/__init__.py": """
import sdk_sync
sdk_sync.initializing.set()
sdk_sync.submodule_loading.wait(2)
from litellm import submodule
ready = True
""",
"litellm/submodule.py": """
import litellm
import sdk_sync
sdk_sync.submodule_loading.set()
ready_when_loaded = getattr(litellm, "ready", False)
""",
"litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n",
"litellm/proxy/__init__.py": "",
"litellm/proxy/client/__init__.py": "",
"litellm/proxy/client/cli/__init__.py": "",
"litellm/proxy/client/cli/main.py": """
import json, threading
import litellm
import sdk_sync
def cli():
def initialize_sdk():
litellm.ready
def import_submodule_once_initialization_starts():
sdk_sync.initializing.wait()
import litellm.submodule
threads = [threading.Thread(target=initialize_sdk, daemon=True),
threading.Thread(target=import_submodule_once_initialization_starts, daemon=True)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(10)
print(json.dumps({"finished": [not thread.is_alive() for thread in threads]}))
litellm_proxy_cli = cli
""",
}
_SDK_WHOSE_INIT_FAILS: Final = {
"litellm/__init__.py": 'raise ImportError("a dependency of the SDK is broken")\n',
"litellm/utils.py": "",
"litellm/_process_defaults.py": "def apply_process_defaults() -> None: ...\n",
"litellm/proxy/__init__.py": "",
"litellm/proxy/client/__init__.py": "",
"litellm/proxy/client/cli/__init__.py": "",
"litellm/proxy/client/cli/main.py": """
import json
import litellm
def error_from(use):
try:
use()
except Exception as error:
return f"{type(error).__name__}: {error}"
return None
def import_sdk_code():
import litellm.utils
def cli():
print(json.dumps([error_from(lambda: litellm.completion), error_from(lambda: litellm.completion), error_from(import_sdk_code)]))
litellm_proxy_cli = cli
""",
}
_RUN_LITE_AGAINST_SDK_AT: Final = """
import sys
sys.path.insert(0, sys.argv[1])
from litellm_cli.main import cli
cli()
"""
_PYPROJECT: Final = Path(__file__).resolve().parents[2] / "pyproject.toml"
_RUN_ENTRY_POINT_HELP: Final = """
import contextlib, io, json, sys
from {module} import {function}
imported_sdk = "litellm" in sys.modules
sys.argv = ["litellm", "--help"]
output = io.StringIO()
with contextlib.redirect_stdout(output):
try:
{function}()
except SystemExit as stopped:
code = stopped.code
else:
code = 0
print(json.dumps({{"imported_sdk": imported_sdk, "exit": code, "help": output.getvalue()}}))
"""
_WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG: Final = """
import json, logging, os
os.environ["LITELLM_WORKER_LOG_CONFIG"] = json.dumps({{"version": 1, "root": {{"level": "INFO"}}}})
from litellm_cli.main import proxy_worker_app
proxy_worker_app()
print(json.dumps({{name: logging.getLogger(name).disabled for name in {names!r}}}))
"""
_COMMAND_MODULES_LOADED_BY: Final = """
import json, sys
from litellm_cli.main import cli
sys.argv = ["lite", *sys.argv[1:]]
try:
cli()
except SystemExit:
pass
prefix = "litellm.proxy.client.cli.commands."
print(json.dumps({"modules": sorted(name[len(prefix):] for name in sys.modules if name.startswith(prefix))}))
"""
def _run(source: str, argv: tuple[str, ...], home: Path) -> dict[str, object]:
checkout: Final = str(_PYPROJECT.parent)
bootstrap: Final = (
"import pathlib,sys\n"
f"sys.path.insert(0, {checkout!r})\n"
f"sys.path.insert(0, str(pathlib.Path({checkout!r}) / '.github/packaging'))\n"
"from cli_package import editable_support\n"
f"exec(editable_support(pathlib.Path({checkout!r}) / 'litellm/cli')[0][1])\n"
)
result: Final = subprocess.run(
[sys.executable, "-I", "-c", bootstrap + source, *argv],
capture_output=True,
text=True,
timeout=120,
env={**os.environ, "HOME": str(home), "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
)
assert result.returncode == 0, result.stderr
return json.loads(result.stdout.rstrip("\n").rpartition("\n")[2])
@pytest.fixture(scope="module")
def cli_runs(tmp_path_factory: pytest.TempPathFactory) -> Mapping[str, dict[str, object]]:
argv: Final = (json.dumps(_HELP_ARGV + _COMMAND_ARGV),)
return {
"thin": _run(_RUN_EACH_ARGV.format(module="litellm_cli.main"), argv, tmp_path_factory.mktemp("thin")),
"full": _run(
_RUN_EACH_ARGV.format(module="litellm.proxy.client.cli.main"), argv, tmp_path_factory.mktemp("full")
),
}
def test_litellm_command_imports_the_proxy_only_when_it_runs(tmp_path: Path):
module, _, function = tomllib.loads(_PYPROJECT.read_text())["project"]["scripts"]["litellm"].partition(":")
run: Final = _run(_RUN_ENTRY_POINT_HELP.format(module=module, function=function), (), tmp_path)
assert run["imported_sdk"] is False
assert run["exit"] == 0
assert "--num_workers" in str(run["help"])
def test_a_worker_log_config_silences_the_sdk_loggers_but_not_the_ones_the_proxy_app_creates(tmp_path: Path):
disabled: Final = _run(
_WORKER_LOGGERS_UNDER_A_DEFAULT_LOG_CONFIG.format(names=("LiteLLM Proxy", "litellm.proxy.health_check")),
(),
tmp_path,
)
assert disabled == {"LiteLLM Proxy": True, "litellm.proxy.health_check": False}
def test_lite_runs_commands_without_initializing_the_sdk(cli_runs: Mapping[str, dict[str, object]]):
assert cli_runs["thin"]["sdk_initialized"] is False
assert cli_runs["full"]["sdk_initialized"] is True
@pytest.mark.parametrize("argv", _HELP_ARGV + _COMMAND_ARGV, ids=" ".join)
def test_lite_prints_what_the_full_sdk_cli_prints(argv: tuple[str, ...], cli_runs: Mapping[str, dict[str, object]]):
thin: Final = cli_runs["thin"]["runs"]
full: Final = cli_runs["full"]["runs"]
assert isinstance(thin, dict) and isinstance(full, dict)
assert thin[" ".join(argv)] == full[" ".join(argv)]
assert thin[" ".join(argv)]["output"]
@pytest.mark.parametrize(
("argv", "loaded", "not_loaded"),
[
(("logout", "--help"), {"auth"}, {"agents", "autoroute.commands", "models", "keys"}),
(("models", "list", "--help"), {"models"}, {"agents", "autoroute.commands", "keys"}),
(("claude", "--help"), {"agents"}, {"autoroute.commands", "models", "keys"}),
],
)
def test_lite_imports_only_the_command_it_runs(
argv: tuple[str, ...], loaded: set[str], not_loaded: set[str], tmp_path: Path
):
modules: Final = _run(_COMMAND_MODULES_LOADED_BY, argv, tmp_path)["modules"]
assert isinstance(modules, list)
assert loaded <= set(modules)
assert not not_loaded & set(modules)
@pytest.fixture(scope="module")
def sdk_used_after_full_cli(tmp_path_factory: pytest.TempPathFactory) -> dict[str, object]:
return _run(
_USE_SDK_AFTER_HELP.format(module="litellm.proxy.client.cli.main", first_use="pass"),
(),
tmp_path_factory.mktemp("full"),
)
@pytest.mark.parametrize(
"first_use",
[
"import litellm; litellm.completion",
"import litellm; litellm.types.utils.ModelResponse",
"from litellm import acompletion",
"from litellm.router import Router",
"from litellm.types import utils",
"import litellm.responses.main",
'importlib.import_module("litellm.router")',
'importlib.import_module(".utils", "litellm")',
'exec("from .types import utils", {"__name__": "litellm.lazy_import", "__package__": "litellm"})',
],
)
def test_sdk_initializes_in_place_on_first_use_after_lite_imports(
first_use: str, sdk_used_after_full_cli: dict[str, object], tmp_path: Path
):
thin: Final = _run(_USE_SDK_AFTER_HELP.format(module="litellm_cli.main", first_use=first_use), (), tmp_path)
assert thin["before"] is False
assert sdk_used_after_full_cli["before"] is True
assert {**thin, "before": True} == sdk_used_after_full_cli
assert (thin["router"], thin["completion"], thin["acompletion"], thin["responses_is_function"]) == (
"Router",
True,
True,
True,
)
def test_star_import_after_lite_imports_gets_the_whole_sdk(tmp_path: Path):
assert _run(_STAR_IMPORT_AFTER_HELP, (), tmp_path) == {"before": False, "completion": True, "acompletion": True}
def test_threads_that_use_the_sdk_at_once_all_wait_for_one_initialization(tmp_path: Path):
used: Final = _run(_FIRST_USE_FROM_THREADS, (), tmp_path)
assert used == {"seen": [True, True, True, True]}
def _write_sdk(root: Path, files: Mapping[str, str]) -> Path:
for relative_path, source in files.items():
(root / relative_path).parent.mkdir(parents=True, exist_ok=True)
(root / relative_path).write_text(source)
return root
def test_importing_sdk_code_while_another_thread_initializes_the_sdk_does_not_hang(tmp_path: Path):
sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_SUBMODULE_READS_IT_WHILE_LOADING)
assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == {"finished": [True, True]}
def test_every_use_after_a_failed_sdk_init_raises_the_init_error(tmp_path: Path):
sdk_root: Final = _write_sdk(tmp_path / "sdk", _SDK_WHOSE_INIT_FAILS)
assert _run(_RUN_LITE_AGAINST_SDK_AT, (str(sdk_root),), tmp_path) == ["ImportError: a dependency of the SDK is broken"] * 3
def test_lite_keeps_an_sdk_that_was_already_imported(
monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]

View file

@ -3,15 +3,12 @@ Router behavior the router-code differential fuzz run found changed once the Rus
only core: count settings below zero, an allowed_fails_policy write that is no policy, a model
response object as a deployment's mock_response, and endpoints read or patched on the Router class.
"""
import asyncio
import gc
import json
import socket
import threading
import weakref
from collections.abc import Iterator, Mapping
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from unittest import mock
@ -245,95 +242,6 @@ def upstream() -> Iterator[str]:
thread.join()
def _served(name: str, deployment_id: str, model: str, api_base: str) -> dict[str, object]:
return {
"model_name": name,
"litellm_params": {"model": model, "api_key": "k", "api_base": api_base},
"model_info": {"id": deployment_id},
}
def _request_logging(call_type: str, *, stream: bool) -> object:
from litellm.litellm_core_utils.litellm_logging import Logging
return Logging(
model="g",
messages=MESSAGES,
stream=stream,
call_type=call_type,
start_time=datetime.now(), # noqa: DTZ005 # litellm's own naive start time
litellm_call_id="call-1",
function_id="f",
)
async def _recorded() -> None:
"""Let the background recordings of a stream's outcome run."""
for _ in range(50):
await asyncio.sleep(0.02)
@pytest.mark.asyncio
async def test_a_responses_stream_cut_after_it_opened_counts_one_fail_with_a_shared_logging_object(
upstream: str,
) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/cut/v1")], num_retries=0, disable_cooldowns=True
)
stream: Final = await router.aresponses(
model="g", input="hi", stream=True, litellm_logging_obj=_request_logging("aresponses", stream=True)
)
with pytest.raises(Exception): # noqa: B017, PT011 # the transport's read error, whatever class maps it
async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True
pass
await _recorded()
assert router.health().store.get_cache(key="d1:fails") == 1
@pytest.mark.asyncio
async def test_a_blocked_response_raised_for_a_fallback_counts_no_success_with_a_shared_logging_object(
upstream: str,
) -> None:
router: Final = Router(
model_list=[
_served("a", "a1", "openai/gpt-4o", f"{upstream}/filtered/v1"),
_served("safe", "s1", "openai/gpt-4o", f"{upstream}/ok/v1"),
],
content_policy_fallbacks=[{"a": ["safe"]}],
num_retries=0,
disable_cooldowns=True,
)
response: Final = await router.acompletion(
model="a", messages=MESSAGES, litellm_logging_obj=_request_logging("acompletion", stream=False)
)
await _recorded()
assert response.choices[0].message.content == "ok"
assert router.health().store.get_cache(key="a1:successes") is None
assert router.health().store.get_cache(key="s1:successes") == 1
@pytest.mark.asyncio
async def test_a_text_completion_stream_that_fails_before_its_first_chunk_counts_nothing(upstream: str) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-3.5-turbo-instruct", f"{upstream}/down/v1")],
num_retries=0,
disable_cooldowns=True,
)
stream: Final = await router.atext_completion(model="g", prompt="hi", stream=True)
with pytest.raises(litellm.BadGatewayError):
async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True
pass
await _recorded()
assert router.health().store.get_cache(key="d1:successes") is None
assert router.health().store.get_cache(key="d1:fails") is None
@pytest.mark.asyncio
async def test_a_text_completion_without_a_prompt_raises_type_error_before_routing() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0)
@ -389,20 +297,6 @@ async def test_a_refused_pick_by_deployment_id_names_the_deployments_group() ->
await router.async_get_available_deployment(model="p2", messages=MESSAGES, request_kwargs={})
@pytest.mark.asyncio
async def test_a_request_refused_before_the_provider_counts_against_no_deployment(upstream: str) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/ok/v1")], num_retries=0, disable_cooldowns=True
)
with pytest.raises(litellm.UnsupportedParamsError):
await router.acompletion(model="g", messages=MESSAGES, reasoning_effort="low")
await _recorded()
assert router.health().store.get_cache(key="d1:fails") is None
assert not [key for key in router.health().store.in_memory_cache.cache_dict if ":rpm" in key]
@pytest.mark.asyncio
async def test_the_async_pick_read_on_the_class_is_a_function_taking_the_router() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")])

View file

@ -1,4 +1,3 @@
import subprocess
import sys
import pytest
@ -7,19 +6,6 @@ import litellm.router_host as router_host_module
from litellm.router_host import RouterHost, router_host
def test_router_uses_the_proxy_host_without_importing_the_proxy_first():
# A fresh interpreter, so no earlier test has imported litellm.proxy.
script = (
"import sys\n"
"from litellm.router import Router\n"
"from litellm.router_host import router_host\n"
"assert 'litellm.proxy' not in sys.modules, 'importing the router imported the proxy'\n"
"print(type(router_host()).__name__)\n"
)
result = subprocess.run([sys.executable, "-P", "-c", script], capture_output=True, text=True, check=True)
assert result.stdout.strip() == "ProxyRouterHost"
@pytest.fixture(autouse=True)
def nothing_installed():
saved = router_host_module._installed[0]

View file

@ -1,6 +1,6 @@
# Windows tests
The `windows-sdk` job in `.github/workflows/test-merge-smoke.yml` is the only place these run, on a Windows runner with no proxy dependencies installed. They prove the SDK works on Windows without them, so a test here must not import uvicorn, fastapi, psutil or anything else outside the `dev` group, and must not need the `litellm-gateway.exe` binary
The `windows-sdk` job in `.github/workflows/test-merge-smoke.yml` is the only place these run, on pushes to main and manual runs but not on pull requests, on a Windows runner with no proxy dependencies installed. They prove the SDK works on Windows without them, so a test here must not import uvicorn, fastapi, psutil or anything else outside the `dev` group, and must not need the `litellm-gateway.exe` binary
A test that cannot run off Windows carries a `skipif` on `sys.platform`, so collecting this folder on another OS never fails