chore: import upstream snapshot with attribution

python/tests/integration/memory/azure_cosmos_db/conftest.py

@@ -0,0 +1,77 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from dataclasses import field
+from typing import Annotated, Any
+from uuid import uuid4
+
+from pydantic import BaseModel
+from pytest import fixture
+
+from semantic_kernel.data.vector import VectorStoreField, vectorstoremodel
+
+
+@fixture
+def data_record() -> dict[str, Any]:
+    return {
+        "id": "e6103c03-487f-4d7d-9c23-4723651c17f4",
+        "description": "This is a test record",
+        "product_type": "test",
+        "vector": [0.1, 0.2, 0.3, 0.4, 0.5],
+    }
+
+
+@fixture
+def record_type() -> type:
+    @vectorstoremodel
+    class TestDataModelType(BaseModel):
+        vector: Annotated[
+            list[float] | None,
+            VectorStoreField(
+                "vector",
+                index_kind="flat",
+                dimensions=5,
+                distance_function="cosine_similarity",
+                type="float",
+            ),
+        ] = None
+        id: Annotated[str, VectorStoreField("key")] = field(default_factory=lambda: str(uuid4()))
+        product_type: Annotated[str, VectorStoreField("data")] = "N/A"
+        description: Annotated[
+            str, VectorStoreField("data", has_embedding=True, embedding_property_name="vector", type="str")
+        ] = "N/A"
+
+    return TestDataModelType
+
+
+@fixture
+def data_record_with_key_as_key_field() -> dict[str, Any]:
+    return {
+        "key": "e6103c03-487f-4d7d-9c23-4723651c17f4",
+        "description": "This is a test record",
+        "product_type": "test",
+        "vector": [0.1, 0.2, 0.3, 0.4, 0.5],
+    }
+
+
+@fixture
+def record_type_with_key_as_key_field() -> type:
+    @vectorstoremodel
+    class TestDataModelType(BaseModel):
+        vector: Annotated[
+            list[float] | None,
+            VectorStoreField(
+                "vector",
+                index_kind="flat",
+                dimensions=5,
+                distance_function="cosine_similarity",
+                type="float",
+            ),
+        ] = None
+        key: Annotated[str, VectorStoreField("key")] = field(default_factory=lambda: str(uuid4()))
+        product_type: Annotated[str, VectorStoreField("data")] = "N/A"
+        description: Annotated[
+            str, VectorStoreField("data", has_embedding=True, embedding_property_name="vector", type="str")
+        ] = "N/A"
+
+    return TestDataModelType

python/tests/integration/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql.py

