feat: resolve pgvector vector_size automatically (#588)

Author: mhordynski

docs/how-to/vector_stores/use_pgVector_store.md

@@ -62,7 +62,6 @@ async with asyncpg.create_pool(dsn=DB) as pool:
 
 The connection pool created with asyncpg.create_pool will be used to initialize an instance of PgVectorStore.
 
-
 ```python
 import asyncpg
 from ragbits.core.vector_stores.pgvector import PgVectorStore
@@ -71,12 +70,13 @@ async def main() -> None:
   DB = "postgresql://ragbits_user:ragbits_password@localhost:5432/ragbits_db"
   async with asyncpg.create_pool(dsn=DB) as pool:
     embedder = LiteLLMEmbedder(model="text-embedding-3-small")
-    vector_store = PgVectorStore(embedder=embedder, client=pool, table_name="test_table", vector_size=1536)
+    vector_store = PgVectorStore(embedder=embedder, client=pool, table_name="test_table")
 ```
 
+
 !!! note
-    Ensure that the vector size is correctly configured when initializing PgVectorStore,
-    as it must match the expected dimensions of the stored embeddings.
+    PgVectorStore will automatically determine the vector dimensions from the embedder.
+    If you prefer explicit control or need to override the automatic detection, you can provide the `vector_size` parameter to PgVectorStore initializer.
 
 ## pgVectorStore in Ragbits
 Example:
@@ -91,19 +91,19 @@ async def main() -> None:
   DB = "postgresql://ragbits_user:ragbits_password@localhost:5432/ragbits_db"
   async with asyncpg.create_pool(dsn=DB) as pool:
     embedder = LiteLLMEmbedder(model="text-embedding-3-small")
-    vector_store = PgVectorStore(embedder=embedder, client=pool, table_name="test_table", vector_size=3)
-    data = [VectorStoreEntry(id="test_id_1", key="test_key_1", vector=[0.1, 0.2, 0.3],
+    vector_store = PgVectorStore(embedder=embedder, client=pool, table_name="test_table")
+    data = [VectorStoreEntry(id="test_id_1", text="test text 1",
             metadata={"key1": "value1", "content": "test 1"}),
-            VectorStoreEntry(id="test_id_2", key="test_key_2", vector=[0.4, 0.5, 0.6],
+            VectorStoreEntry(id="test_id_2", text="test text 2",
                               metadata={"key2": "value2", "content": "test 2"})]
 
     await vector_store.store(data)
     all_entries = await vector_store.list()
     print("All entries ", all_entries)
     list_result = await vector_store.list({"content": "test 1"})
     print("Entries with  {content: test 1}", list_result)
-    retrieve_result = await vector_store.retrieve(vector=[0.39, 0.55, 0.6])
-    print("Entries similar to [0.17, 0.23, 0.314] ", retrieve_result)
+    retrieve_result = await vector_store.retrieve("similar test query")
+    print("Entries similar to query", retrieve_result)
     await vector_store.remove(["test_id_1", "test_id_2"])
     after_remove = await vector_store.list()
     print("Entries after remove ", after_remove)

examples/document-search/pgvector.py

@@ -84,7 +84,7 @@ async def main() -> None:
         embedder = LiteLLMEmbedder(
             model_name="text-embedding-3-small",
         )
-        vector_store = PgVectorStore(embedder=embedder, client=pool, table_name="example", vector_size=1536)
+        vector_store = PgVectorStore(embedder=embedder, client=pool, table_name="example")
         document_search = DocumentSearch(
             vector_store=vector_store,
         )

packages/ragbits-core/CHANGELOG.md

@@ -2,7 +2,8 @@
 
 ## Unreleased
 
-- feat: add get_vector_size method to all Embedders (#587)
+- Resolve vector_size by PgVectorStore automatically (#588)
+- Add get_vector_size method to all Embedders (#587)
 
 ## 0.19.1 (2025-05-27)
 

packages/ragbits-core/src/ragbits/core/vector_stores/pgvector.py

@@ -7,7 +7,7 @@
 from pydantic.json import pydantic_encoder
 
 from ragbits.core.audit.traces import trace
-from ragbits.core.embeddings.base import Embedder, SparseVector
+from ragbits.core.embeddings.base import Embedder, SparseVector, VectorSize
 from ragbits.core.embeddings.sparse.base import SparseEmbedder
 from ragbits.core.vector_stores.base import (
     EmbeddingType,
@@ -53,8 +53,8 @@ def __init__(
         self,
         client: asyncpg.Pool,
         table_name: str,
-        vector_size: int,
         embedder: Embedder,
+        vector_size: int | None = None,
         embedding_type: EmbeddingType = EmbeddingType.TEXT,
         distance_method: str | None = None,
         hnsw_params: dict | None = None,
@@ -66,8 +66,8 @@ def __init__(
         Args:
             client: The pgVector database connection pool.
             table_name: The name of the table.
-            vector_size: The size of the vectors.
             embedder: The embedder to use for converting entries to vectors.
+            vector_size: The size of the vectors. If None, will be determined automatically from the embedder.
             embedding_type: Which part of the entry to embed, either text or image. The other part will be ignored.
             distance_method: The distance method to use, default is "cosine" for dense vectors
                 and "sparsevec_l2" for sparse vectors.
@@ -84,7 +84,7 @@ def __init__(
 
         if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", table_name):
             raise ValueError(f"Invalid table name: {table_name}")
-        if not isinstance(vector_size, int) or vector_size <= 0:
+        if vector_size is not None and (not isinstance(vector_size, int) or vector_size <= 0):
             raise ValueError("Vector size must be a positive integer.")
 
         if hnsw_params is None:
@@ -103,6 +103,7 @@ def __init__(
         self._client = client
         self._table_name = table_name
         self._vector_size = vector_size
+        self._vector_size_info: VectorSize | None = None
         self._distance_method = distance_method
         self._hnsw_params = hnsw_params
 
@@ -113,6 +114,32 @@ def __reduce__(self) -> tuple:
         # TODO: To be implemented. Required for Ray processing.
         raise NotImplementedError
 
+    async def _get_vector_size_info(self) -> VectorSize:
+        """
+        Get vector size information from the embedder if not already cached.
+
+        Returns:
+            VectorSize information including size and sparsity.
+        """
+        if self._vector_size_info is None:
+            self._vector_size_info = await self._embedder.get_vector_size()
+            # Update _vector_size for backward compatibility if it wasn't provided
+            if self._vector_size is None:
+                self._vector_size = self._vector_size_info.size
+        return self._vector_size_info
+
+    async def _get_vector_size(self) -> int:
+        """
+        Get the vector size, either from the constructor parameter or from the embedder.
+
+        Returns:
+            The vector size as an integer.
+        """
+        if self._vector_size is not None:
+            return self._vector_size
+        vector_size_info = await self._get_vector_size_info()
+        return vector_size_info.size
+
     def _vector_to_string(self, vector: list[float] | SparseVector) -> str:
         """
         Converts a vector to a string representation.
@@ -124,8 +151,13 @@ def _vector_to_string(self, vector: list[float] | SparseVector) -> str:
             str: The string representation of the vector.
         """
         if isinstance(vector, SparseVector):
+            # For sparse vectors, we need the vector size to be available
+            # This will be resolved when this method is called from async context
+            vector_size = self._vector_size
+            if vector_size is None:
+                raise RuntimeError("Vector size must be determined before converting sparse vectors to string")
             points_str = ",".join(f"{i}:{v}" for i, v in zip(vector.indices, vector.values, strict=False))
-            return f"{{{points_str}}}/{self._vector_size}"
+            return f"{{{points_str}}}/{vector_size}"
         return json.dumps(vector)
 
     @staticmethod
@@ -234,23 +266,25 @@ async def create_table(self) -> None:
         """
         Create a pgVector table with an HNSW index for given similarity.
         """
+        vector_size = await self._get_vector_size()
         with trace(
             table_name=self._table_name,
             distance_method=self._distance_method,
-            vector_size=self._vector_size,
+            vector_size=vector_size,
             hnsw_index_parameters=self._hnsw_params,
         ):
             distance = DISTANCE_OPS[self._distance_method].function_name
             create_vector_extension = "CREATE EXTENSION IF NOT EXISTS vector;"
             # _table_name and has been validated in the class constructor, and it is a valid table name.
-            # _vector_size has been validated in the class constructor, and it is a valid vector size.
+            # vector_size has been validated in the class constructor or obtained from embedder,
+            # and it is a valid vector size.
 
             is_sparse = isinstance(self._embedder, SparseEmbedder)
             vector_func = "VECTOR" if not is_sparse else "SPARSEVEC"
 
             create_table_query = f"""
             CREATE TABLE {self._table_name}
-            (id UUID, text TEXT, image_bytes BYTEA, vector {vector_func}({self._vector_size}), metadata JSONB);
+            (id UUID, text TEXT, image_bytes BYTEA, vector {vector_func}({vector_size}), metadata JSONB);
             """
             # _hnsw_params has been validated in the class constructor, and it is valid dict[str,int].
             create_index_query = f"""
@@ -283,6 +317,10 @@ async def store(self, entries: list[VectorStoreEntry]) -> None:
         """
         if not entries:
             return
+
+        # Ensure vector size is determined before processing
+        vector_size = await self._get_vector_size()
+
         # _table_name has been validated in the class constructor, and it is a valid table name.
         insert_query = f"""
         INSERT INTO {self._table_name} (id, text, image_bytes, vector, metadata)
@@ -291,7 +329,7 @@ async def store(self, entries: list[VectorStoreEntry]) -> None:
         with trace(
             table_name=self._table_name,
             entries=entries,
-            vector_size=self._vector_size,
+            vector_size=vector_size,
             embedder=repr(self._embedder),
             embedding_type=self._embedding_type,
         ):
@@ -359,11 +397,14 @@ async def retrieve(
         """
         merged_options = (self.default_options | options) if options else self.default_options
 
+        # Ensure vector size is determined before processing
+        vector_size = await self._get_vector_size()
+
         with trace(
             text=text,
             options=merged_options.dict(),
             table_name=self._table_name,
-            vector_size=self._vector_size,
+            vector_size=vector_size,
             distance_method=self._distance_method,
             embedder=repr(self._embedder),
             embedding_type=self._embedding_type,

packages/ragbits-core/tests/unit/vector_stores/test_pgvector.py

@@ -358,21 +358,39 @@ async def test_list_with_nested_where_clause(
 async def test_retrieve_with_where_clause_and_score_threshold(
     mock_pgvector_store: PgVectorStore, mock_db_pool: tuple[MagicMock, AsyncMock]
 ) -> None:
-    data = DATA_JSON_EXAMPLE
-    query = "SQL RETRIEVE QUERY"
     _, mock_conn = mock_db_pool
-    where_clause: dict[str, str | int | float | bool | dict[Any, Any]] = {"key1": "value1"}
+    mock_conn.fetch.return_value = []
+    where = cast(WhereQuery, {"id": "test_id"})
+    await mock_pgvector_store.retrieve(text="test", options=VectorStoreOptions(score_threshold=0.1, where=where))
+    mock_conn.fetch.assert_called_once()
+    calls = mock_conn.fetch.mock_calls
+    assert calls[0].args[0].startswith("SELECT *, vector <=> $1 as distance, 1 - (vector <=> $1) as score FROM")
+    assert calls[0].args[0].endswith("WHERE score >= $2 AND metadata @> $3 ORDER BY distance LIMIT $4;")
 
-    with patch.object(mock_pgvector_store, "_create_retrieve_query") as mock_create_retrieve_query:
-        mock_conn.fetch = AsyncMock(return_value=data)
-        mock_create_retrieve_query.return_value = (query, [["[0.1, 0.2, 0.3]", 0.1, 1]])
-        results = await mock_pgvector_store.retrieve(
-            text="some_text", options=VectorStoreOptions(where=where_clause, score_threshold=0.5)
-        )
-        mock_create_retrieve_query.assert_called_once()
-        mock_conn.fetch.assert_called_once()
-        assert len(results) == 2
-        assert isinstance(results[0], VectorStoreResult)
-        assert isinstance(results[1], VectorStoreResult)
-        assert results[0].entry.id == UUID("8c7d6b27-4ef1-537c-ad7c-676edb8bc8a8")
-        assert results[1].entry.id == UUID("9c7d6b27-4ef1-537c-ad7c-676edb8bc8a8")
+
+@pytest.mark.asyncio
+async def test_auto_vector_size_determination(mock_db_pool: tuple[MagicMock, AsyncMock]) -> None:
+    """Test that PgVectorStore can determine vector size automatically from embedder."""
+    mock_pool, _mock_conn = mock_db_pool
+    mock_embedder = AsyncMock()
+
+    # Mock the get_vector_size method to return a VectorSize with size 5
+    from ragbits.core.embeddings.base import VectorSize
+
+    mock_embedder.get_vector_size.return_value = VectorSize(size=5, is_sparse=False)
+
+    # Create PgVectorStore without providing vector_size
+    store = PgVectorStore(client=mock_pool, table_name=TEST_TABLE_NAME, embedder=mock_embedder)
+
+    # The vector size should be None initially
+    assert store._vector_size is None
+
+    # When we call _get_vector_size(), it should determine the size from embedder
+    vector_size = await store._get_vector_size()
+    assert vector_size == 5
+
+    # Now _vector_size should be cached
+    assert store._vector_size == 5
+
+    # Verify the embedder's get_vector_size was called
+    mock_embedder.get_vector_size.assert_called_once()