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"])