@@ -0,0 +1,231 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+import platform
+from collections.abc import Callable
+from typing import Any
+
+import pytest
+from azure.cosmos.aio import CosmosClient
+from azure.cosmos.partition_key import PartitionKey
+
+from semantic_kernel.connectors.azure_cosmos_db import CosmosNoSqlCompositeKey, CosmosNoSqlStore
+from semantic_kernel.data.vector import VectorStore
+from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorException
+from tests.integration.memory.vector_store_test_base import VectorStoreTestBase
+
+
+@pytest.mark.skipif(
+    platform.system() != "Windows",
+    reason="The Azure Cosmos DB Emulator is only available on Windows.",
+)
+class TestCosmosDBNoSQL(VectorStoreTestBase):
+    """Test Cosmos DB NoSQL store functionality."""
+
+    async def test_list_collection_names(
+        self,
+        stores: dict[str, Callable[[], VectorStore]],
+        record_type: type,
+    ):
+        """Test list collection names."""
+        async with stores["azure_cosmos_db_no_sql"]() as store:
+            assert await store.list_collection_names() == []
+
+            collection_name = "list_collection_names"
+            collection = store.get_collection(collection_name=collection_name, record_type=record_type)
+            await collection.ensure_collection_exists()
+
+            collection_names = await store.list_collection_names()
+            assert collection_name in collection_names
+
+            await collection.ensure_collection_deleted()
+            assert await collection.collection_exists() is False
+            collection_names = await store.list_collection_names()
+            assert collection_name not in collection_names
+
+    async def test_collection_not_created(
+        self,
+        stores: dict[str, Callable[[], VectorStore]],
+        record_type: type,
+        data_record: dict[str, Any],
+    ):
+        """Test get without collection."""
+        async with stores["azure_cosmos_db_no_sql"]() as store:
+            collection_name = "collection_not_created"
+            collection = store.get_collection(collection_name=collection_name, record_type=record_type)
+
+            assert await collection.collection_exists() is False
+
+            with pytest.raises(
+                MemoryConnectorException, match="The collection does not exist yet. Create the collection first."
+            ):
+                await collection.upsert(record_type(**data_record))
+
+            with pytest.raises(
+                MemoryConnectorException, match="The collection does not exist yet. Create the collection first."
+            ):
+                await collection.get(data_record["id"])
+
+            with pytest.raises(MemoryConnectorException):
+                await collection.delete(data_record["id"])
+
+            with pytest.raises(MemoryConnectorException, match="Container could not be deleted."):
+                await collection.ensure_collection_deleted()
+
+    async def test_custom_partition_key(
+        self,
+        stores: dict[str, Callable[[], VectorStore]],
+        record_type: type,
+        data_record: dict[str, Any],
+    ):
+        """Test custom partition key."""
+        async with stores["azure_cosmos_db_no_sql"]() as store:
+            collection_name = "custom_partition_key"
+            collection = store.get_collection(
+                collection_name=collection_name,
+                record_type=record_type,
+                partition_key=PartitionKey(path="/product_type"),
+            )
+
+            composite_key = CosmosNoSqlCompositeKey(key=data_record["id"], partition_key=data_record["product_type"])
+
+            # Upsert
+            await collection.ensure_collection_exists()
+            await collection.upsert(record_type(**data_record))
+
+            # Verify
+            record = await collection.get(composite_key)
+            assert record is not None
+            assert isinstance(record, record_type)
+
+            # Remove
+            await collection.delete(composite_key)
+            record = await collection.get(composite_key)
+            assert record is None
+
+            # Remove collection
+            await collection.ensure_collection_deleted()
+            assert await collection.collection_exists() is False
+
+    async def test_get_include_vector(
+        self,
+        stores: dict[str, Callable[[], VectorStore]],
+        record_type: type,
+        data_record: dict[str, Any],
+    ):
+        """Test get with include_vector."""
+        async with stores["azure_cosmos_db_no_sql"]() as store:
+            collection_name = "get_include_vector"
+            collection = store.get_collection(collection_name=collection_name, record_type=record_type)
+
+            # Upsert
+            await collection.ensure_collection_exists()
+            await collection.upsert(record_type(**data_record))
+
+            # Verify
+            record = await collection.get(data_record["id"], include_vectors=True)
+            assert record is not None
+            assert isinstance(record, record_type)
+            assert record.vector == data_record["vector"]
+
+            # Remove
+            await collection.delete(data_record["id"])
+            record = await collection.get(data_record["id"])
+            assert record is None
+
+            # Remove collection
+            await collection.ensure_collection_deleted()
+            assert await collection.collection_exists() is False
+
+    async def test_get_not_include_vector(
+        self,
+        stores: dict[str, Callable[[], VectorStore]],
+        record_type: type,
+        data_record: dict[str, Any],
+    ):
+        """Test get with include_vector."""
+        async with stores["azure_cosmos_db_no_sql"]() as store:
+            collection_name = "get_not_include_vector"
+            collection = store.get_collection(collection_name=collection_name, record_type=record_type)
+
+            # Upsert
+            await collection.ensure_collection_exists()
+            await collection.upsert(record_type(**data_record))
+
+            # Verify
+            record = await collection.get(data_record["id"], include_vectors=False)
+            assert record is not None
+            assert isinstance(record, record_type)
+            assert record.vector is None
+
+            # Remove
+            await collection.delete(data_record["id"])
+            record = await collection.get(data_record["id"])
+            assert record is None
+
+            # Remove collection
+            await collection.ensure_collection_deleted()
+            assert await collection.collection_exists() is False
+
+    async def test_collection_with_key_as_key_field(
+        self,
+        stores: dict[str, Callable[[], VectorStore]],
+        record_type_with_key_as_key_field: type,
+        data_record_with_key_as_key_field: dict[str, Any],
+    ):
+        """Test collection with key as key field."""
+        async with stores["azure_cosmos_db_no_sql"]() as store:
+            collection_name = "collection_with_key_as_key_field"
+            collection = store.get_collection(
+                collection_name=collection_name, record_type=record_type_with_key_as_key_field
+            )
+
+            # Upsert
+            await collection.ensure_collection_exists()
+            result = await collection.upsert(record_type_with_key_as_key_field(**data_record_with_key_as_key_field))
+            assert data_record_with_key_as_key_field["key"] == result
+
+            # Verify
+            record = await collection.get(data_record_with_key_as_key_field["key"])
+            assert record is not None
+            assert isinstance(record, record_type_with_key_as_key_field)
+            assert record.key == data_record_with_key_as_key_field["key"]
+
+            # Remove
+            await collection.delete(data_record_with_key_as_key_field["key"])
+            record = await collection.get(data_record_with_key_as_key_field["key"])
+            assert record is None
+
+            # Remove collection
+            await collection.ensure_collection_deleted()
+            assert await collection.collection_exists() is False
+
+    async def test_custom_client(
+        self,
+        record_type: type,
+    ):
+        """Test list collection names."""
+        url = os.environ.get("AZURE_COSMOS_DB_NO_SQL_URL")
+        key = os.environ.get("AZURE_COSMOS_DB_NO_SQL_KEY")
+
+        async with (
+            CosmosClient(url, key) as custom_client,
+            CosmosNoSqlStore(
+                database_name="test_database",
+                cosmos_client=custom_client,
+                create_database=True,
+            ) as store,
+        ):
+            assert await store.list_collection_names() == []
+
+            collection_name = "list_collection_names"
+            collection = store.get_collection(collection_name=collection_name, record_type=record_type)
+            await collection.ensure_collection_exists()
+
+            collection_names = await store.list_collection_names()
+            assert collection_name in collection_names
+
+            await collection.ensure_collection_deleted()
+            assert await collection.collection_exists() is False
+            collection_names = await store.list_collection_names()
+            assert collection_name not in collection_names

