LCORE-4709: Implemented sentence transformers provider

src/providers/inference/registry.py

@@ -8,6 +8,9 @@
 from providers.inference.bedrock.provider import BedrockInferenceProvider
 from providers.inference.errors import ProviderNotFoundError
 from providers.inference.openai.provider import OpenAIInferenceProvider
+from providers.inference.sentence_transformers.provider import (
+    SentenceTransformersInferenceProvider,
+)
 from providers.inference.vertexai.provider import VertexAIInferenceProvider
 from providers.inference.watsonx.provider import WatsonxInferenceProvider
 from utils.types import Singleton
@@ -24,6 +27,7 @@
     "vertexai": VertexAIInferenceProvider,
     "bedrock": BedrockInferenceProvider,
     "watsonx": WatsonxInferenceProvider,
+    "sentence_transformers": SentenceTransformersInferenceProvider,
 }
 
 

src/providers/inference/sentence_transformers/README.md

@@ -4,3 +4,10 @@
 
 Sentence-transformers inference provider package.
 
+## [model_cache.py](model_cache.py)
+
+In-process cache for sentence-transformers embedding and rerank models.
+
+## [provider.py](provider.py)
+
+Local sentence-transformers inference provider (embedding and rerank).

src/providers/inference/sentence_transformers/__init__.py

@@ -1 +1,7 @@
 """Sentence-transformers inference provider package."""
+
+from providers.inference.sentence_transformers.provider import (
+    SentenceTransformersInferenceProvider,
+)
+
+__all__ = ["SentenceTransformersInferenceProvider"]

src/providers/inference/sentence_transformers/model_cache.py

@@ -0,0 +1,72 @@
+"""In-process cache for sentence-transformers embedding and rerank models."""
+
+from __future__ import annotations
+
+from typing import Literal, Optional, overload
+
+from pydantic_ai.embeddings import EmbeddingModel
+from sentence_transformers import CrossEncoder
+
+
+class SentenceTransformersModelCache:
+    """Process-local cache of loaded embedding and rerank model instances."""
+
+    def __init__(self) -> None:
+        """Initialize empty embedding and rerank caches."""
+        self._embeddings: dict[str, EmbeddingModel] = {}
+        self._rerankers: dict[str, CrossEncoder] = {}
+
+    @overload
+    def get(
+        self,
+        provider_model_id: str,
+        model_type: Literal["embedding"],
+    ) -> Optional[EmbeddingModel]: ...
+
+    @overload
+    def get(
+        self,
+        provider_model_id: str,
+        model_type: Literal["rerank"],
+    ) -> Optional[CrossEncoder]: ...
+
+    def get(
+        self,
+        provider_model_id: str,
+        model_type: Literal["embedding", "rerank"],
+    ) -> Optional[EmbeddingModel | CrossEncoder]:
+        """Return a cached model when present.
+
+        Args:
+            provider_model_id: Cache key; the Hugging Face model id.
+            model_type: Capability whose cache to read (embedding or rerank).
+
+        Returns:
+            The cached model, or None if it has not been stored.
+
+        """
+        match model_type:
+            case "embedding":
+                return self._embeddings.get(provider_model_id)
+            case "rerank":
+                return self._rerankers.get(provider_model_id)
+
+    def put(
+        self,
+        provider_model_id: str,
+        model: EmbeddingModel | CrossEncoder,
+    ) -> None:
+        """Store a model in the cache.
+
+        The target store is chosen from the instance type: CrossEncoder goes
+        to the rerank cache, anything else to the embedding cache.
+
+        Args:
+            provider_model_id: Cache key; the Hugging Face model id.
+            model: Model instance to cache.
+
+        """
+        if isinstance(model, CrossEncoder):
+            self._rerankers[provider_model_id] = model
+        else:
+            self._embeddings[provider_model_id] = model

src/providers/inference/sentence_transformers/provider.py

