feat: add get_vector_size method to embedder interface (#587)

Author: mhordynski

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- feat: add get_vector_size method to all Embedders (#587)
+
 ## 0.19.1 (2025-05-27)
 
 ## 0.19.0 (2025-05-27)

packages/ragbits-core/src/ragbits/core/embeddings/__init__.py

@@ -1,4 +1,4 @@
-from .base import Embedder, EmbedderOptionsT, SparseVector
+from .base import Embedder, EmbedderOptionsT, SparseVector, VectorSize
 from .dense import DenseEmbedder, LiteLLMEmbedder, NoopEmbedder
 from .sparse import BagOfTokens, BagOfTokensOptions, SparseEmbedder, SparseEmbedderOptionsT
 
@@ -13,4 +13,5 @@
     "SparseEmbedder",
     "SparseEmbedderOptionsT",
     "SparseVector",
+    "VectorSize",
 ]

packages/ragbits-core/src/ragbits/core/embeddings/base.py

@@ -24,6 +24,16 @@ def __repr__(self) -> str:
         return f"SparseVector(indices={self.indices}, values={self.values})"
 
 
+class VectorSize(BaseModel):
+    """Information about vector dimensions returned by an embedder"""
+
+    size: int
+    """The size/dimension of the vector"""
+
+    is_sparse: bool = False
+    """Whether this represents a sparse vector (where size is vocabulary size) or dense vector"""
+
+
 class Embedder(ConfigurableComponent[EmbedderOptionsT], ABC):
     """
     Abstract class that defines a common interface for both sparse and dense embedding models.
@@ -48,6 +58,15 @@ async def embed_text(
             List of embeddings for the given strings.
         """
 
+    @abstractmethod
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get information about the vector size/dimensions returned by this embedder.
+
+        Returns:
+            VectorSize object containing dimension information and whether vectors are sparse.
+        """
+
     def image_support(self) -> bool:  # noqa: PLR6301
         """
         Check if the model supports image embeddings.

packages/ragbits-core/src/ragbits/core/embeddings/dense/base.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-from ragbits.core.embeddings.base import Embedder, EmbedderOptionsT
+from ragbits.core.embeddings.base import Embedder, EmbedderOptionsT, VectorSize
 
 
 class DenseEmbedder(Embedder[EmbedderOptionsT], ABC):  # noqa: F821
@@ -21,6 +21,15 @@ async def embed_text(self, data: list[str], options: EmbedderOptionsT | None = N
             List of embeddings for the given strings.
         """
 
+    @abstractmethod
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get information about the dense vector size/dimensions returned by this embedder.
+
+        Returns:
+            VectorSize object with is_sparse=False and the embedding dimension.
+        """
+
     async def embed_image(self, images: list[bytes], options: EmbedderOptionsT | None = None) -> list[list[float]]:
         """
         Creates embeddings for the given images.

packages/ragbits-core/src/ragbits/core/embeddings/dense/fastembed.py

@@ -3,6 +3,7 @@
 from fastembed import TextEmbedding
 
 from ragbits.core.audit.traces import trace
+from ragbits.core.embeddings.base import VectorSize
 from ragbits.core.embeddings.dense.base import DenseEmbedder, EmbedderOptionsT
 from ragbits.core.options import Options
 
@@ -43,6 +44,26 @@ def __reduce__(self) -> tuple[Callable, tuple[str, bool, FastEmbedOptions | None
         """
         return (self.__class__, (self.model_name, self.use_gpu, self.default_options))
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this FastEmbed model.
+
+        Returns:
+            VectorSize object with the model's embedding dimension.
+        """
+        # Get model info from FastEmbed's supported models list
+        supported_models = self._model.list_supported_models()
+        model_info = next((model for model in supported_models if model["model"] == self.model_name), None)
+
+        if model_info and "dim" in model_info:
+            vector_size = model_info["dim"]
+        else:
+            # Fallback to the original method if metadata is not available
+            sample_embedding = await self.embed_text(["sample"])
+            vector_size = len(sample_embedding[0])
+
+        return VectorSize(size=vector_size, is_sparse=False)
+
     async def embed_text(self, data: list[str], options: EmbedderOptionsT | None = None) -> list[list[float]]:
         """
         Embeds a list of strings into a list of embeddings.

packages/ragbits-core/src/ragbits/core/embeddings/dense/litellm.py

@@ -1,9 +1,10 @@
-from typing import Any
+from typing import Any, cast
 
 import litellm
 from typing_extensions import Self
 
 from ragbits.core.audit.traces import trace
+from ragbits.core.embeddings.base import VectorSize
 from ragbits.core.embeddings.dense.base import DenseEmbedder
 from ragbits.core.embeddings.exceptions import (
     EmbeddingConnectionError,
@@ -68,6 +69,29 @@ def __init__(
         self.api_version = api_version
         self.router = router
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this LiteLLM model.
+
+        If dimensions are specified in default options, use that value.
+        Otherwise, embed a sample text to determine the dimension.
+
+        Returns:
+            VectorSize object with the model's embedding dimension.
+        """
+        # Check if dimensions are explicitly set in default options
+        if (
+            self.default_options
+            and self.default_options.dimensions is not NOT_GIVEN
+            and self.default_options.dimensions is not None
+        ):
+            # We've checked that dimensions is not None and not NOT_GIVEN, so it must be int
+            return VectorSize(size=cast(int, self.default_options.dimensions), is_sparse=False)
+
+        # If no dimensions specified, embed a sample text to determine size
+        sample_embedding = await self.embed_text(["sample"])
+        return VectorSize(size=len(sample_embedding[0]), is_sparse=False)
+
     async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | None = None) -> list[list[float]]:
         """
         Creates embeddings for the given strings.

packages/ragbits-core/src/ragbits/core/embeddings/dense/local.py

@@ -2,6 +2,7 @@
 from typing import Any
 
 from ragbits.core.audit.traces import trace
+from ragbits.core.embeddings.base import VectorSize
 from ragbits.core.embeddings.dense.base import DenseEmbedder
 from ragbits.core.options import Options
 
@@ -55,6 +56,19 @@ def __init__(
         self.model_name = model_name
         self.model = SentenceTransformer(self.model_name, **model_kwargs)
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this local SentenceTransformer model.
+
+        Returns:
+            VectorSize object with the model's embedding dimension.
+        """
+        dimension = self.model.get_sentence_embedding_dimension()
+        if dimension is None:
+            sample_embedding = await self.embed_text(["sample"])
+            dimension = len(sample_embedding[0])
+        return VectorSize(size=dimension, is_sparse=False)
+
     async def embed_text(self, data: list[str], options: LocalEmbedderOptions | None = None) -> list[list[float]]:
         """
         Calls the appropriate encoder endpoint with the given data and options.

packages/ragbits-core/src/ragbits/core/embeddings/dense/noop.py

@@ -1,4 +1,5 @@
 from ragbits.core.audit.traces import trace
+from ragbits.core.embeddings.base import VectorSize
 from ragbits.core.embeddings.dense.base import DenseEmbedder
 from ragbits.core.options import Options, OptionsT
 
@@ -37,6 +38,15 @@ def __init__(
         self.return_cycle = 0
         self.image_return_cycle = 0
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this NoopEmbedder.
+
+        Returns:
+            VectorSize object with the dimension of the first embedding vector.
+        """
+        return VectorSize(size=len(self.return_values[0][0]), is_sparse=False)
+
     async def embed_text(self, data: list[str], options: Options | None = None) -> list[list[float]]:  # noqa: PLR6301
         """
         Embeds a list of strings into a list of vectors.

packages/ragbits-core/src/ragbits/core/embeddings/dense/vertex_multimodal.py

@@ -1,6 +1,7 @@
 import asyncio
 import base64
 
+from ragbits.core.embeddings.base import VectorSize
 from ragbits.core.embeddings.dense.litellm import LiteLLMEmbedderOptions
 
 try:
@@ -69,6 +70,18 @@ def __init__(
         if model_name not in supported_models:
             raise ValueError(f"Model {model_name} is not supported by VertexAI multimodal embeddings")
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this VertexAI multimodal model.
+
+        Embeds a sample text to determine the dimension.
+
+        Returns:
+            VectorSize object with the model's embedding dimension.
+        """
+        sample_embedding = await self.embed_text(["sample"])
+        return VectorSize(size=len(sample_embedding[0]), is_sparse=False)
+
     async def _embed(self, data: list[dict], options: LiteLLMEmbedderOptions | None = None) -> list[dict]:
         """
         Creates embeddings for the given data. The format is defined in the VertexAI API:

packages/ragbits-core/src/ragbits/core/embeddings/sparse/bag_of_tokens.py

@@ -3,7 +3,7 @@
 import tiktoken
 
 from ragbits.core.audit.traces import trace
-from ragbits.core.embeddings.base import SparseVector
+from ragbits.core.embeddings.base import SparseVector, VectorSize
 from ragbits.core.embeddings.sparse.base import SparseEmbedder
 from ragbits.core.options import Options
 from ragbits.core.types import NOT_GIVEN, NotGiven
@@ -22,6 +22,33 @@ class BagOfTokens(SparseEmbedder[BagOfTokensOptions]):
 
     options_cls = BagOfTokensOptions
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this BagOfTokens model.
+
+        For BagOfTokens, this returns the tokenizer vocabulary size.
+
+        Returns:
+            VectorSize object with is_sparse=True and the vocabulary size.
+        """
+        merged_options = self.default_options
+
+        if merged_options.encoding_name and merged_options.model_name:
+            raise ValueError("Please specify only one of encoding_name or model_name")
+        if not (merged_options.encoding_name or merged_options.model_name):
+            raise ValueError("Either encoding_name or model_name needs to be specified")
+
+        if merged_options.encoding_name:
+            encoder = tiktoken.get_encoding(encoding_name=merged_options.encoding_name)
+        elif merged_options.model_name:
+            encoder = tiktoken.encoding_for_model(model_name=merged_options.model_name)
+        else:
+            raise ValueError("Either encoding_name or model_name needs to be specified")
+
+        # Get the vocabulary size from the encoder
+        vocab_size = encoder.n_vocab
+        return VectorSize(size=vocab_size, is_sparse=True)
+
     async def embed_text(self, texts: list[str], options: BagOfTokensOptions | None = None) -> list[SparseVector]:
         """
         Transforms a list of texts into sparse vectors using bag-of-tokens representation.

packages/ragbits-core/src/ragbits/core/embeddings/sparse/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import TypeVar
 
-from ragbits.core.embeddings.base import Embedder, SparseVector
+from ragbits.core.embeddings.base import Embedder, SparseVector, VectorSize
 from ragbits.core.options import Options
 
 SparseEmbedderOptionsT = TypeVar("SparseEmbedderOptionsT", bound=Options)
@@ -23,6 +23,15 @@ async def embed_text(self, texts: list[str], options: SparseEmbedderOptionsT | N
             list of sparse embeddings.
         """
 
+    @abstractmethod
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get information about the sparse vector size/dimensions returned by this embedder.
+
+        Returns:
+            VectorSize object with is_sparse=True and the vocabulary size.
+        """
+
     async def embed_image(
         self, images: list[bytes], options: SparseEmbedderOptionsT | None = None
     ) -> list[SparseVector]:

packages/ragbits-core/src/ragbits/core/embeddings/sparse/fastembed.py

@@ -3,7 +3,7 @@
 from fastembed import SparseTextEmbedding
 
 from ragbits.core.audit.traces import trace
-from ragbits.core.embeddings.base import EmbedderOptionsT, SparseVector
+from ragbits.core.embeddings.base import EmbedderOptionsT, SparseVector, VectorSize
 from ragbits.core.embeddings.dense.fastembed import FastEmbedOptions
 from ragbits.core.embeddings.sparse.base import SparseEmbedder
 
@@ -35,6 +35,29 @@ def __reduce__(self) -> tuple[Callable, tuple[str, bool, FastEmbedOptions | None
         """
         return (self.__class__, (self.model_name, self.use_gpu, self.default_options))
 
+    async def get_vector_size(self) -> VectorSize:
+        """
+        Get the vector size for this FastEmbed sparse model.
+
+        For sparse models, this returns the vocabulary size.
+
+        Returns:
+            VectorSize object with is_sparse=True and the vocabulary size.
+        """
+        # Get model info from FastEmbed's supported models list
+        supported_models = self._model.list_supported_models()
+        model_info = next((model for model in supported_models if model["model"] == self.model_name), None)
+
+        if model_info and "vocab_size" in model_info:
+            vocab_size = model_info["vocab_size"]
+        else:
+            sample_embedding = await self.embed_text(["sample text with various tokens"])
+            vocab_size = (
+                max(sample_embedding[0].indices) + 1 if sample_embedding and sample_embedding[0].indices else 30000
+            )
+
+        return VectorSize(size=vocab_size, is_sparse=True)
+
     async def embed_text(self, data: list[str], options: EmbedderOptionsT | None = None) -> list[SparseVector]:
         """
         Embeds a list of strings into a list of sparse embeddings.