feat: integration tests for pgvector (#495)

Author: kzamlynska

.github/workflows/ci.yml

@@ -155,6 +155,21 @@ jobs:
       matrix:
         python-version: ["3.10", "3.11", "3.12"]
 
+    services:
+      postgres:
+        image: postgres:16
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: test_db
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U postgres"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
     steps:
       - uses: actions/checkout@v4
 
@@ -182,7 +197,14 @@ jobs:
           path: ~/nltk_data
           key: nltk-${{ runner.os }}
 
+      - name: Install PostgreSQL + pgvector
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y postgresql-16-pgvector
+
       - name: Run Tests With Coverage
+        env:
+          DATABASE_URL: postgresql://postgres:postgres@localhost:5432/test_db
         run: |
           # run with coverage to not execute tests twice
           uv run coverage run -m pytest -v -p no:warnings --junitxml=report.xml

.libraries-whitelist.txt

@@ -5,4 +5,7 @@ chroma-hnswlib
 rouge
 distilabel
 rerankers
-py_rust_stemmers
+py_rust_stemmers
+mirakuru
+psycopg
+pytest-postgresql

examples/document-search/chroma.py

@@ -73,7 +73,7 @@ async def main() -> None:
     Run the example.
     """
     embedder = LiteLLMEmbedder(
-        model="text-embedding-3-small",
+        model_name="text-embedding-3-small",
         default_options=LiteLLMEmbedderOptions(
             dimensions=1024,
             timeout=1000,

examples/document-search/distributed.py

@@ -84,7 +84,7 @@ async def main() -> None:
     Run the example.
     """
     embedder = LiteLLMEmbedder(
-        model="text-embedding-3-small",
+        model_name="text-embedding-3-small",
     )
     vector_store = QdrantVectorStore(
         client=AsyncQdrantClient(

examples/document-search/otel.py

@@ -99,7 +99,7 @@ async def main() -> None:
     Run the example.
     """
     embedder = LiteLLMEmbedder(
-        model="text-embedding-3-small",
+        model_name="text-embedding-3-small",
     )
     vector_store = ChromaVectorStore(
         client=EphemeralClient(),

examples/document-search/pgvector.py

@@ -82,7 +82,7 @@ async def main() -> None:
     database_url = "postgresql://ragbits_example:ragbits_example@localhost/ragbits_example"
     async with asyncpg.create_pool(dsn=database_url) as pool:
         embedder = LiteLLMEmbedder(
-            model="text-embedding-3-small",
+            model_name="text-embedding-3-small",
         )
         vector_store = PgVectorStore(embedder=embedder, client=pool, table_name="example", vector_size=1536)
         document_search = DocumentSearch(

examples/document-search/qdrant.py

@@ -72,7 +72,7 @@ async def main() -> None:
     Run the example.
     """
     embedder = LiteLLMEmbedder(
-        model="text-embedding-3-small",
+        model_name="text-embedding-3-small",
     )
     vector_store = QdrantVectorStore(
         client=AsyncQdrantClient(location=":memory:"),

packages/ragbits-core/CHANGELOG.md

@@ -1,6 +1,9 @@
 # CHANGELOG
 
 ## Unreleased
+
+- Image embeddings in PgVectorStore (#495)
+- Add PgVectorStore to vector store integration tests (#495)
 - Add new fusion strategies for the hybrid vector store: RRF and DBSF (#413)
 - move sources from ragbits-document-search to ragbits-core (#496)
 - adding connection check to Azure get_blob_service (#502)

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

@@ -178,6 +178,15 @@ def _create_list_query(
         ]
         return query, values
 
+    async def _check_table_exists(self) -> bool:
+        check_table_existence = """
+                SELECT EXISTS (
+                SELECT FROM information_schema.tables
+                WHERE table_name = $1
+            ); """
+        async with self._client.acquire() as conn:
+            return await conn.fetchval(check_table_existence, self._table_name)
+
     async def create_table(self) -> None:
         """
         Create a pgVector table with an HNSW index for given similarity.
@@ -188,43 +197,36 @@ async def create_table(self) -> None:
             vector_size=self._vector_size,
             hnsw_index_parameters=self._hnsw_params,
         ):
-            check_table_existence = """
-                    SELECT EXISTS (
-                    SELECT FROM information_schema.tables
-                    WHERE table_name = $1
-                ); """
             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.
 
             create_table_query = f"""
             CREATE TABLE {self._table_name}
-            (id UUID, key TEXT, vector VECTOR({self._vector_size}), metadata JSONB);
+            (id UUID, text TEXT, image_bytes BYTEA, vector VECTOR({self._vector_size}), metadata JSONB);
             """
             # _hnsw_params has been validated in the class constructor, and it is valid dict[str,int].
             create_index_query = f"""
                     CREATE INDEX {self._table_name + "_hnsw_idx"} ON {self._table_name}
                     USING hnsw (vector {distance})
                     WITH (m = {self._hnsw_params["m"]}, ef_construction = {self._hnsw_params["ef_construction"]});
                     """
-
+            if await self._check_table_exists():
+                print(f"Table {self._table_name} already exist!")
+                return
             async with self._client.acquire() as conn:
                 await conn.execute(create_vector_extension)
-                exists = await conn.fetchval(check_table_existence, self._table_name)
 
-                if not exists:
-                    try:
-                        async with conn.transaction():
-                            await conn.execute(create_table_query)
-                            await conn.execute(create_index_query)
+                try:
+                    async with conn.transaction():
+                        await conn.execute(create_table_query)
+                        await conn.execute(create_index_query)
 
-                        print("Table and index created!")
-                    except Exception as e:
-                        print(f"Failed to create table and index: {e}")
-                        raise
-                else:
-                    print("Table already exists!")
+                    print("Table and index created!")
+                except Exception as e:
+                    print(f"Failed to create table and index: {e}")
+                    raise
 
     async def store(self, entries: list[VectorStoreEntry]) -> None:
         """
@@ -237,8 +239,8 @@ async def store(self, entries: list[VectorStoreEntry]) -> None:
             return
         # _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, key, vector, metadata)
-        VALUES ($1, $2, $3, $4)
+        INSERT INTO {self._table_name} (id, text, image_bytes, vector, metadata)
+        VALUES ($1, $2, $3, $4, $5)
         """  # noqa S608
         with trace(
             table_name=self._table_name,
@@ -248,30 +250,28 @@ async def store(self, entries: list[VectorStoreEntry]) -> None:
             embedding_type=self._embedding_type,
         ):
             embeddings = await self._create_embeddings(entries)
-
-            try:
-                async with self._client.acquire() as conn:
-                    for entry in entries:
-                        if entry.id not in embeddings:
-                            continue
-
-                        await conn.execute(
-                            insert_query,
-                            str(entry.id),
-                            entry.text,
-                            str(embeddings[entry.id]),
-                            json.dumps(entry.metadata, default=pydantic_encoder),
-                        )
-            except asyncpg.exceptions.UndefinedTableError:
+            exists = await self._check_table_exists()
+            if not exists:
                 print(f"Table {self._table_name} does not exist. Creating the table.")
                 try:
                     await self.create_table()
                 except Exception as e:
                     print(f"Failed to handle missing table: {e}")
                     return
 
-                print("Table created successfully. Inserting entries...")
-                await self.store(entries)
+            async with self._client.acquire() as conn:
+                for entry in entries:
+                    if entry.id not in embeddings:
+                        continue
+
+                    await conn.execute(
+                        insert_query,
+                        str(entry.id),
+                        entry.text,
+                        entry.image_bytes,
+                        str(embeddings[entry.id]),
+                        json.dumps(entry.metadata, default=pydantic_encoder),
+                    )
 
     async def remove(self, ids: list[UUID]) -> None:
         """
@@ -296,34 +296,6 @@ async def remove(self, ids: list[UUID]) -> None:
                 print(f"Table {self._table_name} does not exist.")
                 return
 
-    async def _fetch_records(self, query: str, values: list[Any]) -> list[VectorStoreEntry]:
-        """
-        Fetch records from the pgVector collection.
-
-        Args:
-            query: sql query
-            values: list of values to be used in the query.
-
-        Returns:
-            list of VectorStoreEntry objects.
-        """
-        try:
-            async with self._client.acquire() as conn:
-                results = await conn.fetch(query, *values)
-
-            return [
-                VectorStoreEntry(
-                    id=record["id"],
-                    text=record["key"],
-                    metadata=json.loads(record["metadata"]),
-                )
-                for record in results
-            ]
-
-        except asyncpg.exceptions.UndefinedTableError:
-            print(f"Table {self._table_name} does not exist.")
-            return []
-
     async def retrieve(
         self,
         text: str,
@@ -362,7 +334,8 @@ async def retrieve(
                     VectorStoreResult(
                         entry=VectorStoreEntry(
                             id=record["id"],
-                            text=record["key"],
+                            text=record["text"],
+                            image_bytes=record["image_bytes"],
                             metadata=json.loads(record["metadata"]),
                         ),
                         vector=json.loads(record["vector"]),
@@ -393,5 +366,20 @@ async def list(
         """
         with trace(table=self._table_name, query=where, limit=limit, offset=offset) as outputs:
             list_query, values = self._create_list_query(where, limit, offset)
-            outputs.listed_entries = await self._fetch_records(list_query, values)
+            try:
+                async with self._client.acquire() as conn:
+                    results = await conn.fetch(list_query, *values)
+                outputs.listed_entries = [
+                    VectorStoreEntry(
+                        id=record["id"],
+                        text=record["text"],
+                        image_bytes=record["image_bytes"],
+                        metadata=json.loads(record["metadata"]),
+                    )
+                    for record in results
+                ]
+
+            except asyncpg.exceptions.UndefinedTableError:
+                print(f"Table {self._table_name} does not exist.")
+                outputs.listed_entries = []
             return outputs.listed_entries

packages/ragbits-core/tests/integration/vector_stores/test_vector_store.py

@@ -3,8 +3,10 @@
 from typing import cast
 from uuid import UUID
 
+import asyncpg
 import pytest
 from chromadb import EphemeralClient
+from psycopg import Connection
 from qdrant_client import AsyncQdrantClient
 
 from ragbits.core.embeddings.noop import NoopEmbedder
@@ -18,6 +20,7 @@
 )
 from ragbits.core.vector_stores.chroma import ChromaVectorStore
 from ragbits.core.vector_stores.in_memory import InMemoryVectorStore
+from ragbits.core.vector_stores.pgvector import PgVectorStore
 from ragbits.core.vector_stores.qdrant import QdrantVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
@@ -34,22 +37,31 @@
 IMAGES_PATH = Path(__file__).parent.parent.parent / "assets" / "img"
 
 
-# TODO: Add PgVectorStore
+@pytest.fixture
+async def pgvector_test_db(postgresql: Connection) -> asyncpg.pool:
+    dsn = f"postgresql://{postgresql.info.user}:{postgresql.info.password}@{postgresql.info.host}:{postgresql.info.port}/{postgresql.info.dbname}"
+    async with asyncpg.create_pool(dsn) as pool:
+        yield pool
+
+
 @pytest.fixture(
     name="vector_store_cls",
     params=[
-        lambda: partial(InMemoryVectorStore),
-        lambda: partial(ChromaVectorStore, client=EphemeralClient(), index_name="test_index_name"),
-        lambda: partial(QdrantVectorStore, client=AsyncQdrantClient(":memory:"), index_name="test_index_name"),
+        lambda _: partial(InMemoryVectorStore),
+        lambda _: partial(ChromaVectorStore, client=EphemeralClient(), index_name="test_index_name"),
+        lambda _: partial(QdrantVectorStore, client=AsyncQdrantClient(":memory:"), index_name="test_index_name"),
+        lambda pg_pool: partial(PgVectorStore, client=pg_pool, table_name="test_index_name", vector_size=3),
     ],
-    ids=["InMemoryVectorStore", "ChromaVectorStore", "QdrantVectorStore"],
+    ids=["InMemoryVectorStore", "ChromaVectorStore", "QdrantVectorStore", "PgVectorStore"],
 )
-def vector_store_cls_fixture(request: pytest.FixtureRequest) -> type[VectorStoreWithExternalEmbedder]:
+def vector_store_cls_fixture(
+    request: pytest.FixtureRequest, pgvector_test_db: asyncpg.pool
+) -> type[VectorStoreWithExternalEmbedder]:
     """
     Returns vector stores classes with different backends, with backend-specific parameters already set,
     but parameters common to VectorStoreWithExternalEmbedder left to be set.
     """
-    return request.param()
+    return request.param(pgvector_test_db)
 
 
 @pytest.fixture(name="vector_store", params=[EmbeddingType.TEXT, EmbeddingType.IMAGE], ids=["Text", "Image"])