@@ -0,0 +1,185 @@
+"""Local sentence-transformers inference provider (embedding and rerank)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from pydantic_ai.embeddings.sentence_transformers import (
+    SentenceTransformerEmbeddingModel,
+)
+from pydantic_ai.models import Model
+from pydantic_ai.providers import Provider
+from sentence_transformers import CrossEncoder, SentenceTransformer
+
+from log import get_logger
+from models.config import SentenceTransformersInferenceProviderEntry
+from providers.inference.base import InferenceProvider
+from providers.inference.catalog import CatalogModel, DiscoveredModel
+from providers.inference.errors import ModelTypeNotSupportedError
+from providers.inference.sentence_transformers.model_cache import (
+    SentenceTransformersModelCache,
+)
+
+if TYPE_CHECKING:
+    from pydantic_ai.embeddings import EmbeddingModel
+
+    from providers.inference.types import ModelType
+
+logger = get_logger(__name__)
+
+
+class SentenceTransformersInferenceProvider(
+    InferenceProvider[SentenceTransformersInferenceProviderEntry, Provider[Any]]
+):
+    """Inference provider for local sentence-transformers embedding and rerank.
+
+    This provider has no HTTP client and no remote discovery. The configured
+    allowed_models list is the catalog. Chat models are never exposed.
+    """
+
+    def __init__(self, entry: SentenceTransformersInferenceProviderEntry) -> None:
+        """Initialize with an empty in-process model cache.
+
+        Args:
+            entry: Provider configuration entry with a non-empty allow-list.
+
+        """
+        super().__init__(entry)
+        self._cache = SentenceTransformersModelCache()
+
+    def _build_provider(self) -> Optional[Provider[Any]]:
+        """Return no Pydantic AI provider; models run in-process.
+
+        Returns:
+            Always None.
+
+        """
+        return None
+
+    async def _discover_models(self) -> list[DiscoveredModel]:
+        """Remote discovery is unsupported for sentence-transformers.
+
+        Returns:
+            An empty list. Catalog construction uses allowed_models only.
+
+        """
+        return []
+
+    async def _load_catalog(self) -> list[CatalogModel]:
+        """Build the catalog exclusively from configured allowed models.
+
+        Equivalent to the base-class skip_model_discovery path. Entries typed
+        as llm are omitted so this provider never exposes chat models.
+
+        Returns:
+            Catalog rows for embedding and rerank models only.
+
+        """
+        catalog: list[CatalogModel] = []
+        for allowed in self._entry.allowed_models:
+            model_type = allowed.model_type or self._infer_model_type(
+                allowed.provider_model_id
+            )
+            if model_type == "llm":
+                logger.warning(
+                    "Omitting configured model %r with provider model ID %r "
+                    "from sentence-transformers provider %r; llm is unsupported",
+                    allowed.name,
+                    allowed.provider_model_id,
+                    self.provider_id,
+                )
+                continue
+            catalog.append(
+                CatalogModel(
+                    name=allowed.name,
+                    provider_model_id=allowed.provider_model_id,
+                    model_type=model_type,
+                    metadata=dict(allowed.metadata),
+                )
+            )
+        return catalog
+
+    def _create_llm(self, model: CatalogModel) -> Model:
+        """Reject LLM construction; this provider never serves chat models.
+
+        Args:
+            model: Catalog entry that requested LLM support.
+
+        Raises:
+            ModelTypeNotSupportedError: Always, because llm is unsupported.
+
+        """
+        _ = model
+        raise ModelTypeNotSupportedError("llm", self.provider_id)
+
+    def _create_embedding(self, model: CatalogModel) -> EmbeddingModel:
+        """Create or reuse a cached Pydantic AI sentence-transformers embedding model.
+
+        Args:
+            model: Catalog entry describing the embedding model to instantiate.
+
+        Returns:
+            A cached Pydantic AI embedding model for the provider model id.
+
+        """
+        provider_model_id = model.provider_model_id
+        cached = self._cache.get(provider_model_id, "embedding")
+        if cached is not None:
+            return cached
+
+        config = self._entry.config
+        embedding = SentenceTransformerEmbeddingModel(
+            SentenceTransformer(
+                provider_model_id,
+                device=config.device,
+                cache_folder=config.cache_dir,
+            )
+        )
+        self._cache.put(provider_model_id, embedding)
+        return embedding
+
+    def _create_rerank(self, model: CatalogModel) -> CrossEncoder:
+        """Create or reuse a cached sentence-transformers cross-encoder.
+
+        Args:
+            model: Catalog entry describing the reranker to instantiate.
+
+        Returns:
+            A cached CrossEncoder for the provider model id.
+
+        """
+        provider_model_id = model.provider_model_id
+        cached = self._cache.get(provider_model_id, "rerank")
+        if cached is not None:
+            return cached
+
+        config = self._entry.config
+        reranker = CrossEncoder(
+            provider_model_id,
+            device=config.device,
+            cache_folder=config.cache_dir,
+        )
+        self._cache.put(provider_model_id, reranker)
+        return reranker
+
+    def _infer_model_type(
+        self,
+        provider_model_id: str,
+        **kwargs: Any,
+    ) -> ModelType:
+        """Infer embedding vs rerank when allowed_models model_type is omitted.
+
+        Args:
+            provider_model_id: Hugging Face model id or local path.
+            **kwargs: Unused; accepted for base-class compatibility.
+
+        Returns:
+            rerank when the id looks like a cross-encoder or reranker,
+            otherwise embedding. Never llm.
+
+        """
+        _ = kwargs
+        normalized = provider_model_id.lower().replace("-", "").replace("_", "")
+        if "crossencoder" in normalized or "rerank" in normalized:
+            return "rerank"
+        return "embedding"

tests/unit/providers/inference/sentence_transformers/README.md

@@ -0,0 +1,13 @@
+# List of source files stored in `tests/unit/providers/inference/sentence_transformers` directory
+
+## [__init__.py](__init__.py)
+
+Unit tests for the sentence-transformers inference provider.
+
+## [test_model_cache.py](test_model_cache.py)
+
+Unit tests for the in-process sentence-transformers model cache.
+
+## [test_provider.py](test_provider.py)
+
+Unit tests for the sentence-transformers inference provider.

tests/unit/providers/inference/sentence_transformers/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for the sentence-transformers inference provider."""

