mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
25109a523b
* test(proxy): move utils, agent_endpoints and endpoint tests into tests/unit/proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): package moved unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): exclude proxy-db-owned files from the misc target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop the redundant fixture docstrings in the proxy conftest 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>
594 lines
19 KiB
Python
594 lines
19 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
import litellm
|
|
from litellm import ModelResponse
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import (
|
|
create_model_info_response,
|
|
get_available_models_for_user,
|
|
hash_token,
|
|
is_known_model,
|
|
is_known_vector_store_index,
|
|
model_dump_with_preserved_fields,
|
|
validate_model_access,
|
|
)
|
|
|
|
|
|
def normalize(value):
|
|
return value
|
|
|
|
|
|
def _router_with_models(model_names):
|
|
router = MagicMock()
|
|
router.get_model_names.return_value = model_names
|
|
router.get_model_access_groups.return_value = {}
|
|
return router
|
|
|
|
|
|
def test_is_known_model_happy_path_returns_true_when_in_router():
|
|
router = _router_with_models(["gpt-4o", "claude-haiku"])
|
|
summary = {
|
|
"result": is_known_model("gpt-4o", router),
|
|
"model": "gpt-4o",
|
|
"router_models": ["gpt-4o", "claude-haiku"],
|
|
}
|
|
assert summary == {
|
|
"result": True,
|
|
"model": "gpt-4o",
|
|
"router_models": ["gpt-4o", "claude-haiku"],
|
|
}
|
|
|
|
|
|
def test_is_known_model_returns_false_when_not_in_router():
|
|
router = _router_with_models(["gpt-4o"])
|
|
summary = {
|
|
"result": is_known_model("claude-haiku", router),
|
|
"model": "claude-haiku",
|
|
"router_models": ["gpt-4o"],
|
|
}
|
|
assert summary == {
|
|
"result": False,
|
|
"model": "claude-haiku",
|
|
"router_models": ["gpt-4o"],
|
|
}
|
|
|
|
|
|
def test_is_known_model_error_path_none_model():
|
|
router = _router_with_models(["gpt-4o"])
|
|
assert is_known_model(None, router) is False
|
|
|
|
|
|
def test_is_known_model_error_path_none_router():
|
|
assert is_known_model("gpt-4o", None) is False
|
|
|
|
|
|
def test_is_known_vector_store_index_happy_path(monkeypatch):
|
|
registry = MagicMock()
|
|
registry.get_vector_store_indexes.return_value = ["index-a", "index-b"]
|
|
monkeypatch.setattr(litellm, "vector_store_index_registry", registry)
|
|
summary = {
|
|
"result": is_known_vector_store_index("index-a"),
|
|
"indexes": ["index-a", "index-b"],
|
|
"input": "index-a",
|
|
}
|
|
assert summary == {
|
|
"result": True,
|
|
"indexes": ["index-a", "index-b"],
|
|
"input": "index-a",
|
|
}
|
|
|
|
|
|
def test_is_known_vector_store_index_returns_false_when_missing(monkeypatch):
|
|
registry = MagicMock()
|
|
registry.get_vector_store_indexes.return_value = ["index-a"]
|
|
monkeypatch.setattr(litellm, "vector_store_index_registry", registry)
|
|
summary = {
|
|
"result": is_known_vector_store_index("missing"),
|
|
"indexes": ["index-a"],
|
|
"input": "missing",
|
|
}
|
|
assert summary == {
|
|
"result": False,
|
|
"indexes": ["index-a"],
|
|
"input": "missing",
|
|
}
|
|
|
|
|
|
def test_is_known_vector_store_index_error_path_no_registry(monkeypatch):
|
|
monkeypatch.setattr(litellm, "vector_store_index_registry", None)
|
|
assert is_known_vector_store_index("anything") is False
|
|
|
|
|
|
def test_create_model_info_response_happy_path_no_metadata():
|
|
result = create_model_info_response(model_id="gpt-4o", provider="openai")
|
|
snapshot = {
|
|
"id": result["id"],
|
|
"object": result["object"],
|
|
"owned_by": result["owned_by"],
|
|
"created_is_int": isinstance(result["created"], int),
|
|
"metadata_absent": "metadata" not in result,
|
|
"max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) and result["max_input_tokens"] > 0,
|
|
"max_output_tokens_positive_int": isinstance(result["max_output_tokens"], int)
|
|
and result["max_output_tokens"] > 0,
|
|
}
|
|
assert snapshot == {
|
|
"id": "gpt-4o",
|
|
"object": "model",
|
|
"owned_by": "openai",
|
|
"created_is_int": True,
|
|
"metadata_absent": True,
|
|
"max_input_tokens_positive_int": True,
|
|
"max_output_tokens_positive_int": True,
|
|
}
|
|
|
|
|
|
def test_create_model_info_response_with_metadata_default_general(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_all_fallbacks",
|
|
lambda **_kwargs: [{"model": "fallback-1"}],
|
|
)
|
|
result = create_model_info_response(
|
|
model_id="gpt-4o",
|
|
provider="openai",
|
|
include_metadata=True,
|
|
)
|
|
snapshot = {
|
|
"id": result["id"],
|
|
"owned_by": result["owned_by"],
|
|
"object": result["object"],
|
|
"fallbacks": result["metadata"]["fallbacks"],
|
|
}
|
|
assert snapshot == {
|
|
"id": "gpt-4o",
|
|
"owned_by": "openai",
|
|
"object": "model",
|
|
"fallbacks": [{"model": "fallback-1"}],
|
|
}
|
|
|
|
|
|
def test_create_model_info_response_with_explicit_fallback_type(monkeypatch):
|
|
captured = {}
|
|
|
|
def _capture(model, llm_router, fallback_type):
|
|
captured["fallback_type"] = fallback_type
|
|
return ["x"]
|
|
|
|
monkeypatch.setattr("litellm.proxy.auth.model_checks.get_all_fallbacks", _capture)
|
|
result = create_model_info_response(
|
|
model_id="gpt-4o",
|
|
provider="openai",
|
|
include_metadata=True,
|
|
fallback_type="context_window",
|
|
)
|
|
snapshot = {
|
|
"id": result["id"],
|
|
"fallbacks": result["metadata"]["fallbacks"],
|
|
"captured_fallback_type": captured["fallback_type"],
|
|
"owned_by": result["owned_by"],
|
|
}
|
|
assert snapshot == {
|
|
"id": "gpt-4o",
|
|
"fallbacks": ["x"],
|
|
"captured_fallback_type": "context_window",
|
|
"owned_by": "openai",
|
|
}
|
|
|
|
|
|
def test_create_model_info_response_invalid_fallback_type_raises():
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
create_model_info_response(
|
|
model_id="gpt-4o",
|
|
provider="openai",
|
|
include_metadata=True,
|
|
fallback_type="bogus",
|
|
)
|
|
assert exc_info.value.status_code == 400
|
|
assert "Invalid fallback_type" in str(exc_info.value.detail)
|
|
|
|
|
|
def test_validate_model_access_happy_path_single_model_in_list():
|
|
summary = {
|
|
"result": validate_model_access("gpt-4o", ["gpt-4o", "claude-haiku"]),
|
|
"model": "gpt-4o",
|
|
"available": ["gpt-4o", "claude-haiku"],
|
|
}
|
|
assert summary == {
|
|
"result": None,
|
|
"model": "gpt-4o",
|
|
"available": ["gpt-4o", "claude-haiku"],
|
|
}
|
|
|
|
|
|
def test_validate_model_access_happy_path_batch_all_accessible():
|
|
summary = {
|
|
"result": validate_model_access("gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"]),
|
|
"input": "gpt-4o,claude-haiku",
|
|
"available": ["gpt-4o", "claude-haiku", "gemini"],
|
|
}
|
|
assert summary == {
|
|
"result": None,
|
|
"input": "gpt-4o,claude-haiku",
|
|
"available": ["gpt-4o", "claude-haiku", "gemini"],
|
|
}
|
|
|
|
|
|
def test_validate_model_access_single_model_not_accessible_raises():
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
validate_model_access("missing-model", ["gpt-4o"])
|
|
assert exc_info.value.status_code == 404
|
|
assert "missing-model" in str(exc_info.value.detail)
|
|
|
|
|
|
def test_validate_model_access_batch_partial_inaccessible_raises():
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
validate_model_access("gpt-4o,unknown-x", ["gpt-4o"])
|
|
assert exc_info.value.status_code == 404
|
|
assert "unknown-x" in str(exc_info.value.detail)
|
|
assert "gpt-4o" not in str(exc_info.value.detail).split("not accessible:")[1]
|
|
|
|
|
|
def _make_model_response():
|
|
return ModelResponse(
|
|
id="resp-123",
|
|
choices=[
|
|
{
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "do_thing", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
"index": 0,
|
|
"finish_reason": "tool_calls",
|
|
}
|
|
],
|
|
model="gpt-4o",
|
|
)
|
|
|
|
|
|
def test_model_dump_with_preserved_fields_restores_none_content():
|
|
resp = _make_model_response()
|
|
result = model_dump_with_preserved_fields(resp)
|
|
message = result["choices"][0]["message"]
|
|
snapshot = {
|
|
"content_is_none": message["content"] is None,
|
|
"role": message["role"],
|
|
"has_tool_calls": "tool_calls" in message,
|
|
"model": result["model"],
|
|
}
|
|
assert snapshot == {
|
|
"content_is_none": True,
|
|
"role": "assistant",
|
|
"has_tool_calls": True,
|
|
"model": "gpt-4o",
|
|
}
|
|
|
|
|
|
def test_model_dump_with_preserved_fields_no_choices_returns_plain_dump():
|
|
class _Bare:
|
|
def model_dump(self, **_kwargs):
|
|
return {"id": "x", "object": "y", "extra": "z"}
|
|
|
|
bare = _Bare()
|
|
result = model_dump_with_preserved_fields(bare)
|
|
assert result == {"id": "x", "object": "y", "extra": "z"}
|
|
|
|
|
|
def test_model_dump_with_preserved_fields_error_path_invalid_obj_raises():
|
|
with pytest.raises(AttributeError):
|
|
model_dump_with_preserved_fields(None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_available_models_for_user_happy_path_returns_complete_list(
|
|
monkeypatch,
|
|
):
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_key_models",
|
|
lambda **_k: ["gpt-4o"],
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_team_models",
|
|
lambda **_k: ["claude-haiku"],
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_complete_model_list",
|
|
lambda **_k: ["gpt-4o", "claude-haiku", "gemini"],
|
|
)
|
|
router = _router_with_models(["gpt-4o", "claude-haiku", "gemini"])
|
|
user_api_key_dict = UserAPIKeyAuth(
|
|
api_key="sk-test-key",
|
|
user_id="user-1",
|
|
team_id=None,
|
|
team_models=[],
|
|
)
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=user_api_key_dict,
|
|
llm_router=router,
|
|
general_settings={},
|
|
user_model=None,
|
|
)
|
|
summary = {
|
|
"result_sorted": sorted(result),
|
|
"count": len(result),
|
|
"user_id": user_api_key_dict.user_id,
|
|
"router_set": True,
|
|
}
|
|
assert summary == {
|
|
"result_sorted": ["claude-haiku", "gemini", "gpt-4o"],
|
|
"count": 3,
|
|
"user_id": "user-1",
|
|
"router_set": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_available_models_for_user_with_none_router(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_key_models",
|
|
lambda **_k: [],
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_team_models",
|
|
lambda **_k: [],
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_complete_model_list",
|
|
lambda **_k: ["user-model"],
|
|
)
|
|
user_api_key_dict = UserAPIKeyAuth(
|
|
api_key="sk-test-key",
|
|
user_id="user-1",
|
|
team_id=None,
|
|
team_models=[],
|
|
)
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=user_api_key_dict,
|
|
llm_router=None,
|
|
general_settings={},
|
|
user_model="user-model",
|
|
)
|
|
summary = {
|
|
"result": result,
|
|
"router_is_none": True,
|
|
"user_model": "user-model",
|
|
"count": len(result),
|
|
}
|
|
assert summary == {
|
|
"result": ["user-model"],
|
|
"router_is_none": True,
|
|
"user_model": "user-model",
|
|
"count": 1,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_available_models_for_user_error_path_complete_list_raises(
|
|
monkeypatch,
|
|
):
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_key_models",
|
|
lambda **_k: [],
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.auth.model_checks.get_team_models",
|
|
lambda **_k: [],
|
|
)
|
|
|
|
def _boom(**_kwargs):
|
|
raise RuntimeError("downstream failure")
|
|
|
|
monkeypatch.setattr("litellm.proxy.auth.model_checks.get_complete_model_list", _boom)
|
|
user_api_key_dict = UserAPIKeyAuth(
|
|
api_key="sk-test-key",
|
|
user_id="user-1",
|
|
team_id=None,
|
|
team_models=[],
|
|
)
|
|
with pytest.raises(RuntimeError):
|
|
await get_available_models_for_user(
|
|
user_api_key_dict=user_api_key_dict,
|
|
llm_router=None,
|
|
general_settings={},
|
|
user_model=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_available_models_for_user_resolves_team_access_group_models(
|
|
monkeypatch,
|
|
):
|
|
from litellm.models.access_group import LiteLLM_AccessGroupTable
|
|
from litellm.models.team import LiteLLM_TeamTableCachedObj
|
|
|
|
team = LiteLLM_TeamTableCachedObj(
|
|
team_id="team-1",
|
|
models=["no-default-models"],
|
|
access_group_ids=["ag-1"],
|
|
)
|
|
access_group = LiteLLM_AccessGroupTable(
|
|
access_group_id="ag-1",
|
|
access_group_name="repro-group",
|
|
access_model_names=["model-a", "model-b"],
|
|
assigned_team_ids=["team-1"],
|
|
)
|
|
|
|
async def _get_team_object(**_kwargs):
|
|
return team
|
|
|
|
async def _get_access_object(**_kwargs):
|
|
return access_group
|
|
|
|
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object)
|
|
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object)
|
|
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-test-key",
|
|
user_id="user-1",
|
|
team_id="team-1",
|
|
models=["all-team-models"],
|
|
team_models=["no-default-models"],
|
|
),
|
|
llm_router=_router_with_models(["model-a", "model-b", "model-c"]),
|
|
general_settings={},
|
|
user_model=None,
|
|
prisma_client=MagicMock(),
|
|
proxy_logging_obj=MagicMock(),
|
|
user_api_key_cache=MagicMock(),
|
|
)
|
|
assert sorted(result) == ["model-a", "model-b"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_available_models_for_user_without_access_groups_grants_nothing(
|
|
monkeypatch,
|
|
):
|
|
from litellm.models.team import LiteLLM_TeamTableCachedObj
|
|
|
|
async def _get_team_object(**_kwargs):
|
|
return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"])
|
|
|
|
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object)
|
|
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-test-key",
|
|
user_id="user-1",
|
|
team_id="team-1",
|
|
models=["all-team-models"],
|
|
team_models=["no-default-models"],
|
|
),
|
|
llm_router=_router_with_models(["model-a", "model-b"]),
|
|
general_settings={},
|
|
user_model=None,
|
|
prisma_client=MagicMock(),
|
|
proxy_logging_obj=MagicMock(),
|
|
user_api_key_cache=MagicMock(),
|
|
)
|
|
assert result == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_available_models_for_user_resolves_key_access_group_models(
|
|
monkeypatch,
|
|
):
|
|
from litellm.models.access_group import LiteLLM_AccessGroupTable
|
|
from litellm.models.team import LiteLLM_TeamTableCachedObj
|
|
|
|
async def _get_team_object(**_kwargs):
|
|
return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"])
|
|
|
|
async def _get_access_object(**_kwargs):
|
|
return LiteLLM_AccessGroupTable(
|
|
access_group_id="ag-1",
|
|
access_group_name="key-group",
|
|
access_model_names=["model-b"],
|
|
assigned_key_ids=[hash_token("sk-test-key")],
|
|
)
|
|
|
|
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object)
|
|
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
|
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-test-key",
|
|
user_id="user-1",
|
|
team_id="team-1",
|
|
models=["no-default-models"],
|
|
team_models=["no-default-models"],
|
|
access_group_ids=["ag-1"],
|
|
),
|
|
llm_router=_router_with_models(["model-a", "model-b"]),
|
|
general_settings={},
|
|
user_model=None,
|
|
prisma_client=MagicMock(),
|
|
proxy_logging_obj=MagicMock(),
|
|
user_api_key_cache=MagicMock(),
|
|
)
|
|
assert result == ["model-b"]
|
|
|
|
|
|
def _agent_ceiling(models: frozenset[str] | None):
|
|
from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling
|
|
|
|
async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None:
|
|
if models is None:
|
|
return None
|
|
return AgentAccessGroupCeiling(
|
|
access_group_ids=("ag-agent",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset()
|
|
)
|
|
|
|
return resolve
|
|
|
|
|
|
def _agent_key(models: list[str]) -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(api_key="sk-agent-key", user_id="user-1", agent_id="agent-1", models=models)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_key_listing_is_capped_to_its_access_groups():
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=_agent_key(["model-a", "model-b", "model-c"]),
|
|
llm_router=_router_with_models(["model-a", "model-b", "model-c"]),
|
|
general_settings={},
|
|
user_model=None,
|
|
resolve_agent_ceiling=_agent_ceiling(frozenset({"model-b", "model-d"})),
|
|
)
|
|
assert result == ["model-b"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_key_listing_is_empty_when_its_groups_grant_no_model():
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=_agent_key(["model-a"]),
|
|
llm_router=_router_with_models(["model-a"]),
|
|
general_settings={},
|
|
user_model=None,
|
|
resolve_agent_ceiling=_agent_ceiling(frozenset()),
|
|
)
|
|
assert result == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_ceiling_expands_a_model_access_group_name_for_listing():
|
|
router = _router_with_models(["model-a", "model-b"])
|
|
router.get_model_access_groups.return_value = {"fast-models": ["model-b"]}
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=_agent_key(["model-a", "model-b"]),
|
|
llm_router=router,
|
|
general_settings={},
|
|
user_model=None,
|
|
resolve_agent_ceiling=_agent_ceiling(frozenset({"fast-models"})),
|
|
)
|
|
assert result == ["model-b"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups():
|
|
router = _router_with_models(["model-a", "model-b"])
|
|
plain_key = await get_available_models_for_user(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-plain", user_id="user-1", models=["model-a", "model-b"]),
|
|
llm_router=router,
|
|
general_settings={},
|
|
user_model=None,
|
|
resolve_agent_ceiling=_agent_ceiling(frozenset({"model-a"})),
|
|
)
|
|
agent_without_groups = await get_available_models_for_user(
|
|
user_api_key_dict=_agent_key(["model-a", "model-b"]),
|
|
llm_router=router,
|
|
general_settings={},
|
|
user_model=None,
|
|
resolve_agent_ceiling=_agent_ceiling(None),
|
|
)
|
|
assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"])
|