Python: Add vector search to Postgres connector (microsoft#10213)

### Motivation and Context

Following up on microsoft#8951, this PR adds an implementation of
`VectorSearchBase` to `PostgresCollection`. This implementation provides
vectorized search and does not implement text search or vectorizable
text search.

Unit and integration tests are added, and the
`python/samples/getting_started/third_party/postgres-memory.ipynb`
notebook was expanded to include vector search in the example.

### Contribution Checklist

- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄

Author: lossyrob

python/samples/getting_started/third_party/postgres-memory.ipynb

@@ -5,6 +5,10 @@
 # Limitation based on pgvector documentation https://github.com/pgvector/pgvector#what-if-i-want-to-index-vectors-with-more-than-2000-dimensions
 MAX_DIMENSIONALITY = 2000
 
+# The name of the column that returns distance value in the database.
+# It is used in the similarity search query. Must not conflict with model property.
+DISTANCE_COLUMN_NAME = "sk_pg_distance"
+
 # Environment Variables
 PGHOST_ENV_VAR = "PGHOST"
 PGPORT_ENV_VAR = "PGPORT"

python/semantic_kernel/connectors/memory/postgres/constants.py

@@ -4,6 +4,7 @@
 
 from psycopg.conninfo import conninfo_to_dict
 from psycopg_pool import AsyncConnectionPool
+from psycopg_pool.abc import ACT
 from pydantic import Field, SecretStr
 
 from semantic_kernel.connectors.memory.postgres.constants import (
@@ -14,10 +15,7 @@
     PGSSL_MODE_ENV_VAR,
     PGUSER_ENV_VAR,
 )
-from semantic_kernel.exceptions.memory_connector_exceptions import (
-    MemoryConnectorConnectionException,
-    MemoryConnectorInitializationError,
-)
+from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorConnectionException
 from semantic_kernel.kernel_pydantic import KernelBaseSettings
 from semantic_kernel.utils.experimental_decorator import experimental_class
 
@@ -89,30 +87,34 @@ def get_connection_args(self) -> dict[str, Any]:
         if self.password:
             result["password"] = self.password.get_secret_value()
 
-        # Ensure required values
-        if "host" not in result:
-            raise MemoryConnectorInitializationError("host is required. Please set PGHOST or connection_string.")
-        if "dbname" not in result:
-            raise MemoryConnectorInitializationError(
-                "database is required. Please set PGDATABASE or connection_string."
-            )
-        if "user" not in result:
-            raise MemoryConnectorInitializationError("user is required. Please set PGUSER or connection_string.")
-        if "password" not in result:
-            raise MemoryConnectorInitializationError(
-                "password is required. Please set PGPASSWORD or connection_string."
-            )
-
         return result
 
-    async def create_connection_pool(self) -> AsyncConnectionPool:
-        """Creates a connection pool based off of settings."""
+    async def create_connection_pool(
+        self, connection_class: type[ACT] | None = None, **kwargs: Any
+    ) -> AsyncConnectionPool:
+        """Creates a connection pool based off of settings.
+
+        Args:
+            connection_class: The connection class to use.
+            kwargs: Additional keyword arguments to pass to the connection class.
+
+        Returns:
+            The connection pool.
+        """
         try:
+            # Only pass connection_class if it specified, or else allow psycopg to use the default connection class
+            extra_args: dict[str, Any] = {} if connection_class is None else {"connection_class": connection_class}
+
             pool = AsyncConnectionPool(
                 min_size=self.min_pool,
                 max_size=self.max_pool,
                 open=False,
-                kwargs=self.get_connection_args(),
+                # kwargs are passed to the connection class
+                kwargs={
+                    **self.get_connection_args(),
+                    **kwargs,
+                },
+                **extra_args,
             )
             await pool.open()
         except Exception as e:

python/semantic_kernel/connectors/memory/postgres/postgres_collection.py

@@ -4,21 +4,20 @@
 import sys
 from typing import Any, TypeVar
 
-if sys.version_info >= (3, 12):
-    from typing import override  # pragma: no cover
-else:
-    from typing_extensions import override  # pragma: no cover
-
 from psycopg import sql
 from psycopg_pool import AsyncConnectionPool
 
 from semantic_kernel.connectors.memory.postgres.postgres_collection import PostgresCollection
 from semantic_kernel.connectors.memory.postgres.postgres_memory_store import DEFAULT_SCHEMA
-from semantic_kernel.data.record_definition.vector_store_model_definition import VectorStoreRecordDefinition
-from semantic_kernel.data.vector_storage.vector_store import VectorStore
-from semantic_kernel.data.vector_storage.vector_store_record_collection import VectorStoreRecordCollection
+from semantic_kernel.data import VectorStore, VectorStoreRecordCollection, VectorStoreRecordDefinition
 from semantic_kernel.utils.experimental_decorator import experimental_class
 
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
+
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 TModel = TypeVar("TModel")

python/semantic_kernel/connectors/memory/postgres/postgres_settings.py

@@ -52,7 +52,9 @@ def python_type_to_postgres(python_type_str: str) -> str | None:
     return None
 
 
-def convert_row_to_dict(row: tuple[Any, ...], fields: list[tuple[str, VectorStoreRecordField]]) -> dict[str, Any]:
+def convert_row_to_dict(
+    row: tuple[Any, ...], fields: list[tuple[str, VectorStoreRecordField | None]]
+) -> dict[str, Any]:
     """Convert a row from a PostgreSQL query to a dictionary.
 
     Uses the field information to map the row values to the corresponding field names.
@@ -65,11 +67,12 @@ def convert_row_to_dict(row: tuple[Any, ...], fields: list[tuple[str, VectorStor
         A dictionary representation of the row.
     """
 
-    def _convert(v: Any | None, field: VectorStoreRecordField) -> Any | None:
+    def _convert(v: Any | None, field: VectorStoreRecordField | None) -> Any | None:
         if v is None:
             return None
-        if isinstance(field, VectorStoreRecordVectorField):
-            # psycopg returns vector as a string
+        if isinstance(field, VectorStoreRecordVectorField) and isinstance(v, str):
+            # psycopg returns vector as a string if pgvector is not loaded.
+            # If pgvector is registered with the connection, no conversion is required.
             return json.loads(v)
         return v
 
@@ -109,6 +112,8 @@ def get_vector_index_ops_str(distance_function: DistanceFunction) -> str:
         >>> get_vector_index_ops_str(DistanceFunction.COSINE)
         'vector_cosine_ops'
     """
+    if distance_function == DistanceFunction.COSINE_DISTANCE:
+        return "vector_cosine_ops"
     if distance_function == DistanceFunction.COSINE_SIMILARITY:
         return "vector_cosine_ops"
     if distance_function == DistanceFunction.DOT_PROD:
@@ -121,6 +126,38 @@ def get_vector_index_ops_str(distance_function: DistanceFunction) -> str:
     raise ValueError(f"Unsupported distance function: {distance_function}")
 
 
+def get_vector_distance_ops_str(distance_function: DistanceFunction) -> str:
+    """Get the PostgreSQL distance operator string for a given distance function.
+
+    Args:
+        distance_function: The distance function for which the operator string is needed.
+
+    Note:
+        For the COSINE_SIMILARITY and DOT_PROD distance functions,
+        there is additional query steps to retrieve the correct distance.
+        For dot product, take -1 * inner product, as <#> returns the negative inner product
+        since Postgres only supports ASC order index scans on operators
+        For cosine similarity, take 1 - cosine distance.
+
+    Returns:
+        The PostgreSQL distance operator string for the given distance function.
+
+    Raises:
+        ValueError: If the distance function is unsupported.
+    """
+    if distance_function == DistanceFunction.COSINE_DISTANCE:
+        return "<=>"
+    if distance_function == DistanceFunction.COSINE_SIMILARITY:
+        return "<=>"
+    if distance_function == DistanceFunction.DOT_PROD:
+        return "<#>"
+    if distance_function == DistanceFunction.EUCLIDEAN_DISTANCE:
+        return "<->"
+    if distance_function == DistanceFunction.MANHATTAN:
+        return "<+>"
+    raise ValueError(f"Unsupported distance function: {distance_function}")
+
+
 async def ensure_open(connection_pool: AsyncConnectionPool) -> AsyncConnectionPool:
     """Ensure the connection pool is open.
 

python/semantic_kernel/connectors/memory/postgres/postgres_store.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import uuid
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from contextlib import asynccontextmanager
 from typing import Annotated, Any
 
@@ -11,6 +11,7 @@
 from pydantic import BaseModel
 
 from semantic_kernel.connectors.memory.postgres import PostgresSettings, PostgresStore
+from semantic_kernel.connectors.memory.postgres.postgres_collection import PostgresCollection
 from semantic_kernel.data import (
     DistanceFunction,
     IndexKind,
@@ -20,6 +21,7 @@
     VectorStoreRecordVectorField,
     vectorstoremodel,
 )
+from semantic_kernel.data.vector_search.vector_search_options import VectorSearchOptions
 from semantic_kernel.exceptions.memory_connector_exceptions import (
     MemoryConnectorConnectionException,
     MemoryConnectorInitializationError,
@@ -49,13 +51,13 @@
 class SimpleDataModel(BaseModel):
     id: Annotated[int, VectorStoreRecordKeyField()]
     embedding: Annotated[
-        list[float],
+        list[float] | None,
         VectorStoreRecordVectorField(
             index_kind=IndexKind.HNSW,
             dimensions=3,
             distance_function=DistanceFunction.COSINE_SIMILARITY,
         ),
-    ]
+    ] = None
     data: Annotated[
         dict[str, Any],
         VectorStoreRecordDataField(has_embedding=True, embedding_property_name="embedding", property_type="JSONB"),
@@ -97,7 +99,9 @@ async def vector_store() -> AsyncGenerator[PostgresStore, None]:
 
 
 @asynccontextmanager
-async def create_simple_collection(vector_store: PostgresStore):
+async def create_simple_collection(
+    vector_store: PostgresStore,
+) -> AsyncGenerator[PostgresCollection[int, SimpleDataModel], None]:
     """Returns a collection with a unique name that is deleted after the context.
 
     This can be moved to use a fixture with scope=function and loop_scope=session
@@ -107,6 +111,7 @@ async def create_simple_collection(vector_store: PostgresStore):
     suffix = str(uuid.uuid4()).replace("-", "")[:8]
     collection_id = f"test_collection_{suffix}"
     collection = vector_store.get_collection(collection_id, SimpleDataModel)
+    assert isinstance(collection, PostgresCollection)
     await collection.create_collection()
     try:
         yield collection
@@ -213,6 +218,7 @@ async def test_upsert_get_and_delete_batch(vector_store: PostgresStore):
         # this should return only the two existing records.
         result = await simple_collection.get_batch([1, 2, 3])
         assert result is not None
+        assert isinstance(result, Sequence)
         assert len(result) == 2
         assert result[0] is not None
         assert result[0].id == record1.id
@@ -226,3 +232,28 @@ async def test_upsert_get_and_delete_batch(vector_store: PostgresStore):
         await simple_collection.delete_batch([1, 2])
         result_after_delete = await simple_collection.get_batch([1, 2])
         assert result_after_delete is None
+
+
+async def test_search(vector_store: PostgresStore):
+    async with create_simple_collection(vector_store) as simple_collection:
+        records = [
+            SimpleDataModel(id=1, embedding=[1.0, 0.0, 0.0], data={"key": "value1"}),
+            SimpleDataModel(id=2, embedding=[0.8, 0.2, 0.0], data={"key": "value2"}),
+            SimpleDataModel(id=3, embedding=[0.6, 0.0, 0.4], data={"key": "value3"}),
+            SimpleDataModel(id=4, embedding=[1.0, 1.0, 0.0], data={"key": "value4"}),
+            SimpleDataModel(id=5, embedding=[0.0, 1.0, 1.0], data={"key": "value5"}),
+            SimpleDataModel(id=6, embedding=[1.0, 0.0, 1.0], data={"key": "value6"}),
+        ]
+
+        await simple_collection.upsert_batch(records)
+
+        try:
+            search_results = await simple_collection.vectorized_search(
+                [1.0, 0.0, 0.0], options=VectorSearchOptions(top=3, include_total_count=True)
+            )
+            assert search_results is not None
+            assert search_results.total_count == 3
+            assert {result.record.id async for result in search_results.results} == {1, 2, 3}
+
+        finally:
+            await simple_collection.delete_batch([r.id for r in records])

python/semantic_kernel/connectors/memory/postgres/utils.py

@@ -5,6 +5,7 @@
 from typing import Annotated, Any
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
+import pytest
 import pytest_asyncio
 from psycopg import AsyncConnection, AsyncCursor
 from psycopg_pool import AsyncConnectionPool
@@ -13,6 +14,8 @@
 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
     OpenAIEmbeddingPromptExecutionSettings,
 )
+from semantic_kernel.connectors.memory.postgres.constants import DISTANCE_COLUMN_NAME
+from semantic_kernel.connectors.memory.postgres.postgres_collection import PostgresCollection
 from semantic_kernel.connectors.memory.postgres.postgres_settings import PostgresSettings
 from semantic_kernel.connectors.memory.postgres.postgres_store import PostgresStore
 from semantic_kernel.data.const import DistanceFunction, IndexKind
@@ -22,6 +25,7 @@
     VectorStoreRecordKeyField,
     VectorStoreRecordVectorField,
 )
+from semantic_kernel.data.vector_search.vector_search_options import VectorSearchOptions
 
 
 @fixture(scope="function")
@@ -76,6 +80,9 @@ class SimpleDataModel:
     ]
 
 
+# region VectorStore Tests
+
+
 async def test_vector_store_defaults(vector_store: PostgresStore) -> None:
     assert vector_store.connection_pool is not None
     async with vector_store.connection_pool.connection() as conn:
@@ -236,7 +243,130 @@ async def test_get_records(vector_store: PostgresStore, mock_cursor: Mock) -> No
     assert records[2].data == {"key": "value3"}
 
 
-# Test settings
+# endregion
+
+# region Vector Search tests
+
+
+@pytest.mark.parametrize(
+    "distance_function, operator, subquery_distance, include_vectors, include_total_count",
+    [
+        (DistanceFunction.COSINE_SIMILARITY, "<=>", f'1 - subquery."{DISTANCE_COLUMN_NAME}"', False, False),
+        (DistanceFunction.COSINE_DISTANCE, "<=>", None, False, False),
+        (DistanceFunction.DOT_PROD, "<#>", f'-1 * subquery."{DISTANCE_COLUMN_NAME}"', True, False),
+        (DistanceFunction.EUCLIDEAN_DISTANCE, "<->", None, False, True),
+        (DistanceFunction.MANHATTAN, "<+>", None, True, True),
+    ],
+)
+async def test_vector_search(
+    vector_store: PostgresStore,
+    mock_cursor: Mock,
+    distance_function: DistanceFunction,
+    operator: str,
+    subquery_distance: str | None,
+    include_vectors: bool,
+    include_total_count: bool,
+) -> None:
+    @vectorstoremodel
+    @dataclass
+    class SimpleDataModel:
+        id: Annotated[int, VectorStoreRecordKeyField()]
+        embedding: Annotated[
+            list[float],
+            VectorStoreRecordVectorField(
+                embedding_settings={"embedding": OpenAIEmbeddingPromptExecutionSettings(dimensions=1536)},
+                index_kind=IndexKind.HNSW,
+                dimensions=1536,
+                distance_function=distance_function,
+                property_type="float",
+            ),
+        ]
+        data: Annotated[
+            dict[str, Any],
+            VectorStoreRecordDataField(has_embedding=True, embedding_property_name="embedding", property_type="JSONB"),
+        ]
+
+    collection = vector_store.get_collection("test_collection", SimpleDataModel)
+    assert isinstance(collection, PostgresCollection)
+
+    search_results = await collection.vectorized_search(
+        [1.0, 2.0, 3.0],
+        options=VectorSearchOptions(
+            top=10, skip=5, include_vectors=include_vectors, include_total_count=include_total_count
+        ),
+    )
+    if include_total_count:
+        # Including total count issues query directly
+        assert mock_cursor.execute.call_count == 1
+    else:
+        # Total count is not included, query is issued when iterating over results
+        assert mock_cursor.execute.call_count == 0
+        async for _ in search_results.results:
+            pass
+        assert mock_cursor.execute.call_count == 1
+
+    execute_args, _ = mock_cursor.execute.call_args
+
+    assert (search_results.total_count is not None) == include_total_count
+
+    statement = execute_args[0]
+    statement_str = statement.as_string()
+
+    expected_columns = '"id", "data"'
+    if include_vectors:
+        expected_columns = '"id", "embedding", "data"'
+
+    expected_statement = (
+        f'SELECT {expected_columns}, "embedding" {operator} %s as "{DISTANCE_COLUMN_NAME}" '
+        'FROM "public"."test_collection" '
+        f'ORDER BY "{DISTANCE_COLUMN_NAME}" LIMIT 10 OFFSET 5'
+    )
+
+    if subquery_distance:
+        expected_statement = (
+            f'SELECT subquery.*, {subquery_distance} AS "{DISTANCE_COLUMN_NAME}" FROM ('
+            + expected_statement
+            + ") AS subquery"
+        )
+
+    assert statement_str == expected_statement
+
+
+async def test_model_post_init_conflicting_distance_column_name(vector_store: PostgresStore) -> None:
+    @vectorstoremodel
+    @dataclass
+    class ConflictingDataModel:
+        id: Annotated[int, VectorStoreRecordKeyField()]
+        sk_pg_distance: Annotated[
+            float, VectorStoreRecordDataField()
+        ]  # Note: test depends on value of DISTANCE_COLUMN_NAME constant
+
+        embedding: Annotated[
+            list[float],
+            VectorStoreRecordVectorField(
+                embedding_settings={"embedding": OpenAIEmbeddingPromptExecutionSettings(dimensions=1536)},
+                index_kind=IndexKind.HNSW,
+                dimensions=1536,
+                distance_function=DistanceFunction.COSINE_SIMILARITY,
+                property_type="float",
+            ),
+        ]
+        data: Annotated[
+            dict[str, Any],
+            VectorStoreRecordDataField(has_embedding=True, embedding_property_name="embedding", property_type="JSONB"),
+        ]
+
+    collection = vector_store.get_collection("test_collection", ConflictingDataModel)
+    assert isinstance(collection, PostgresCollection)
+
+    # Ensure that the distance column name has been changed to avoid conflict
+    assert collection._distance_column_name != DISTANCE_COLUMN_NAME
+    assert collection._distance_column_name.startswith(f"{DISTANCE_COLUMN_NAME}_")
+
+
+# endregion
+
+# region Settings tests
 
 
 def test_settings_connection_string(monkeypatch) -> None:
@@ -290,3 +420,6 @@ def test_settings_env_vars(monkeypatch) -> None:
     assert conn_info["dbname"] == "dbname"
     assert conn_info["user"] == "user"
     assert conn_info["password"] == "password"
+
+
+# endregion