mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-03 12:24:08 +00:00
Sameer KankuteandYuneng Jiang
5cea68f3ae
* fix(proxy): authorize batch files using upload target_model_names (LIT-3593)
After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)
Restores the reverse-lookup for the JSONL body.model fallback path so that
legacy/pre-target_model_names managed files still map stripped provider IDs
back to proxy aliases before auth. Also cleans up redundant `or None`.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)"
This reverts commit 30d2e96f77ef521ccaaf2193fe554980380eb669.
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
(cherry picked from commit 2cd7e87485)
421 lines
14 KiB
Python
421 lines
14 KiB
Python
"""
|
|
VERIA-39 regression tests:
|
|
|
|
- The batch input-file token counter must measure embeddings (`input`)
|
|
and text-completion (`prompt`) payloads, not only chat (`messages`).
|
|
- The batch rate-limiter pre-call hook must reject batch files that name
|
|
models the caller is not authorized to use.
|
|
"""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Token counter — covers all three batch payload shapes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_token_counter_counts_chat_messages():
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{
|
|
"body": {
|
|
"model": "gpt-4o-mini",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}
|
|
}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens > 0
|
|
|
|
|
|
def test_token_counter_counts_text_completion_prompt():
|
|
"""Pre-fix this returned 0 tokens (the function only inspected
|
|
`messages`), letting `prompt`-style batches slip past TPM limits."""
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens > 0
|
|
|
|
|
|
def test_token_counter_counts_embedding_input_string():
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{"body": {"model": "text-embedding-3-small", "input": "hello world"}}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens > 0
|
|
|
|
|
|
def test_token_counter_counts_embedding_input_list():
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{
|
|
"body": {
|
|
"model": "text-embedding-3-small",
|
|
"input": ["hello", "world"],
|
|
}
|
|
}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens > 0
|
|
|
|
|
|
def test_token_counter_counts_text_completion_prompt_list():
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{
|
|
"body": {
|
|
"model": "gpt-3.5-turbo-instruct",
|
|
"prompt": ["alpha", "beta"],
|
|
}
|
|
}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens > 0
|
|
|
|
|
|
def test_token_counter_counts_pre_tokenized_prompt_int_list():
|
|
"""OpenAI's text-completion API accepts a single pre-tokenized prompt as
|
|
a list of ints. Each int is one token; pre-fix this shape was silently
|
|
counted as zero, leaving a TPM bypass."""
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{
|
|
"body": {
|
|
"model": "gpt-3.5-turbo-instruct",
|
|
"prompt": [1, 2, 3, 4, 5],
|
|
}
|
|
}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens == 5
|
|
|
|
|
|
def test_token_counter_counts_pre_tokenized_prompt_list_of_int_lists():
|
|
"""Multiple pre-tokenized prompts (`list[list[int]]`) — the most
|
|
important bypass shape. A 1000-token batch must report 1000 tokens,
|
|
not zero."""
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{
|
|
"body": {
|
|
"model": "gpt-3.5-turbo-instruct",
|
|
"prompt": [[1] * 250, [2] * 250, [3] * 500],
|
|
}
|
|
}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens == 1000
|
|
|
|
|
|
def test_token_counter_counts_pre_tokenized_input_for_embeddings():
|
|
"""Same shape applies to embeddings (`input`)."""
|
|
from litellm.batches.batch_utils import _get_batch_job_input_file_usage
|
|
|
|
usage = _get_batch_job_input_file_usage(
|
|
file_content_dictionary=[
|
|
{
|
|
"body": {
|
|
"model": "text-embedding-3-small",
|
|
"input": [[1, 2, 3], [4, 5, 6]],
|
|
}
|
|
}
|
|
]
|
|
)
|
|
assert usage.prompt_tokens == 6
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model extractor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_model_extractor_returns_distinct_models():
|
|
from litellm.batches.batch_utils import _get_models_from_batch_input_file_content
|
|
|
|
models = _get_models_from_batch_input_file_content(
|
|
[
|
|
{"body": {"model": "gpt-4o", "messages": []}},
|
|
{"body": {"model": "gpt-4o", "messages": []}}, # duplicate
|
|
{"body": {"model": "gpt-4o-mini", "messages": []}},
|
|
{"body": {}}, # missing model
|
|
]
|
|
)
|
|
assert models == ["gpt-4o", "gpt-4o-mini"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pre-call hook model validation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_rejects_unauthorized_model_in_batch_file():
|
|
"""Pre-fix the hook only validated the outer `model` parameter and
|
|
forwarded the file as-is. With this fix, a model named inside the
|
|
JSONL that the caller cannot use must trigger a 403."""
|
|
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
|
|
|
rate_limiter = _PROXY_BatchRateLimiter(
|
|
internal_usage_cache=MagicMock(),
|
|
parallel_request_limiter=MagicMock(),
|
|
)
|
|
|
|
# Simulated decoded batch file: caller is restricted to gpt-3.5
|
|
# but the JSONL points at gpt-4o.
|
|
file_dict = [
|
|
{"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "x"}]}}
|
|
]
|
|
|
|
user = UserAPIKeyAuth(
|
|
api_key="sk-restricted",
|
|
user_id="alice",
|
|
models=["gpt-3.5-turbo"],
|
|
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
|
)
|
|
|
|
# `can_key_call_model` raises a ProxyException for non-allowed models.
|
|
async def _raise_unauthorized(**kwargs):
|
|
raise Exception(
|
|
f"Key not allowed to access model. This key only has access to models={kwargs['valid_token'].models}"
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.auth.auth_checks.can_key_call_model",
|
|
new=AsyncMock(side_effect=_raise_unauthorized),
|
|
),
|
|
patch("litellm.proxy.proxy_server.llm_router", None),
|
|
):
|
|
with pytest.raises(HTTPException) as exc:
|
|
await rate_limiter._enforce_batch_file_model_access(
|
|
user_api_key_dict=user,
|
|
file_content_as_dict=file_dict,
|
|
)
|
|
|
|
assert exc.value.status_code == 403
|
|
assert "gpt-4o" in str(exc.value.detail)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_allows_authorized_model_in_batch_file():
|
|
"""If every model in the JSONL is on the caller's allowlist, the hook
|
|
must not raise."""
|
|
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
|
|
|
rate_limiter = _PROXY_BatchRateLimiter(
|
|
internal_usage_cache=MagicMock(),
|
|
parallel_request_limiter=MagicMock(),
|
|
)
|
|
|
|
file_dict = [
|
|
{
|
|
"body": {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "x"}],
|
|
}
|
|
}
|
|
]
|
|
|
|
user = UserAPIKeyAuth(
|
|
api_key="sk-ok",
|
|
user_id="alice",
|
|
models=["gpt-3.5-turbo"],
|
|
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.auth.auth_checks.can_key_call_model",
|
|
new=AsyncMock(return_value=True),
|
|
),
|
|
patch("litellm.proxy.proxy_server.llm_router", None),
|
|
):
|
|
# Should not raise
|
|
await rate_limiter._enforce_batch_file_model_access(
|
|
user_api_key_dict=user,
|
|
file_content_as_dict=file_dict,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias():
|
|
"""After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5).
|
|
Auth must check target_model_names from the unified file id, not reverse-map
|
|
the stripped id."""
|
|
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
|
|
|
rate_limiter = _PROXY_BatchRateLimiter(
|
|
internal_usage_cache=MagicMock(),
|
|
parallel_request_limiter=MagicMock(),
|
|
)
|
|
proxy_alias = "openai/openai/gpt-5.5-batch"
|
|
file_dict = [
|
|
{"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}}
|
|
]
|
|
user = UserAPIKeyAuth(
|
|
api_key="sk-ok",
|
|
user_id="alice",
|
|
models=[proxy_alias],
|
|
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
|
)
|
|
mock_router = MagicMock()
|
|
mock_router.model_list = []
|
|
can_key_call_model = AsyncMock(return_value=True)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.auth.auth_checks.can_key_call_model",
|
|
new=can_key_call_model,
|
|
),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
await rate_limiter._enforce_batch_file_model_access(
|
|
user_api_key_dict=user,
|
|
file_content_as_dict=file_dict,
|
|
target_model_names=[proxy_alias],
|
|
)
|
|
|
|
can_key_call_model.assert_awaited_once()
|
|
assert can_key_call_model.await_args.kwargs["model"] == proxy_alias
|
|
mock_router.resolve_model_name_from_model_id.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"model_list_order",
|
|
[
|
|
[
|
|
"openai/openai/gpt-5.5",
|
|
"openai/openai/gpt-5.5-batch",
|
|
"us/azure/openai/gpt-5.5",
|
|
],
|
|
[
|
|
"us/azure/openai/gpt-5.5",
|
|
"openai/openai/gpt-5.5",
|
|
"openai/openai/gpt-5.5-batch",
|
|
],
|
|
[
|
|
"openai/openai/gpt-5.5-batch",
|
|
"us/azure/openai/gpt-5.5",
|
|
"openai/openai/gpt-5.5",
|
|
],
|
|
],
|
|
)
|
|
async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup(
|
|
model_list_order,
|
|
):
|
|
"""LIT-3593: three deployments strip to gpt-5.5; auth must use the upload
|
|
target alias from target_model_names, not first-match reverse lookup."""
|
|
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
|
|
|
rate_limiter = _PROXY_BatchRateLimiter(
|
|
internal_usage_cache=MagicMock(),
|
|
parallel_request_limiter=MagicMock(),
|
|
)
|
|
batch_alias = "openai/openai/gpt-5.5-batch"
|
|
deployment_templates = {
|
|
"openai/openai/gpt-5.5": {
|
|
"model_name": "openai/openai/gpt-5.5",
|
|
"litellm_params": {"model": "openai/gpt-5.5"},
|
|
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
|
|
},
|
|
"openai/openai/gpt-5.5-batch": {
|
|
"model_name": "openai/openai/gpt-5.5-batch",
|
|
"litellm_params": {"model": "openai/gpt-5.5"},
|
|
"model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"},
|
|
},
|
|
"us/azure/openai/gpt-5.5": {
|
|
"model_name": "us/azure/openai/gpt-5.5",
|
|
"litellm_params": {"model": "azure/gpt-5.5"},
|
|
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
|
|
},
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_router.model_list = [deployment_templates[name] for name in model_list_order]
|
|
|
|
def _resolve(model_id):
|
|
for deployment in mock_router.model_list:
|
|
actual_model = deployment.get("litellm_params", {}).get("model")
|
|
if actual_model == model_id or (
|
|
actual_model and actual_model.endswith(f"/{model_id}")
|
|
):
|
|
return deployment.get("model_name")
|
|
return None
|
|
|
|
mock_router.resolve_model_name_from_model_id.side_effect = _resolve
|
|
|
|
file_dict = [
|
|
{"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}}
|
|
]
|
|
user = UserAPIKeyAuth(
|
|
api_key="sk-ok",
|
|
user_id="alice",
|
|
models=[batch_alias],
|
|
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
|
)
|
|
can_key_call_model = AsyncMock(return_value=True)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.auth.auth_checks.can_key_call_model",
|
|
new=can_key_call_model,
|
|
),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
await rate_limiter._enforce_batch_file_model_access(
|
|
user_api_key_dict=user,
|
|
file_content_as_dict=file_dict,
|
|
target_model_names=[batch_alias],
|
|
)
|
|
|
|
can_key_call_model.assert_awaited_once()
|
|
assert can_key_call_model.await_args.kwargs["model"] == batch_alias
|
|
mock_router.resolve_model_name_from_model_id.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_skips_check_when_no_models_present():
|
|
"""Files without any `body.model` (corrupt or empty) must not 500;
|
|
the rate limiter logs a warning elsewhere and proceeds."""
|
|
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
|
|
|
rate_limiter = _PROXY_BatchRateLimiter(
|
|
internal_usage_cache=MagicMock(),
|
|
parallel_request_limiter=MagicMock(),
|
|
)
|
|
user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice")
|
|
|
|
# Should not raise even though `can_key_call_model` is the default
|
|
# (would fail). The early-return on empty models keeps the call out
|
|
# entirely.
|
|
await rate_limiter._enforce_batch_file_model_access(
|
|
user_api_key_dict=user,
|
|
file_content_as_dict=[],
|
|
)
|
|
await rate_limiter._enforce_batch_file_model_access(
|
|
user_api_key_dict=user,
|
|
file_content_as_dict=[{"body": {}}],
|
|
)
|