From b530b0cf3ab68cb848e6bd7be43de54a29c9f62c Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Sat, 24 Oct 2026 16:56:58 +0200 Subject: [PATCH] Implemented vLLM provider --- src/providers/inference/vllm/provider.py | 118 +++++++++++++ tests/unit/providers/inference/vllm/README.md | 9 + .../unit/providers/inference/vllm/__init__.py | 1 + .../providers/inference/vllm/test_provider.py | 163 ++++++++++++++++++ 4 files changed, 291 insertions(+) create mode 100644 src/providers/inference/vllm/provider.py create mode 100644 tests/unit/providers/inference/vllm/README.md create mode 100644 tests/unit/providers/inference/vllm/__init__.py create mode 100644 tests/unit/providers/inference/vllm/test_provider.py diff --git a/src/providers/inference/vllm/provider.py b/src/providers/inference/vllm/provider.py new file mode 100644 index 000000000..947a2b0ff --- /dev/null +++ b/src/providers/inference/vllm/provider.py @@ -0,0 +1,118 @@ +"""vLLM runtime inference provider backed by Pydantic AI.""" + +from __future__ import annotations + +from typing import Any + +from pydantic_ai.embeddings import EmbeddingModel +from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel +from pydantic_ai.models import Model +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.vllm import VLLMProvider +from sentence_transformers import CrossEncoder + +from models.config import VllmInferenceProviderEntry +from providers.inference.base import InferenceProvider +from providers.inference.catalog import CatalogModel, DiscoveredModel +from providers.inference.errors import ModelTypeNotSupportedError +from providers.inference.model_type_inference import ( + infer_model_type_from_pai_known_names, +) +from providers.inference.types import ModelType + + +class VllmInferenceProvider( + InferenceProvider[VllmInferenceProviderEntry, VLLMProvider] +): + """Inference provider for vLLM OpenAI-compatible HTTP APIs.""" + + def _build_provider(self) -> VLLMProvider: + """Build the Pydantic AI vLLM provider from configuration. + + Returns: + A Pydantic AI vLLM provider using the configured base URL and + optional API key. + + """ + config = self._entry.config + api_key = ( + config.api_key.get_secret_value() if config.api_key is not None else None + ) + return VLLMProvider(api_key=api_key, base_url=config.base_url) + + async def _discover_models(self) -> list[DiscoveredModel]: + """List models from the vLLM OpenAI-compatible models endpoint. + + Returns: + Discovered models with inferred types, before allow-list filtering. + + """ + response = await self._get_provider.client.models.list() + return [ + DiscoveredModel( + provider_model_id=model.id, + model_type=self._infer_model_type(model.id), + ) + for model in response.data + ] + + def _create_llm(self, model: CatalogModel) -> Model: + """Create a Chat Completions model for a catalog entry. + + Args: + model: Catalog entry describing the LLM to instantiate. + + Returns: + A Pydantic AI OpenAI chat model bound to this provider. + + """ + return OpenAIChatModel( + model.provider_model_id, + provider=self._get_provider, + ) + + def _create_embedding(self, model: CatalogModel) -> EmbeddingModel: + """Create an embedding model for a catalog entry. + + Args: + model: Catalog entry describing the embedding model to instantiate. + + Returns: + A Pydantic AI OpenAI embedding model bound to this provider. + + """ + return OpenAIEmbeddingModel( + model.provider_model_id, + provider=self._get_provider, + ) + + def _create_rerank(self, model: CatalogModel) -> CrossEncoder: + """Reject rerank requests; this provider has no rerank API. + + Args: + model: Catalog entry that requested rerank support. + + Raises: + ModelTypeNotSupportedError: Always, because rerank is unsupported. + + """ + _ = model + raise ModelTypeNotSupportedError("rerank", self.provider_id) + + def _infer_model_type( + self, + provider_model_id: str, + **kwargs: Any, + ) -> ModelType: + """Infer model type from Pydantic AI known LLM and embedding names. + + Args: + provider_model_id: Model id from vLLM or from an allow-list entry. + **kwargs: Unused; accepted for base-class compatibility. + + Returns: + embedding when the id is a known embedding model, otherwise llm. + + """ + _ = kwargs + return infer_model_type_from_pai_known_names(provider_model_id) or "llm" diff --git a/tests/unit/providers/inference/vllm/README.md b/tests/unit/providers/inference/vllm/README.md new file mode 100644 index 000000000..949cd2e52 --- /dev/null +++ b/tests/unit/providers/inference/vllm/README.md @@ -0,0 +1,9 @@ +# List of source files stored in `tests/unit/providers/inference/vllm` directory + +## [__init__.py](__init__.py) + +Unit tests for the vLLM inference provider. + +## [test_provider.py](test_provider.py) + +Unit tests for the vLLM inference provider. diff --git a/tests/unit/providers/inference/vllm/__init__.py b/tests/unit/providers/inference/vllm/__init__.py new file mode 100644 index 000000000..785e62255 --- /dev/null +++ b/tests/unit/providers/inference/vllm/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the vLLM inference provider.""" diff --git a/tests/unit/providers/inference/vllm/test_provider.py b/tests/unit/providers/inference/vllm/test_provider.py new file mode 100644 index 000000000..c247eb705 --- /dev/null +++ b/tests/unit/providers/inference/vllm/test_provider.py @@ -0,0 +1,163 @@ +"""Unit tests for the vLLM inference provider.""" + +# pylint: disable=redefined-outer-name,protected-access + +from typing import Any + +import pytest +from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.vllm import VLLMProvider + +from models.config import VllmInferenceProviderEntry +from providers.inference.catalog import CatalogModel +from providers.inference.errors import ModelTypeNotSupportedError +from providers.inference.vllm.provider import VllmInferenceProvider + +BASE_URL = "https://example.test/v1" + + +def _entry(**config: Any) -> VllmInferenceProviderEntry: + """Build a vLLM provider configuration.""" + return VllmInferenceProviderEntry.model_validate( + { + "type": "vllm", + "id": "vllm", + "config": { + "base_url": BASE_URL, + "api_key": "test-key", + **config, + }, + } + ) + + +@pytest.fixture +def provider() -> VllmInferenceProvider: + """Create a vLLM provider.""" + return VllmInferenceProvider(_entry()) + + +def test_build_provider_uses_config(provider: VllmInferenceProvider) -> None: + """Builds a Pydantic AI provider with configured credentials and URL.""" + pai_provider = provider._build_provider() + + assert isinstance(pai_provider, VLLMProvider) + assert str(pai_provider.client.base_url).rstrip("/") == BASE_URL + assert pai_provider.client.api_key == "test-key" + + +def test_build_provider_allows_missing_api_key() -> None: + """Builds a Pydantic AI provider when no API key is configured.""" + pai_provider = VllmInferenceProvider(_entry(api_key=None))._build_provider() + + assert isinstance(pai_provider, VLLMProvider) + assert str(pai_provider.client.base_url).rstrip("/") == BASE_URL + assert pai_provider.client.api_key == "api-key-not-set" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_id", "expected_type"), + [ + ("gpt-5.6-luna", "llm"), + ("text-embedding-3-small", "embedding"), + ("meta-llama/Llama-3.2-1B-Instruct", "llm"), + ("custom-model", "llm"), + ], +) +async def test_discover_models_infers_model_types( + provider: VllmInferenceProvider, + mocker: Any, + model_id: str, + expected_type: str, +) -> None: + """Discovers vLLM model IDs and applies the model-type inference fallback.""" + response = mocker.MagicMock() + response.data = [mocker.MagicMock(id=model_id)] + + pai_provider = mocker.MagicMock(spec=VLLMProvider) + pai_provider.client.models.list = mocker.AsyncMock(return_value=response) + provider._provider = pai_provider + + discovered = await provider._discover_models() + + assert len(discovered) == 1 + assert discovered[0].provider_model_id == model_id + assert discovered[0].model_type == expected_type + pai_provider.client.models.list.assert_awaited_once_with() + + +def test_create_llm_uses_openai_chat_model( + provider: VllmInferenceProvider, +) -> None: + """Creates a Pydantic AI chat model using the provider model ID.""" + pai_provider = provider._build_provider() + provider._provider = pai_provider + + model = provider._create_llm( + CatalogModel( + name="chat", + provider_model_id="meta-llama/Llama-3.2-1B-Instruct", + model_type="llm", + ) + ) + + assert isinstance(model, OpenAIChatModel) + assert model.model_name == "meta-llama/Llama-3.2-1B-Instruct" + assert model.provider is pai_provider + + +def test_create_embedding_uses_openai_embedding_model( + provider: VllmInferenceProvider, +) -> None: + """Creates a Pydantic AI embedding model using the provider model ID.""" + pai_provider = provider._build_provider() + provider._provider = pai_provider + + model = provider._create_embedding( + CatalogModel( + name="embed", + provider_model_id="text-embedding-3-small", + model_type="embedding", + ) + ) + + assert isinstance(model, OpenAIEmbeddingModel) + assert model.model_name == "text-embedding-3-small" + assert model._provider is pai_provider + + +def test_create_rerank_is_unsupported( + provider: VllmInferenceProvider, +) -> None: + """Rejects reranking because the vLLM provider does not support it.""" + with pytest.raises(ModelTypeNotSupportedError) as exc_info: + provider._create_rerank( + CatalogModel( + name="rerank", + provider_model_id="rerank-model", + model_type="rerank", + ) + ) + + assert exc_info.value.model_type == "rerank" + assert exc_info.value.provider_id == "vllm" + + +@pytest.mark.parametrize( + ("model_id", "expected_type"), + [ + ("text-embedding-3-small", "embedding"), + ("gpt-5.6-luna", "llm"), + ("meta-llama/Llama-3.2-1B-Instruct", "llm"), + ("custom-model", "llm"), + ], +) +def test_infer_model_type_uses_known_names_with_llm_fallback( + provider: VllmInferenceProvider, + model_id: str, + expected_type: str, +) -> None: + """Infers embedding from known names; unknown ids (including HF paths) are llm.""" + assert provider._infer_model_type(model_id) == expected_type