mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
24584d3d3d
* test(proxy): move proxy_server, _experimental and db tests into tests/unit/proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): keep tuple identity in proxy state restore and fix misc target paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
415 lines
16 KiB
Python
415 lines
16 KiB
Python
"""Behavior pins for ``proxy_server.py`` llm-utils routes.
|
|
|
|
Pins (PR2):
|
|
- POST /utils/token_counter
|
|
- GET /utils/supported_openai_params
|
|
- GET /utils/model_info
|
|
- POST /utils/transform_request
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
|
|
import litellm
|
|
from litellm.litellm_core_utils import get_llm_provider_logic
|
|
from litellm.proxy import proxy_server
|
|
|
|
from .conftest import normalize # type: ignore[import-not-found]
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# POST /utils/token_counter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def patched_token_counter(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
|
|
monkeypatch.setattr(
|
|
litellm.utils,
|
|
"_select_tokenizer",
|
|
lambda model, custom_tokenizer=None: {
|
|
"type": "openai_tokenizer",
|
|
"tokenizer": None,
|
|
},
|
|
)
|
|
monkeypatch.setattr(litellm, "token_counter", lambda **kwargs: 7)
|
|
yield
|
|
|
|
|
|
def test_token_counter_happy_path(client, auth_as, patched_token_counter):
|
|
"""Pins ``POST /utils/token_counter``."""
|
|
payload = {"model": "gpt-4", "prompt": "Hi there"}
|
|
with auth_as():
|
|
response = client.post("/utils/token_counter", json=payload)
|
|
assert response.status_code == 200
|
|
assert normalize(response.json()) == {
|
|
"total_tokens": 7,
|
|
"request_model": "gpt-4",
|
|
"model_used": "gpt-4",
|
|
"tokenizer_type": "openai_tokenizer",
|
|
"original_response": None,
|
|
"error": False,
|
|
"error_message": None,
|
|
"status_code": None,
|
|
}
|
|
|
|
|
|
def test_token_counter_counts_off_the_event_loop(client, auth_as, patched_token_counter, monkeypatch):
|
|
"""
|
|
A large prompt must not stall the proxy: the count runs in a worker thread, where there
|
|
is no running event loop, rather than on the loop serving other requests.
|
|
"""
|
|
counted_off_loop = []
|
|
|
|
def recording_counter(**kwargs):
|
|
try:
|
|
asyncio.get_running_loop()
|
|
counted_off_loop.append(False)
|
|
except RuntimeError:
|
|
counted_off_loop.append(True)
|
|
return 7
|
|
|
|
monkeypatch.setattr(litellm, "token_counter", recording_counter)
|
|
|
|
with auth_as():
|
|
response = client.post("/utils/token_counter", json={"model": "gpt-4", "prompt": "Hi there"})
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["total_tokens"] == 7
|
|
assert counted_off_loop == [True]
|
|
|
|
|
|
def test_token_counter_missing_input_returns_400(
|
|
client, auth_as, patched_token_counter
|
|
):
|
|
"""Pins ``POST /utils/token_counter`` (error: missing input)."""
|
|
with auth_as():
|
|
response = client.post("/utils/token_counter", json={"model": "gpt-4"})
|
|
assert response.status_code == 400
|
|
assert "prompt or messages or contents" in response.text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GET /utils/supported_openai_params
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def patched_supported_params(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"get_llm_provider",
|
|
lambda model: (model, "openai", None, None),
|
|
)
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"get_supported_openai_params",
|
|
lambda model, custom_llm_provider=None: ["max_tokens", "temperature", "top_p"],
|
|
)
|
|
yield
|
|
|
|
|
|
def test_supported_openai_params_happy_path(client, auth_as, patched_supported_params):
|
|
"""Pins ``GET /utils/supported_openai_params``."""
|
|
with auth_as():
|
|
response = client.get(
|
|
"/utils/supported_openai_params", params={"model": "gpt-4"}
|
|
)
|
|
assert response.status_code == 200
|
|
assert normalize(response.json()) == {
|
|
"supported_openai_params": ["max_tokens", "temperature", "top_p"],
|
|
}
|
|
|
|
|
|
def test_supported_openai_params_resolves_router_alias(client, auth_as, monkeypatch):
|
|
"""A router alias absent from the cost map resolves through the deployment's underlying model."""
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "claude-opus-4-6-cached",
|
|
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"},
|
|
}
|
|
]
|
|
)
|
|
monkeypatch.setattr(proxy_server, "llm_router", router)
|
|
|
|
with auth_as():
|
|
response = client.get("/utils/supported_openai_params", params={"model": "claude-opus-4-6-cached"})
|
|
|
|
assert response.status_code == 200
|
|
expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic")
|
|
assert response.json() == {"supported_openai_params": expected}
|
|
assert "max_tokens" in response.json()["supported_openai_params"]
|
|
|
|
|
|
def test_supported_openai_params_declared_prefix_alias_resolves_through_router(client, auth_as, monkeypatch):
|
|
"""Regression: an alias whose name starts with an authenticating provider's prefix skipped
|
|
router resolution and answered with that provider's params instead of the deployment's."""
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "github_copilot/gpt-4o",
|
|
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"},
|
|
}
|
|
]
|
|
)
|
|
monkeypatch.setattr(proxy_server, "llm_router", router)
|
|
|
|
with auth_as():
|
|
response = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"})
|
|
|
|
assert response.status_code == 200
|
|
expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic")
|
|
assert response.json() == {"supported_openai_params": expected}
|
|
|
|
|
|
def test_supported_openai_params_never_runs_oauth_for_authenticating_providers(client, auth_as, monkeypatch, tmp_path):
|
|
"""Regression: github_copilot/chatgpt names answer from their declaration; resolving them
|
|
through ``get_llm_provider`` would run the provider's OAuth device flow and block the event loop."""
|
|
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
|
|
(tmp_path / "access-token").write_text("fake-access-token")
|
|
(tmp_path / "api-key.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"token": "fake-api-key",
|
|
"expires_at": 4102444800,
|
|
"endpoints": {"api": "https://api.githubcopilot.com"},
|
|
}
|
|
)
|
|
)
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "copilot-alias",
|
|
"litellm_params": {"model": "github_copilot/gpt-4o"},
|
|
},
|
|
{
|
|
"model_name": "openai/*",
|
|
"litellm_params": {"model": "openai/*"},
|
|
},
|
|
]
|
|
)
|
|
monkeypatch.setattr(proxy_server, "llm_router", router)
|
|
|
|
resolution_attempts: list[str] = []
|
|
|
|
def _oauth_tripwire(model, *args, **kwargs):
|
|
resolution_attempts.append(model)
|
|
raise AssertionError("get_llm_provider would run the OAuth device flow")
|
|
|
|
monkeypatch.setattr(litellm, "get_llm_provider", _oauth_tripwire)
|
|
monkeypatch.setattr(get_llm_provider_logic, "get_llm_provider", _oauth_tripwire)
|
|
expected = litellm.get_supported_openai_params(model="gpt-4o", custom_llm_provider="github_copilot")
|
|
|
|
with auth_as():
|
|
via_alias = client.get("/utils/supported_openai_params", params={"model": "copilot-alias"})
|
|
via_direct_name = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"})
|
|
|
|
assert via_alias.status_code == 200
|
|
assert via_alias.json() == {"supported_openai_params": expected}
|
|
assert via_direct_name.status_code == 200
|
|
assert via_direct_name.json() == {"supported_openai_params": expected}
|
|
assert resolution_attempts == []
|
|
|
|
|
|
def test_supported_openai_params_invalid_model(client, auth_as, monkeypatch):
|
|
"""Pins ``GET /utils/supported_openai_params`` (error: unknown model)."""
|
|
|
|
def _raise(model):
|
|
raise Exception("unknown")
|
|
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
monkeypatch.setattr(litellm, "get_llm_provider", _raise)
|
|
with auth_as():
|
|
response = client.get("/utils/supported_openai_params", params={"model": "??"})
|
|
assert response.status_code == 400
|
|
assert "Could not map model" in response.text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GET /utils/model_info
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def lookup_fixture_model(monkeypatch):
|
|
entry = {
|
|
"litellm_provider": "openai",
|
|
"mode": "chat",
|
|
"max_input_tokens": 1234,
|
|
"max_output_tokens": 56,
|
|
"input_cost_per_token": 1e-6,
|
|
"output_cost_per_token": 2e-6,
|
|
"supports_vision": True,
|
|
"deprecation_date": "2099-01-01",
|
|
"supports_lookup_fixture_edit": True,
|
|
}
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
monkeypatch.setitem(litellm.model_cost, "lookup-fixture-model", entry)
|
|
litellm.get_model_info.cache_clear()
|
|
litellm.utils._cached_get_model_info_helper.cache_clear()
|
|
yield entry
|
|
litellm.get_model_info.cache_clear()
|
|
litellm.utils._cached_get_model_info_helper.cache_clear()
|
|
|
|
|
|
def test_model_info_lookup_returns_full_cost_map_entry_for_unregistered_model(client, auth_as, lookup_fixture_model):
|
|
"""Every raw cost map field comes back, including ones outside ``ModelInfoBase`` that ``get_model_info`` drops."""
|
|
with auth_as():
|
|
response = client.get(
|
|
"/utils/model_info", params={"model": "lookup-fixture-model", "custom_llm_provider": "openai"}
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body = response.json()
|
|
assert body["model"] == "lookup-fixture-model"
|
|
assert body["custom_llm_provider"] == "openai"
|
|
assert body["model_info"]["key"] == "lookup-fixture-model"
|
|
assert isinstance(body["model_info"]["supported_openai_params"], list)
|
|
assert {k: body["model_info"][k] for k in lookup_fixture_model} == lookup_fixture_model
|
|
|
|
|
|
def test_model_info_lookup_unknown_model_returns_404(client, auth_as, monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
with auth_as():
|
|
response = client.get("/utils/model_info", params={"model": "no-such-model-lit-7476"})
|
|
assert response.status_code == 404, response.text
|
|
assert "is not in the model cost map" in response.text
|
|
|
|
|
|
@respx.mock
|
|
def test_model_info_lookup_returns_404_when_typed_info_has_no_cost_map_entry(
|
|
client, auth_as, monkeypatch, local_model_cost_map
|
|
):
|
|
"""``get_model_info`` synthesizes info for huggingface fallbacks absent from ``model_cost``;
|
|
with no raw entry the route must 404 rather than answer 200 with typed fields only."""
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
respx.get("https://huggingface.co/not-in-map-org/not-in-map-model/raw/main/config.json").mock(
|
|
return_value=httpx.Response(404)
|
|
)
|
|
with auth_as():
|
|
response = client.get("/utils/model_info", params={"model": "huggingface/not-in-map-org/not-in-map-model"})
|
|
assert response.status_code == 404, response.text
|
|
assert "is not in the model cost map" in response.text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# POST /utils/transform_request
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def patched_transform(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
monkeypatch.setattr(proxy_server, "is_request_body_safe", lambda **kwargs: True)
|
|
|
|
def _fake_return_raw_request(endpoint, kwargs):
|
|
return {
|
|
"raw_request_api_base": "https://api.openai.com/v1/chat/completions",
|
|
"raw_request_body": kwargs,
|
|
"raw_request_headers": {"Authorization": "Bearer redacted"},
|
|
}
|
|
|
|
monkeypatch.setattr("litellm.utils.return_raw_request", _fake_return_raw_request)
|
|
yield
|
|
|
|
|
|
def test_transform_request_happy_path(client, auth_as, patched_transform):
|
|
"""Pins ``POST /utils/transform_request``."""
|
|
payload = {"call_type": "completion", "request_body": {"model": "gpt-4"}}
|
|
with auth_as():
|
|
response = client.post("/utils/transform_request", json=payload)
|
|
assert response.status_code == 200
|
|
assert normalize(response.json()) == {
|
|
"raw_request_api_base": "https://api.openai.com/v1/chat/completions",
|
|
"raw_request_body": {"model": "gpt-4"},
|
|
"raw_request_headers": {"Authorization": "Bearer redacted"},
|
|
}
|
|
|
|
|
|
def test_transform_request_unsafe_body(client, auth_as, monkeypatch):
|
|
"""Pins ``POST /utils/transform_request`` (error: unsafe body)."""
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
|
|
def _raise(**kwargs):
|
|
raise ValueError("unsafe model")
|
|
|
|
monkeypatch.setattr(proxy_server, "is_request_body_safe", _raise)
|
|
payload = {"call_type": "completion", "request_body": {"model": "evil"}}
|
|
with auth_as():
|
|
response = client.post("/utils/transform_request", json=payload)
|
|
assert response.status_code == 400
|
|
assert "unsafe" in response.text or "error" in response.text
|
|
|
|
|
|
def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, auth_as, monkeypatch):
|
|
"""The ``litellm.token_counter`` fallback counts the request's tools and system prompt, and Anthropic ``image``/``document`` blocks, instead of 500ing."""
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
|
|
system = [{"type": "text", "text": "You are a terse assistant. Answer in one sentence."}]
|
|
tools = [
|
|
{
|
|
"name": "get_weather",
|
|
"description": "Look up the current weather for a city",
|
|
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
|
}
|
|
]
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "What is in this file?"},
|
|
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}},
|
|
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}},
|
|
],
|
|
}
|
|
]
|
|
|
|
def count(payload: dict) -> int:
|
|
with auth_as():
|
|
response = client.post("/utils/token_counter", json={"model": "claude-fable-5", **payload})
|
|
assert response.status_code == 200, response.text
|
|
return response.json()["total_tokens"]
|
|
|
|
bare = count({"messages": messages})
|
|
full = count({"messages": messages, "tools": tools, "system": system})
|
|
|
|
assert bare == litellm.token_counter(model="claude-fable-5", messages=messages)
|
|
assert full == litellm.token_counter(
|
|
model="claude-fable-5",
|
|
messages=[{"role": "system", "content": system}, *messages],
|
|
tools=tools,
|
|
)
|
|
assert full > bare
|
|
|
|
|
|
def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch):
|
|
"""Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path)."""
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
|
|
prompt = "count the tokens in this sentence please"
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"description": "Look up the current weather for a city",
|
|
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
|
},
|
|
}
|
|
]
|
|
|
|
with auth_as():
|
|
response = client.post(
|
|
"/utils/token_counter", json={"model": "claude-fable-5", "prompt": prompt, "tools": tools}
|
|
)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt)
|