tests/unit/providers/inference/sentence_transformers/test_model_cache.py

@@ -0,0 +1,39 @@
+"""Unit tests for the in-process sentence-transformers model cache."""
+
+from typing import Any
+
+from sentence_transformers import CrossEncoder
+
+from providers.inference.sentence_transformers.model_cache import (
+    SentenceTransformersModelCache,
+)
+
+
+def test_get_and_put_embedding_caches_by_provider_model_id(mocker: Any) -> None:
+    """Stores embeddings by instance type and provider model id."""
+    cache = SentenceTransformersModelCache()
+    embed_a = mocker.Mock(name="embed-a")
+    embed_b = mocker.Mock(name="embed-b")
+
+    assert cache.get("model-a", "embedding") is None
+    cache.put("model-a", embed_a)
+    cache.put("model-b", embed_b)
+
+    assert cache.get("model-a", "embedding") is embed_a
+    assert cache.get("model-b", "embedding") is embed_b
+    assert cache.get("model-a", "rerank") is None
+
+
+def test_get_and_put_rerank_caches_by_provider_model_id(mocker: Any) -> None:
+    """Stores cross-encoders by instance type and provider model id."""
+    cache = SentenceTransformersModelCache()
+    rerank_a = mocker.Mock(name="rerank-a", spec=CrossEncoder)
+    rerank_b = mocker.Mock(name="rerank-b", spec=CrossEncoder)
+
+    assert cache.get("model-a", "rerank") is None
+    cache.put("model-a", rerank_a)
+    cache.put("model-b", rerank_b)
+
+    assert cache.get("model-a", "rerank") is rerank_a
+    assert cache.get("model-b", "rerank") is rerank_b
+    assert cache.get("model-a", "embedding") is None