mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
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:
parent
ac23812b4b
commit
da55893656
68 changed files with 5143 additions and 3821 deletions
4
.github/workflows/test-merge-smoke.yml
vendored
4
.github/workflows/test-merge-smoke.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
14
.github/workflows/test-unit.yml
vendored
14
.github/workflows/test-unit.yml
vendored
|
|
@ -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
107
tests/_module_state.py
Normal 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)
|
||||
|
|
@ -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]]
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -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)
|
||||
119
tests/integration/caching/test_evicted_client_closer_real_io.py
Normal file
119
tests/integration/caching/test_evicted_client_closer_real_io.py
Normal 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()
|
||||
|
|
@ -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
|
||||
272
tests/integration/configuration/test_claude_settings_real_io.py
Normal file
272
tests/integration/configuration/test_claude_settings_real_io.py
Normal 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
|
||||
20
tests/integration/configuration/test_main_real_io.py
Normal file
20
tests/integration/configuration/test_main_real_io.py
Normal 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}"
|
||||
83
tests/integration/configuration/test_runtime_real_io.py
Normal file
83
tests/integration/configuration/test_runtime_real_io.py
Normal 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
|
||||
|
|
@ -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"}),),
|
||||
}
|
||||
)
|
||||
1527
tests/integration/mcp/test_mcp_client_real_io.py
Normal file
1527
tests/integration/mcp/test_mcp_client_real_io.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
|
|
@ -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
|
||||
41
tests/integration/sdk/test_anthropic_common_utils_real_io.py
Normal file
41
tests/integration/sdk/test_anthropic_common_utils_real_io.py
Normal 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]"]}
|
||||
]
|
||||
204
tests/integration/sdk/test_azure_common_utils_real_io.py
Normal file
204
tests/integration/sdk/test_azure_common_utils_real_io.py
Normal 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."),
|
||||
),
|
||||
}
|
||||
39
tests/integration/sdk/test_configuration_real_io.py
Normal file
39
tests/integration/sdk/test_configuration_real_io.py
Normal 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
|
||||
146
tests/integration/sdk/test_context_real_io.py
Normal file
146
tests/integration/sdk/test_context_real_io.py
Normal 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)
|
||||
50
tests/integration/sdk/test_docker_real_io.py
Normal file
50
tests/integration/sdk/test_docker_real_io.py
Normal 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()
|
||||
296
tests/integration/sdk/test_http_handler_real_io.py
Normal file
296
tests/integration/sdk/test_http_handler_real_io.py
Normal 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)
|
||||
37
tests/integration/sdk/test_jwt_signing_real_io.py
Normal file
37
tests/integration/sdk/test_jwt_signing_real_io.py
Normal 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
|
||||
23
tests/integration/sdk/test_lazy_imports_real_io.py
Normal file
23
tests/integration/sdk/test_lazy_imports_real_io.py
Normal 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() == ""
|
||||
384
tests/integration/sdk/test_litellm_cli_real_io.py
Normal file
384
tests/integration/sdk/test_litellm_cli_real_io.py
Normal 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
|
||||
323
tests/integration/sdk/test_long_decode_real_io.py
Normal file
323
tests/integration/sdk/test_long_decode_real_io.py
Normal 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)}
|
||||
|
|
@ -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]
|
||||
28
tests/integration/sdk/test_router_host_real_io.py
Normal file
28
tests/integration/sdk/test_router_host_real_io.py
Normal 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
|
||||
367
tests/integration/sdk/test_streaming_handler_real_io.py
Normal file
367
tests/integration/sdk/test_streaming_handler_real_io.py
Normal 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,
|
||||
)
|
||||
],
|
||||
)
|
||||
207
tests/integration/sdk/test_token_counter_real_io.py
Normal file
207
tests/integration/sdk/test_token_counter_real_io.py
Normal 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}
|
||||
119
tests/integration/sdk/test_tokenizer_real_io.py
Normal file
119
tests/integration/sdk/test_tokenizer_real_io.py
Normal 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()
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]"]}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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")])
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue