import asyncio import threading from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from typing import Final from unittest.mock import MagicMock import httpx import pytest import respx from fastapi import HTTPException import litellm from litellm import ModelResponse, Router 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.registry.return_value.names.return_value = model_names router.registry.return_value.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_lists_team_access_group_models_only_where_auth_grants_them( monkeypatch, ): from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy._types import Member teams = { "team-1": LiteLLM_TeamTableCachedObj(team_id="team-1", models=["model-a"], access_group_ids=["ag-1"]), "team-2": LiteLLM_TeamTableCachedObj( team_id="team-2", models=["model-c"], members_with_roles=[Member(user_id="user-1", role="user")] ), } async def _get_team_object(team_id: str, **_kwargs): return teams[team_id] async def _get_access_object(**_kwargs): return LiteLLM_AccessGroupTable( access_group_id="ag-1", access_group_name="team-group", access_model_names=["model-b"], assigned_team_ids=["team-1"], ) 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) async def _list_for(key_models: list[str], config: dict[str, list[str]] | None = None, team_id: str | None = None): return sorted( 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=key_models, team_models=["model-a"], config=config or {}, ), 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(), team_id=team_id, ) ) assert await _list_for(["model-a"]) == ["model-a"] assert await _list_for(["model-a"], team_id="team-1") == ["model-a"] assert await _list_for(["model-a"], team_id="team-2") == ["model-c"] assert await _list_for(["model-a"], config={"model_list": []}) == ["model-a", "model-b"] assert await _list_for(["all-team-models"]) == ["model-a", "model-b"] assert await _list_for([]) == ["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.registry.return_value.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"]) OPENROUTER_BASE: Final = "https://openrouter.ai/api/v1" @dataclass(frozen=True, slots=True) class _CatalogHeldUntilTheLoopRuns: requested: threading.Event = field(default_factory=threading.Event) released: threading.Event = field(default_factory=threading.Event) def answer(self, _request: httpx.Request) -> httpx.Response: self.requested.set() if not self.released.wait(timeout=30): return httpx.Response(504) return httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) async def release_from_the_loop(self) -> None: assert await asyncio.to_thread(self.requested.wait, 30) self.released.set() def test_get_available_models_for_user_fetches_a_live_catalog_without_holding_the_event_loop_or_its_default_pool( monkeypatch: pytest.MonkeyPatch, ): monkeypatch.setattr(litellm, "check_provider_endpoint", True) catalog: Final = _CatalogHeldUntilTheLoopRuns() router: Final = Router( model_list=[ { "model_name": "openrouter/*", "litellm_params": { "model": "openrouter/*", "api_key": "sk-or-available-models-off-loop", "api_base": OPENROUTER_BASE, }, } ] ) async def list_models_on_a_one_thread_default_pool() -> list[str]: asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) listing: Final = asyncio.ensure_future( get_available_models_for_user( user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-key", user_id="user-1"), llm_router=router, general_settings={}, user_model=None, ) ) await catalog.release_from_the_loop() return await listing with respx.mock() as respx_mock: catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer) models: Final = asyncio.run(list_models_on_a_one_thread_default_pool()) assert catalog_route.calls.last.response.status_code == 200 assert models == ["openrouter/openai/gpt-5.4-mini"]