python/tests/integration/memory/data_records.py

@@ -0,0 +1,145 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+import numpy as np
+
+RAW_RECORD_LIST = {
+    "id": "e6103c03-487f-4d7d-9c23-4723651c17f4",
+    "content": "test content",
+    "vector": [0.1, 0.2, 0.3, 0.4, 0.5],
+}
+
+
+RAW_RECORD_ARRAY = {
+    "id": "e6103c03-487f-4d7d-9c23-4723651c17f4",
+    "content": "test content",
+    "vector": np.array([0.1, 0.2, 0.3, 0.4, 0.5]),
+}
+
+
+# PANDAS_RECORD_DEFINITION = VectorStoreRecordDefinition(
+#     fields={
+#         "vector": VectorStoreRecordVectorField(
+#             name="vector",
+#             index_kind="hnsw",
+#             dimensions=5,
+#             distance_function="cosine_similarity",
+#             property_type="float",
+#         ),
+#         "id": VectorStoreRecordKeyField(name="id"),
+#         "content": VectorStoreRecordDataField(
+#             name="content", has_embedding=True, embedding_property_name="vector", property_type="str"
+#         ),
+#     },
+#     container_mode=True,
+#     to_dict=lambda x: x.to_dict(orient="records"),
+#     from_dict=lambda x, **_: pd.DataFrame(x),
+# )
+
+# A Pandas record definition with flat index kind
+# PANDAS_RECORD_DEFINITION_FLAT = VectorStoreRecordDefinition(
+#     fields={
+#         "vector": VectorStoreRecordVectorField(
+#             name="vector",
+#             index_kind="flat",
+#             dimensions=5,
+#             distance_function="cosine_similarity",
+#             property_type="float",
+#         ),
+#         "id": VectorStoreRecordKeyField(name="id"),
+#         "content": VectorStoreRecordDataField(
+#             name="content", has_embedding=True, embedding_property_name="vector", property_type="str"
+#         ),
+#     },
+#     container_mode=True,
+#     to_dict=lambda x: x.to_dict(orient="records"),
+#     from_dict=lambda x, **_: pd.DataFrame(x),
+# )
+
+
+# @vectorstoremodel
+# @dataclass
+# class TestDataModelArray(distance_function: str):
+#     """A data model where the vector is a numpy array."""
+
+#     vector: Annotated[
+#         np.ndarray | None,
+#         VectorStoreRecordVectorField(
+#             index_kind="hnsw",
+#             dimensions=5,
+#             distance_function=distance_function,
+#             property_type="float",
+#             serialize_function=np.ndarray.tolist,
+#             deserialize_function=np.array,
+#         ),
+#     ] = None
+#     other: str | None = None
+#     id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
+#     content: Annotated[
+#         str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
+#     ] = "content1"
+
+
+# @vectorstoremodel
+# @dataclass
+# class TestDataModelArrayFlat(distance_function:str):
+#     """A data model where the vector is a numpy array and the index kind is IndexKind.Flat."""
+
+#     vector: Annotated[
+#         np.ndarray | None,
+#         VectorStoreRecordVectorField(
+#             index_kind="flat",
+#             dimensions=5,
+#             distance_function=distance_function,
+#             property_type="float",
+#             serialize_function=np.ndarray.tolist,
+#             deserialize_function=np.array,
+#         ),
+#     ] = None
+#     other: str | None = None
+#     id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
+#     content: Annotated[
+#         str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
+#     ] = "content1"
+
+
+# @vectorstoremodel
+# @dataclass
+# class TestDataModelList(distance_function: str):
+#     """A data model where the vector is a list."""
+
+#     vector: Annotated[
+#         list[float] | None,
+#         VectorStoreRecordVectorField(
+#             index_kind="hnsw",
+#             dimensions=5,
+#             distance_function=distance_function,
+#             property_type="float",
+#         ),
+#     ] = None
+#     other: str | None = None
+#     id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
+#     content: Annotated[
+#         str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
+#     ] = "content1"
+
+
+# @vectorstoremodel
+# @dataclass
+# class TestDataModelListFlat:
+#     """A data model where the vector is a list and the index kind is IndexKind.Flat."""
+
+#     vector: Annotated[
+#         list[float] | None,
+#         VectorStoreRecordVectorField(
+#             index_kind="flat",
+#             dimensions=5,
+#             distance_function="cosine_similarity",
+#             property_type="float",
+#         ),
+#     ] = None
+#     other: str | None = None
+#     id: Annotated[str, VectorStoreRecordKeyField()] = field(default_factory=lambda: str(uuid4()))
+#     content: Annotated[
+#         str, VectorStoreRecordDataField(has_embedding=True, embedding_property_name="vector", property_type="str")
+#     ] = "content1"

