diff --git a/python/semantic_kernel/memory/volatile_memory_store.py b/python/semantic_kernel/memory/volatile_memory_store.py index dd1a631d0614..e4b7839c08ff 100644 --- a/python/semantic_kernel/memory/volatile_memory_store.py +++ b/python/semantic_kernel/memory/volatile_memory_store.py @@ -154,9 +154,9 @@ async def get_batch( results = [self._store[collection_name][key] for key in keys if key in self._store[collection_name]] if not with_embeddings: - # create copy of results without embeddings + # create copies of results without embeddings + results = [deepcopy(result) for result in results] for result in results: - result = deepcopy(result) result._embedding = None return results @@ -280,9 +280,9 @@ async def get_nearest_matches( top_results = filtered_results[:limit] if not with_embeddings: - # create copy of results without embeddings + # create copies of results without embeddings + top_results = [deepcopy(result) for result in top_results] for result in top_results: - result = deepcopy(result) result[0]._embedding = None return top_results diff --git a/python/tests/unit/memory/test_volatile_memory_store_embeddings.py b/python/tests/unit/memory/test_volatile_memory_store_embeddings.py new file mode 100644 index 000000000000..0ee23d72a2b3 --- /dev/null +++ b/python/tests/unit/memory/test_volatile_memory_store_embeddings.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft. All rights reserved. + +import warnings + +import numpy as np +import pytest + +from semantic_kernel.memory.memory_record import MemoryRecord +from semantic_kernel.memory.volatile_memory_store import VolatileMemoryStore + + +@pytest.fixture +async def store() -> VolatileMemoryStore: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + memory_store = VolatileMemoryStore() + await memory_store.create_collection("test") + await memory_store.upsert("test", MemoryRecord.local_record("a", "text a", None, None, np.array([1.0, 0.0]))) + return memory_store + + +async def test_get_batch_without_embeddings(store: VolatileMemoryStore): + results = await store.get_batch("test", ["a"], with_embeddings=False) + + assert results[0]._embedding is None + # the stored record must keep its embedding + assert store._store["test"]["a"]._embedding is not None + + +async def test_get_batch_with_embeddings(store: VolatileMemoryStore): + results = await store.get_batch("test", ["a"], with_embeddings=True) + + assert results[0]._embedding is not None + + +async def test_get_nearest_matches_without_embeddings(store: VolatileMemoryStore): + results = await store.get_nearest_matches("test", np.array([1.0, 0.0]), limit=1, with_embeddings=False) + + assert results[0][0]._embedding is None + assert store._store["test"]["a"]._embedding is not None + # a second query still works, so the stored embedding was not cleared + again = await store.get_nearest_matches("test", np.array([1.0, 0.0]), limit=1, with_embeddings=True) + assert again[0][0]._embedding is not None + + +async def test_get_nearest_matches_with_embeddings(store: VolatileMemoryStore): + results = await store.get_nearest_matches("test", np.array([1.0, 0.0]), limit=1, with_embeddings=True) + + assert results[0][0]._embedding is not None