diff --git a/tests/integration/inference_registry/__init__.py b/tests/integration/inference_registry/__init__.py new file mode 100644 index 000000000..ff980c89d --- /dev/null +++ b/tests/integration/inference_registry/__init__.py @@ -0,0 +1 @@ +"""Integration tests for inference provider registry startup.""" diff --git a/tests/integration/inference_registry/_helpers.py b/tests/integration/inference_registry/_helpers.py new file mode 100644 index 000000000..fbaa4978d --- /dev/null +++ b/tests/integration/inference_registry/_helpers.py @@ -0,0 +1,261 @@ +"""Shared helpers for inference registry startup integration tests.""" + +from collections.abc import Iterable +from pathlib import Path +from typing import Any, Optional + +import httpx +import yaml +from botocore.exceptions import ClientError +from openai import APIConnectionError +from pytest_mock import MockerFixture + +from configuration import configuration +from providers.inference.types import ProviderType + +# Shared logical allow-list / discovery fixtures (not provider-realistic ids). +CHAT = { + "name": "chat", + "provider_model_id": "chat-id", + "model_type": "llm", +} +EMBED = { + "name": "embed", + "provider_model_id": "embed-id", + "model_type": "embedding", +} +RERANK = { + "name": "rerank", + "provider_model_id": "rerank-id", + "model_type": "rerank", +} +MISSING = { + "name": "missing", + "provider_model_id": "not-discovered", + "model_type": "llm", +} +DROPPED_LLM = { + "name": "dropped", + "provider_model_id": "dropped-id", + "model_type": "llm", +} +EXTRA_DISCOVERED_ID = "extra-id" +BEDROCK_PROFILE = { + "name": "profile", + "provider_model_id": "profile-id", + "model_type": "llm", +} +BEDROCK_FOUNDATION_ID = "foundation-id" +BEDROCK_UNSUPPORTED_ID = "image-id" + +_SKIP_CONFIGS: dict[ProviderType, dict[str, Any]] = { + "openai": {"api_key": "test-key", "base_url": "https://api.openai.com/v1"}, + "azure": { + "api_key": "test-key", + "base_url": "https://example.openai.azure.com", + "api_version": "2024-12-01-preview", + }, + "vertexai": {"project": "test-project", "location": "us-central1"}, + "bedrock": {"region": "us-east-1", "api_key": "test-bearer-token"}, + "watsonx": { + "api_key": "test-key", + "project_id": "test-project", + "base_url": "https://us-south.ml.cloud.ibm.com", + }, + "vllm": {"base_url": "http://localhost:8000/v1", "api_key": "unused"}, + "sentence_transformers": {"device": "cpu"}, +} + +PROVIDER_SKIP_DISCOVERY_ENTRIES: dict[ProviderType, dict[str, Any]] = { + provider_type: { + "type": provider_type, + "id": provider_type, + "skip_model_discovery": True, + "allowed_models": ( + [EMBED, RERANK] if provider_type == "sentence_transformers" else [CHAT] + ), + "config": config, + } + for provider_type, config in _SKIP_CONFIGS.items() +} + + +def write_config(tmp_path: Path, *, inference: dict[str, Any]) -> Path: + """Write a minimal LCS config with the given inference section.""" + config_path = tmp_path / "lightspeed-stack.yaml" + payload: dict[str, Any] = { + "name": "registry-startup-test", + "service": { + "host": "localhost", + "port": 8080, + "auth_enabled": False, + "workers": 1, + "color_log": True, + "access_log": True, + }, + "ogx": { + "use_as_library_client": False, + "url": "http://localhost:8321", + "api_key": "xyzzy", + "allow_degraded_mode": True, + }, + "user_data_collection": {"feedback_enabled": False}, + "mcp_servers": [], + "inference": inference, + } + config_path.write_text(yaml.dump(payload), encoding="utf-8") + return config_path + + +def load_inference(tmp_path: Path, providers: list[dict[str, Any]]) -> None: + """Write and load an LCS config whose inference.providers is ``providers``.""" + configuration.load_configuration( + str(write_config(tmp_path, inference={"providers": providers})) + ) + + +def provider_entry( + provider_type: ProviderType, + *, + skip_model_discovery: bool = True, + allowed_models: Optional[list[dict[str, Any]]] = None, +) -> dict[str, Any]: + """Shallow-copy a skip-discovery fixture entry with optional overrides.""" + entry = {**PROVIDER_SKIP_DISCOVERY_ENTRIES[provider_type]} + entry["skip_model_discovery"] = skip_model_discovery + if allowed_models is not None: + entry["allowed_models"] = allowed_models + return entry + + +def mock_openai_models_list( + mocker: MockerFixture, + *model_ids: str, + side_effect: Optional[Exception] = None, +) -> Any: + """Mock the foreign OpenAI ``models.list`` API used during discovery.""" + if side_effect is not None: + return mocker.patch( + "openai.resources.models.AsyncModels.list", + new=mocker.AsyncMock(side_effect=side_effect), + ) + + response = mocker.MagicMock() + response.data = [mocker.MagicMock(id=model_id) for model_id in model_ids] + return mocker.patch( + "openai.resources.models.AsyncModels.list", + new=mocker.AsyncMock(return_value=response), + ) + + +def openai_connection_error() -> APIConnectionError: + """Return a realistic OpenAI transport failure from models.list.""" + return APIConnectionError( + message="connection failed", + request=httpx.Request( # pyright: ignore[reportArgumentType] + "GET", "https://api.openai.com/v1/models" + ), + ) + + +def mock_bedrock_control_plane( + mocker: MockerFixture, + foundations: Iterable[dict[str, Any]], + *, + profiles: Optional[Iterable[dict[str, Any]]] = None, + side_effect: Optional[Exception] = None, +) -> Any: + """Mock Bedrock ListFoundationModels / ListInferenceProfiles for discovery.""" + client = mocker.MagicMock() + if side_effect is not None: + client.list_foundation_models.side_effect = side_effect + else: + client.list_foundation_models.return_value = { + "modelSummaries": list(foundations) + } + paginator = mocker.MagicMock() + paginator.paginate.return_value = [ + {"inferenceProfileSummaries": list(profiles or [])} + ] + client.get_paginator.return_value = paginator + return mocker.patch( + "providers.inference.bedrock.provider.BedrockInferenceProvider" + "._create_bedrock_client", + return_value=client, + ) + + +def bedrock_inference_profile( + profile_id: str, + *foundation_model_ids: str, + region: str = "us-east-1", +) -> dict[str, Any]: + """Build a ListInferenceProfiles wire summary for discovery mocks.""" + return { + "inferenceProfileId": profile_id, + "models": [ + {"modelArn": (f"arn:aws:bedrock:{region}::foundation-model/{model_id}")} + for model_id in foundation_model_ids + ], + } + + +def bedrock_access_denied_error() -> ClientError: + """Return a realistic Bedrock ListFoundationModels access failure.""" + return ClientError( + {"Error": {"Code": "AccessDeniedException", "Message": "denied"}}, + "ListFoundationModels", + ) + + +def bedrock_foundation( + model_id: str, + *, + modalities: Optional[list[str]] = None, + inference_types: Optional[list[str]] = None, +) -> dict[str, Any]: + """Build a ListFoundationModels wire summary for discovery mocks.""" + summary: dict[str, Any] = {"modelId": model_id} + if modalities is not None: + summary["outputModalities"] = modalities + if inference_types is not None: + summary["inferenceTypesSupported"] = inference_types + return summary + + +def mock_vertex_models_list( + mocker: MockerFixture, + *model_names: str, +) -> Any: + """Mock Vertex ``aio.models.list``; return the awaitable mock for assertions.""" + + class _AsyncIter: + def __init__(self, items: list[Any]) -> None: + self._items = items + self._iter = iter(items) + + def __aiter__(self) -> "_AsyncIter": + self._iter = iter(self._items) + return self + + async def __anext__(self) -> Any: + try: + return next(self._iter) + except StopIteration as exc: + raise StopAsyncIteration from exc + + models = [] + for name in model_names: + model = mocker.MagicMock() + model.name = name + model.supported_actions = ["generateContent"] + models.append(model) + + list_mock = mocker.AsyncMock(return_value=_AsyncIter(models)) + pai_provider = mocker.MagicMock() + pai_provider.client.aio.models.list = list_mock + mocker.patch( + "providers.inference.vertexai.provider.GoogleCloudProvider", + return_value=pai_provider, + ) + return list_mock diff --git a/tests/integration/inference_registry/conftest.py b/tests/integration/inference_registry/conftest.py new file mode 100644 index 000000000..3d6b17a3a --- /dev/null +++ b/tests/integration/inference_registry/conftest.py @@ -0,0 +1,20 @@ +"""Fixtures for inference registry startup integration tests.""" + +from collections.abc import Generator + +import pytest + +from providers.inference.registry import ProviderRegistry +from utils.types import Singleton + + +@pytest.fixture(autouse=True) +def _reset_provider_registry() -> Generator[None, None, None]: + """Reset the singleton registry between integration tests.""" + # pylint: disable=protected-access + Singleton._instances.pop(ProviderRegistry, None) + + yield + + Singleton._instances.pop(ProviderRegistry, None) + # pylint: enable=protected-access diff --git a/tests/integration/inference_registry/test_startup.py b/tests/integration/inference_registry/test_startup.py new file mode 100644 index 000000000..65668f83e --- /dev/null +++ b/tests/integration/inference_registry/test_startup.py @@ -0,0 +1,361 @@ +"""Integration tests for inference provider registry startup.""" + +from pathlib import Path + +import pytest +from pytest_mock import MockerFixture + +from configuration import configuration +from providers.inference.bedrock.provider import BedrockInferenceProvider +from providers.inference.errors import ModelDiscoveryError +from providers.inference.openai.provider import OpenAIInferenceProvider +from providers.inference.registry import PROVIDER_CLASSES, ProviderRegistry +from providers.inference.sentence_transformers.provider import ( + SentenceTransformersInferenceProvider, +) +from providers.inference.types import ProviderType +from providers.inference.vertexai.provider import VertexAIInferenceProvider +from providers.inference.vllm.provider import VllmInferenceProvider +from tests.integration.inference_registry._helpers import ( + BEDROCK_FOUNDATION_ID, + BEDROCK_PROFILE, + BEDROCK_UNSUPPORTED_ID, + CHAT, + DROPPED_LLM, + EMBED, + EXTRA_DISCOVERED_ID, + MISSING, + PROVIDER_SKIP_DISCOVERY_ENTRIES, + RERANK, + bedrock_access_denied_error, + bedrock_foundation, + bedrock_inference_profile, + load_inference, + mock_bedrock_control_plane, + mock_openai_models_list, + mock_vertex_models_list, + openai_connection_error, + provider_entry, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_type", sorted(PROVIDER_CLASSES)) +async def test_registry_initializes_each_registered_provider_type( + tmp_path: Path, + provider_type: ProviderType, +) -> None: + """Each PROVIDER_CLASSES type initializes from config via skip-discovery.""" + entry = PROVIDER_SKIP_DISCOVERY_ENTRIES[provider_type] + load_inference(tmp_path, [entry]) + + await ProviderRegistry().initialize(configuration.inference) + + registry = ProviderRegistry() + provider = registry.get_provider(entry["id"]) + assert isinstance(provider, PROVIDER_CLASSES[provider_type]) + assert [ + model.model_dump(exclude={"metadata"}) for model in registry.list_models() + ] == entry["allowed_models"] + + +@pytest.mark.asyncio +async def test_registry_skip_discovery_does_not_call_openai_models_list( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """skip_model_discovery=True must not hit the foreign models.list API.""" + list_mock = mock_openai_models_list(mocker, CHAT["provider_model_id"]) + load_inference(tmp_path, [PROVIDER_SKIP_DISCOVERY_ENTRIES["openai"]]) + + await ProviderRegistry().initialize(configuration.inference) + + assert [ + model.model_dump(exclude={"metadata"}) + for model in ProviderRegistry().list_models() + ] == [CHAT] + list_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_registry_initialize_with_empty_providers_succeeds( + tmp_path: Path, +) -> None: + """Empty inference.providers[] leaves an empty but initialized registry.""" + load_inference(tmp_path, []) + + await ProviderRegistry().initialize(configuration.inference) + + assert ProviderRegistry().list_providers() == [] + assert ProviderRegistry().list_models() == [] + + +@pytest.mark.asyncio +async def test_registry_discovery_filters_allow_list_via_base_catalog( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """Base ``_load_catalog`` keep/drop logic runs once through OpenAI.""" + list_mock = mock_openai_models_list( + mocker, + CHAT["provider_model_id"], + EMBED["provider_model_id"], + EXTRA_DISCOVERED_ID, + ) + load_inference( + tmp_path, + [ + provider_entry( + "openai", + skip_model_discovery=False, + allowed_models=[CHAT, EMBED, MISSING], + ) + ], + ) + + await ProviderRegistry().initialize(configuration.inference) + + registry = ProviderRegistry() + assert isinstance(registry.get_provider("openai"), OpenAIInferenceProvider) + assert [model.name for model in registry.list_models()] == [ + CHAT["name"], + EMBED["name"], + ] + assert [model.name for model in registry.list_models(model_type="embedding")] == [ + EMBED["name"] + ] + list_mock.assert_awaited() + + +@pytest.mark.asyncio +async def test_registry_discovery_empty_allow_list_publishes_all_models( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """Empty allow-list publishes every discovered model under its wire id.""" + list_mock = mock_openai_models_list( + mocker, + CHAT["provider_model_id"], + EMBED["provider_model_id"], + ) + load_inference( + tmp_path, + [ + provider_entry( + "openai", + skip_model_discovery=False, + allowed_models=[], + ) + ], + ) + + await ProviderRegistry().initialize(configuration.inference) + + models = ProviderRegistry().list_models() + assert [(model.name, model.provider_model_id) for model in models] == [ + (CHAT["provider_model_id"], CHAT["provider_model_id"]), + (EMBED["provider_model_id"], EMBED["provider_model_id"]), + ] + list_mock.assert_awaited() + + +@pytest.mark.asyncio +async def test_registry_vllm_discovery_builds_catalog_from_models_list( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """vLLM discovery uses the OpenAI-compatible models.list through the registry.""" + list_mock = mock_openai_models_list( + mocker, + CHAT["provider_model_id"], + EXTRA_DISCOVERED_ID, + ) + load_inference( + tmp_path, + [ + provider_entry( + "vllm", + skip_model_discovery=False, + allowed_models=[CHAT], + ) + ], + ) + + await ProviderRegistry().initialize(configuration.inference) + + registry = ProviderRegistry() + assert isinstance(registry.get_provider("vllm"), VllmInferenceProvider) + assert [ + model.model_dump(exclude={"metadata"}) for model in registry.list_models() + ] == [CHAT] + list_mock.assert_awaited() + + +@pytest.mark.asyncio +async def test_registry_bedrock_discovery_builds_catalog_from_control_plane( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """Bedrock keeps on-demand and profile ids; skips unsupported modalities.""" + mock_bedrock_control_plane( + mocker, + [ + bedrock_foundation( + CHAT["provider_model_id"], + modalities=["TEXT"], + inference_types=["ON_DEMAND"], + ), + bedrock_foundation( + EMBED["provider_model_id"], + modalities=["EMBEDDING"], + inference_types=["ON_DEMAND"], + ), + bedrock_foundation( + BEDROCK_FOUNDATION_ID, + modalities=["TEXT"], + inference_types=["INFERENCE_PROFILE"], + ), + bedrock_foundation( + BEDROCK_UNSUPPORTED_ID, + modalities=["IMAGE"], + inference_types=["ON_DEMAND"], + ), + ], + profiles=[ + bedrock_inference_profile( + BEDROCK_PROFILE["provider_model_id"], + BEDROCK_FOUNDATION_ID, + ), + ], + ) + load_inference( + tmp_path, + [ + provider_entry( + "bedrock", + skip_model_discovery=False, + allowed_models=[CHAT, EMBED, BEDROCK_PROFILE], + ) + ], + ) + + await ProviderRegistry().initialize(configuration.inference) + + registry = ProviderRegistry() + assert isinstance(registry.get_provider("bedrock"), BedrockInferenceProvider) + assert [ + model.model_dump(exclude={"metadata"}) for model in registry.list_models() + ] == [CHAT, EMBED, BEDROCK_PROFILE] + + +@pytest.mark.asyncio +async def test_registry_vertex_discovery_keeps_catalog_resource_names( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """Vertex discovery keeps discovered model names as provider model ids.""" + list_mock = mock_vertex_models_list( + mocker, + CHAT["provider_model_id"], + EXTRA_DISCOVERED_ID, + ) + load_inference( + tmp_path, + [ + provider_entry( + "vertexai", + skip_model_discovery=False, + allowed_models=[CHAT], + ) + ], + ) + + await ProviderRegistry().initialize(configuration.inference) + + registry = ProviderRegistry() + assert isinstance(registry.get_provider("vertexai"), VertexAIInferenceProvider) + assert [ + model.model_dump(exclude={"metadata"}) for model in registry.list_models() + ] == [CHAT] + list_mock.assert_awaited() + + +@pytest.mark.asyncio +async def test_registry_sentence_transformers_catalog_drops_llm_entries( + tmp_path: Path, +) -> None: + """Sentence-transformers catalog override keeps embed/rerank and drops llm.""" + load_inference( + tmp_path, + [ + provider_entry( + "sentence_transformers", + skip_model_discovery=False, + allowed_models=[EMBED, RERANK, DROPPED_LLM], + ) + ], + ) + + await ProviderRegistry().initialize(configuration.inference) + + registry = ProviderRegistry() + provider = registry.get_provider("sentence_transformers") + assert isinstance(provider, SentenceTransformersInferenceProvider) + assert [(model.name, model.model_type) for model in registry.list_models()] == [ + (EMBED["name"], EMBED["model_type"]), + (RERANK["name"], RERANK["model_type"]), + ] + + +@pytest.mark.asyncio +async def test_registry_openai_discovery_failure_aborts_initialization( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """OpenAI models.list transport failure aborts registry initialization.""" + mock_openai_models_list(mocker, side_effect=openai_connection_error()) + load_inference( + tmp_path, + [ + provider_entry( + "openai", + skip_model_discovery=False, + allowed_models=[CHAT], + ) + ], + ) + + with pytest.raises(ModelDiscoveryError) as exc_info: + await ProviderRegistry().initialize(configuration.inference) + + assert exc_info.value.provider_id == "openai" + assert ProviderRegistry().list_providers() == [] + + +@pytest.mark.asyncio +async def test_registry_bedrock_discovery_failure_aborts_initialization( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + """Bedrock control-plane AccessDenied aborts registry initialization.""" + mock_bedrock_control_plane( + mocker, + [], + side_effect=bedrock_access_denied_error(), + ) + load_inference( + tmp_path, + [ + provider_entry( + "bedrock", + skip_model_discovery=False, + allowed_models=[CHAT], + ) + ], + ) + + with pytest.raises(ModelDiscoveryError) as exc_info: + await ProviderRegistry().initialize(configuration.inference) + + assert exc_info.value.provider_id == "bedrock" + assert ProviderRegistry().list_providers() == []