python/tests/integration/memory/postgres/test_postgres_int.py

@@ -0,0 +1,259 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import uuid
+from collections.abc import AsyncGenerator, Sequence
+from contextlib import asynccontextmanager
+from typing import Annotated, Any
+
+import pandas as pd
+import pytest
+import pytest_asyncio
+from pydantic import BaseModel
+
+from semantic_kernel.connectors.postgres import PostgresCollection, PostgresSettings, PostgresStore
+from semantic_kernel.data.vector import (
+    DistanceFunction,
+    IndexKind,
+    VectorStoreCollectionDefinition,
+    VectorStoreField,
+    vectorstoremodel,
+)
+from semantic_kernel.exceptions.memory_connector_exceptions import (
+    MemoryConnectorConnectionException,
+    MemoryConnectorInitializationError,
+)
+
+try:
+    import psycopg  # noqa: F401
+    import psycopg_pool  # noqa: F401
+
+    psycopg_pool_installed = True
+except ImportError:
+    psycopg_pool_installed = False
+
+pg_settings: PostgresSettings = PostgresSettings()
+try:
+    connection_params_present = any(pg_settings.get_connection_args().values())
+except MemoryConnectorInitializationError:
+    connection_params_present = False
+
+pytestmark = pytest.mark.skipif(
+    not (psycopg_pool_installed or connection_params_present),
+    reason="psycopg_pool is not installed" if not psycopg_pool_installed else "No connection parameters provided",
+)
+
+
+@vectorstoremodel
+class SimpleDataModel(BaseModel):
+    id: Annotated[int, VectorStoreField("key")]
+    embedding: Annotated[
+        list[float] | str | None,
+        VectorStoreField(
+            "vector",
+            index_kind=IndexKind.HNSW,
+            dimensions=3,
+            distance_function=DistanceFunction.COSINE_SIMILARITY,
+        ),
+    ] = None
+    data: Annotated[
+        dict[str, Any],
+        VectorStoreField("data", type="JSONB"),
+    ]
+
+    def model_post_init(self, context: Any) -> None:
+        if self.embedding is None:
+            self.embedding = self.data
+
+
+def DataModelPandas(record) -> tuple:
+    definition = VectorStoreCollectionDefinition(
+        fields=[
+            VectorStoreField(
+                "vector",
+                name="embedding",
+                index_kind="hnsw",
+                dimensions=3,
+                distance_function="cosine_similarity",
+                type="float",
+            ),
+            VectorStoreField("key", name="id", type="int"),
+            VectorStoreField("data", name="data", type="dict"),
+        ],
+        container_mode=True,
+        to_dict=lambda x: x.to_dict(orient="records"),
+        from_dict=lambda x, **_: pd.DataFrame(x),
+    )
+    df = pd.DataFrame([record])
+    return definition, df
+
+
+@pytest_asyncio.fixture
+async def vector_store() -> AsyncGenerator[PostgresStore, None]:
+    try:
+        async with await pg_settings.create_connection_pool() as pool:
+            yield PostgresStore(connection_pool=pool)
+    except MemoryConnectorConnectionException:
+        pytest.skip("Postgres connection not available")
+        yield None
+        return
+
+
+@asynccontextmanager
+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
+    after upgrade to pytest-asyncio 0.24. With the current version, the fixture
+    would both cache and use the event loop of the declared scope.
+    """
+    suffix = str(uuid.uuid4()).replace("-", "")[:8]
+    collection_id = f"test_collection_{suffix}"
+    collection = vector_store.get_collection(collection_name=collection_id, record_type=SimpleDataModel)
+    assert isinstance(collection, PostgresCollection)
+    await collection.ensure_collection_exists()
+    try:
+        yield collection
+    finally:
+        await collection.ensure_collection_deleted()
+
+
+def test_create_store(vector_store):
+    assert vector_store is not None
+    assert vector_store.connection_pool is not None
+
+
+async def test_ensure_collection_exists_exists_and_delete(vector_store: PostgresStore):
+    suffix = str(uuid.uuid4()).replace("-", "")[:8]
+
+    collection = vector_store.get_collection(collection_name=f"test_collection_{suffix}", record_type=SimpleDataModel)
+
+    does_exist_1 = await collection.collection_exists()
+    assert does_exist_1 is False
+
+    await collection.ensure_collection_exists()
+    does_exist_2 = await collection.collection_exists()
+    assert does_exist_2 is True
+
+    await collection.ensure_collection_deleted()
+    does_exist_3 = await collection.collection_exists()
+    assert does_exist_3 is False
+
+
+async def test_list_collection_names(vector_store):
+    async with create_simple_collection(vector_store) as simple_collection:
+        simple_collection_id = simple_collection.collection_name
+        result = await vector_store.list_collection_names()
+        assert simple_collection_id in result
+
+
+async def test_upsert_get_and_delete(vector_store: PostgresStore):
+    record = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
+    async with create_simple_collection(vector_store) as simple_collection:
+        result_before_upsert = await simple_collection.get(1)
+        assert result_before_upsert is None
+
+        await simple_collection.upsert(record)
+        result = await simple_collection.get(1)
+        assert result is not None
+        assert result.id == record.id
+        assert result.embedding == record.embedding
+        assert result.data == record.data
+
+        # Check that the table has an index
+        connection_pool = simple_collection.connection_pool
+        async with connection_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                "SELECT indexname FROM pg_indexes WHERE tablename = %s", (simple_collection.collection_name,)
+            )
+            rows = await cur.fetchall()
+            index_names = [index[0] for index in rows]
+            assert any("embedding_idx" in index_name for index_name in index_names)
+
+        await simple_collection.delete(1)
+        result_after_delete = await simple_collection.get(1)
+        assert result_after_delete is None
+
+
+async def test_upsert_get_and_delete_pandas(vector_store):
+    record = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
+    definition, df = DataModelPandas(record.model_dump())
+
+    suffix = str(uuid.uuid4()).replace("-", "")[:8]
+    collection = vector_store.get_collection(
+        collection_name=f"test_collection_{suffix}",
+        record_type=pd.DataFrame,
+        definition=definition,
+    )
+    await collection.ensure_collection_exists()
+
+    try:
+        result_before_upsert = await collection.get(1)
+        assert result_before_upsert is None
+
+        await collection.upsert(df)
+        result: pd.DataFrame = await collection.get(1)
+        assert result is not None
+        row = result.iloc[0]
+        assert row.id == record.id
+        assert row.embedding == record.embedding
+        assert row.data == record.data
+
+        await collection.delete(1)
+        result_after_delete = await collection.get(1)
+        assert result_after_delete is None
+    finally:
+        await collection.ensure_collection_deleted()
+
+
+async def test_upsert_get_and_delete_multiple(vector_store: PostgresStore):
+    async with create_simple_collection(vector_store) as simple_collection:
+        record1 = SimpleDataModel(id=1, embedding=[1.1, 2.2, 3.3], data={"key": "value"})
+        record2 = SimpleDataModel(id=2, embedding=[4.4, 5.5, 6.6], data={"key": "value"})
+
+        result_before_upsert = await simple_collection.get([1, 2])
+        assert result_before_upsert is None
+
+        await simple_collection.upsert([record1, record2])
+        # Test get for the two existing keys and one non-existing key;
+        # this should return only the two existing records.
+        result = await simple_collection.get([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
+        assert result[0].embedding == record1.embedding
+        assert result[0].data == record1.data
+        assert result[1] is not None
+        assert result[1].id == record2.id
+        assert result[1].embedding == record2.embedding
+        assert result[1].data == record2.data
+
+        await simple_collection.delete([1, 2])
+        result_after_delete = await simple_collection.get([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(records)
+
+        try:
+            search_results = await simple_collection.search(vector=[1.0, 0.0, 0.0], 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([r.id for r in records])

python/tests/integration/memory/test_vector_store.py

@@ -0,0 +1,363 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import platform
+from collections.abc import Callable
+from typing import Any
+
+import pandas as pd
+import pytest
+
+from semantic_kernel.connectors.redis import RedisCollectionTypes
+from semantic_kernel.data.vector import VectorStore
+from semantic_kernel.exceptions import MemoryConnectorConnectionException
+from tests.integration.memory.data_records import RAW_RECORD_LIST
+from tests.integration.memory.vector_store_test_base import VectorStoreTestBase
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class TestVectorStore(VectorStoreTestBase):
+    """Test vector store functionality.
+
+    This only tests if the vector stores can upsert, get, and delete records.
+    """
+
+    @pytest.mark.parametrize(
+        [
+            "store_id",
+            "collection_name",
+            "collection_options",
+            "record_type",
+            "definition",
+            "distance_function",
+            "index_kind",
+            "vector_property_type",
+            "dimensions",
+            "record",
+        ],
+        [
+            # region Redis
+            pytest.param(
+                "redis",
+                "redis_json_list_data_model",
+                {"collection_type": RedisCollectionTypes.JSON},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="redis_json_list_data_model",
+            ),
+            pytest.param(
+                "redis",
+                "redis_json_pandas_data_model",
+                {"collection_type": RedisCollectionTypes.JSON},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="redis_json_pandas_data_model",
+            ),
+            pytest.param(
+                "redis",
+                "redis_hashset_list_data_model",
+                {"collection_type": RedisCollectionTypes.HASHSET},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="redis_hashset_list_data_model",
+            ),
+            pytest.param(
+                "redis",
+                "redis_hashset_pandas_data_model",
+                {"collection_type": RedisCollectionTypes.HASHSET},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="redis_hashset_pandas_data_model",
+            ),
+            # endregion
+            # region Azure AI Search
+            pytest.param(
+                "azure_ai_search",
+                "azure_ai_search_list_data_model",
+                {},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="azure_ai_search_list_data_model",
+            ),
+            pytest.param(
+                "azure_ai_search",
+                "azure_ai_search_pandas_data_model",
+                {},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="azure_ai_search_pandas_data_model",
+            ),
+            # endregion
+            # region Qdrant
+            pytest.param(
+                "qdrant",
+                "qdrant_list_data_model",
+                {},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="qdrant_list_data_model",
+            ),
+            pytest.param(
+                "qdrant",
+                "qdrant_pandas_data_model",
+                {},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="qdrant_pandas_data_model",
+            ),
+            pytest.param(
+                "qdrant_in_memory",
+                "qdrant_in_memory_list_data_model",
+                {},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="qdrant_in_memory_list_data_model",
+            ),
+            pytest.param(
+                "qdrant_in_memory",
+                "qdrant_in_memory_pandas_data_model",
+                {},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="qdrant_in_memory_pandas_data_model",
+            ),
+            pytest.param(
+                "qdrant",
+                "qdrant_grpc_list_data_model",
+                {"prefer_grpc": True},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="qdrant_grpc_list_data_model",
+            ),
+            pytest.param(
+                "qdrant",
+                "qdrant_grpc_pandas_data_model",
+                {"prefer_grpc": True},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="qdrant_grpc_pandas_data_model",
+            ),
+            # endregion
+            # region Weaviate
+            pytest.param(
+                "weaviate_local",
+                "weaviate_local_list_data_model",
+                {},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                marks=pytest.mark.skipif(
+                    platform.system() != "Linux",
+                    reason="The Weaviate docker image is only available on Linux"
+                    " but some GitHubs job runs in a Windows container.",
+                ),
+                id="weaviate_local_list_data_model",
+            ),
+            pytest.param(
+                "weaviate_local",
+                "weaviate_local_pandas_data_model",
+                {},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                marks=pytest.mark.skipif(
+                    platform.system() != "Linux",
+                    reason="The Weaviate docker image is only available on Linux"
+                    " but some GitHubs job runs in a Windows container.",
+                ),
+                id="weaviate_local_pandas_data_model",
+            ),
+            # endregion
+            # region Azure Cosmos DB
+            pytest.param(
+                "azure_cosmos_db_no_sql",
+                "azure_cosmos_db_no_sql_list_data_model",
+                {},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                "flat",
+                None,
+                5,
+                RAW_RECORD_LIST,
+                marks=pytest.mark.skipif(
+                    platform.system() != "Windows",
+                    reason="The Azure Cosmos DB Emulator is only available on Windows.",
+                ),
+                id="azure_cosmos_db_no_sql_list_data_model",
+            ),
+            pytest.param(
+                "azure_cosmos_db_no_sql",
+                "azure_cosmos_db_no_sql_pandas_data_model",
+                {},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                "flat",
+                None,
+                5,
+                RAW_RECORD_LIST,
+                marks=pytest.mark.skipif(
+                    platform.system() != "Windows",
+                    reason="The Azure Cosmos DB Emulator is only available on Windows.",
+                ),
+                id="azure_cosmos_db_no_sql_pandas_data_model",
+            ),
+            # endregion
+            # region Chroma
+            pytest.param(
+                "chroma",
+                "chroma_list_data_model",
+                {},
+                "dataclass_vector_data_model",
+                None,
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="chroma_list_data_model",
+            ),
+            pytest.param(
+                "chroma",
+                "chroma_pandas_data_model",
+                {},
+                pd.DataFrame,
+                "definition_pandas",
+                None,
+                None,
+                None,
+                5,
+                RAW_RECORD_LIST,
+                id="chroma_pandas_data_model",
+            ),
+            # endregion
+        ],
+    )
+    # region test function
+    async def test_vector_store(
+        self,
+        stores: dict[str, Callable[[], VectorStore]],
+        store_id: str,
+        collection_name: str,
+        collection_options: dict[str, Any],
+        record_type: str | type,
+        definition: str | None,
+        distance_function,
+        index_kind,
+        vector_property_type,
+        dimensions,
+        record: dict[str, Any],
+        request,
+    ):
+        """Test vector store functionality."""
+        if isinstance(record_type, str):
+            record_type = request.getfixturevalue(record_type)
+        if definition is not None:
+            definition = request.getfixturevalue(definition)
+        try:
+            async with (
+                stores[store_id]() as vector_store,
+                vector_store.get_collection(
+                    record_type=record_type,
+                    definition=definition,
+                    collection_name=collection_name,
+                    **collection_options,
+                ) as collection,
+            ):
+                try:
+                    await collection.ensure_collection_deleted()
+                except Exception as exc:
+                    logger.warning(f"Failed to delete collection: {exc}")
+
+                try:
+                    await collection.ensure_collection_exists()
+                except Exception as exc:
+                    pytest.fail(f"Failed to create collection: {exc}")
+
+                # Upsert record
+                await collection.upsert(record_type([record]) if record_type is pd.DataFrame else record_type(**record))
+                # Get record
+                result = await collection.get(record["id"])
+                assert result is not None
+                # Delete record
+                await collection.delete(record["id"])
+                # Get record again, expect None
+                result = await collection.get(record["id"])
+                assert result is None
+
+                try:
+                    await collection.ensure_collection_deleted()
+                except Exception as exc:
+                    pytest.fail(f"Failed to delete collection: {exc}")
+        except MemoryConnectorConnectionException as exc:
+            pytest.xfail(f"Failed to connect to store: {exc}")

python/tests/integration/memory/vector_store_test_base.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Callable
+
+import pytest
+
+from semantic_kernel.data.vector import VectorStore
+
+
+def get_redis_store():
+    from semantic_kernel.connectors.redis import RedisStore
+
+    return RedisStore()
+
+
+def get_azure_ai_search_store():
+    from semantic_kernel.connectors.azure_ai_search import AzureAISearchStore
+
+    return AzureAISearchStore()
+
+
+def get_qdrant_store():
+    from semantic_kernel.connectors.qdrant import QdrantStore
+
+    return QdrantStore()
+
+
+def get_qdrant_store_in_memory():
+    from semantic_kernel.connectors.qdrant import QdrantStore
+
+    return QdrantStore(location=":memory:")
+
+
+def get_weaviate_store():
+    from semantic_kernel.connectors.weaviate import WeaviateStore
+
+    return WeaviateStore(local_host="localhost")
+
+
+def get_azure_cosmos_db_no_sql_store():
+    from semantic_kernel.connectors.azure_cosmos_db import CosmosNoSqlStore
+
+    return CosmosNoSqlStore(database_name="test_database", create_database=True)
+
+
+def get_chroma_store():
+    from semantic_kernel.connectors.chroma import ChromaStore
+
+    return ChromaStore()
+
+
+class VectorStoreTestBase:
+    @pytest.fixture
+    def stores(self) -> dict[str, Callable[[], VectorStore]]:
+        """Return a dictionary of vector stores to test."""
+        return {
+            "redis": get_redis_store,
+            "azure_ai_search": get_azure_ai_search_store,
+            "qdrant": get_qdrant_store,
+            "qdrant_in_memory": get_qdrant_store_in_memory,
+            "weaviate_local": get_weaviate_store,
+            "azure_cosmos_db_no_sql": get_azure_cosmos_db_no_sql_store,
+            "chroma": get_chroma_store,
+        }