chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,428 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import ast
|
||||
import asyncio
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, Any
|
||||
|
||||
from pandas import DataFrame
|
||||
from pydantic import BaseModel, Field
|
||||
from pytest import fixture
|
||||
|
||||
from semantic_kernel.data.vector import (
|
||||
KernelSearchResults,
|
||||
SearchType,
|
||||
VectorSearch,
|
||||
VectorSearchResult,
|
||||
VectorStoreCollection,
|
||||
VectorStoreCollectionDefinition,
|
||||
VectorStoreField,
|
||||
vectorstoremodel,
|
||||
)
|
||||
from semantic_kernel.kernel_types import OptionalOneOrMany
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def _ensure_event_loop():
|
||||
"""Ensure a current event loop exists before each test.
|
||||
|
||||
Works around a pytest-asyncio 0.26 bug on Windows Python 3.10 where
|
||||
asyncio.set_event_loop(None) can be left as state after a previous test's
|
||||
teardown, and _provide_clean_event_loop does not recover because it only
|
||||
creates a fresh loop when old_loop is not None. By guaranteeing a non-None
|
||||
loop at fixture-setup time, _temporary_event_loop_policy saves a valid
|
||||
old_loop, so the teardown path restores a valid loop instead of None.
|
||||
"""
|
||||
try:
|
||||
asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
asyncio.set_event_loop(asyncio.new_event_loop())
|
||||
yield
|
||||
|
||||
|
||||
@fixture
|
||||
def DictVectorStoreRecordCollection() -> type[VectorSearch]:
|
||||
class DictVectorStoreRecordCollection(
|
||||
VectorStoreCollection[str, Any],
|
||||
VectorSearch[str, Any],
|
||||
):
|
||||
supported_search_types = {SearchType.VECTOR}
|
||||
inner_storage: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
async def _inner_delete(self, keys: Sequence[str], **kwargs: Any) -> None:
|
||||
for key in keys:
|
||||
self.inner_storage.pop(key, None)
|
||||
|
||||
async def _inner_get(self, keys: Sequence[str], **kwargs: Any) -> Any | Sequence[Any] | None:
|
||||
return [self.inner_storage[key] for key in keys if key in self.inner_storage]
|
||||
|
||||
async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[str]:
|
||||
updated_keys = []
|
||||
for record in records:
|
||||
key = (
|
||||
record[self._key_field_name]
|
||||
if isinstance(record, Mapping)
|
||||
else getattr(record, self._key_field_name)
|
||||
)
|
||||
self.inner_storage[key] = record
|
||||
updated_keys.append(key)
|
||||
return updated_keys
|
||||
|
||||
def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]:
|
||||
return records
|
||||
|
||||
def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]:
|
||||
return records
|
||||
|
||||
async def ensure_collection_exists(self, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
async def ensure_collection_deleted(self, **kwargs: Any) -> None:
|
||||
self.inner_storage = {}
|
||||
|
||||
async def collection_exists(self, **kwargs: Any) -> bool:
|
||||
return True
|
||||
|
||||
async def _inner_search(
|
||||
self,
|
||||
options: Any = None,
|
||||
keywords: OptionalOneOrMany[str] = None,
|
||||
search_text: str | None = None,
|
||||
vectorizable_text: str | None = None,
|
||||
vector: list[float | int] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
return KernelSearchResults(
|
||||
results=self.generator(),
|
||||
total_count=len(self.inner_storage) if options.include_total_count else None,
|
||||
)
|
||||
|
||||
def _get_record_from_result(self, result: Any) -> Any:
|
||||
return result
|
||||
|
||||
def _get_score_from_result(self, result: Any) -> float | None:
|
||||
return None
|
||||
|
||||
async def generator(self):
|
||||
if self.inner_storage:
|
||||
for record in self.inner_storage.values():
|
||||
yield VectorSearchResult(record=record)
|
||||
|
||||
def _lambda_parser(self, node: ast.AST) -> str:
|
||||
return ""
|
||||
|
||||
return DictVectorStoreRecordCollection
|
||||
|
||||
|
||||
@fixture
|
||||
def definition() -> object:
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", dimensions=5, name="vector"),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def data_model_serialize_definition() -> object:
|
||||
def serialize(record, **kwargs):
|
||||
return record
|
||||
|
||||
def deserialize(records, **kwargs):
|
||||
return records
|
||||
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", dimensions=5, name="vector"),
|
||||
],
|
||||
serialize=serialize,
|
||||
deserialize=deserialize,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def data_model_to_from_dict_definition() -> object:
|
||||
def to_dict(record, **kwargs):
|
||||
return record
|
||||
|
||||
def from_dict(records, **kwargs):
|
||||
return records
|
||||
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", dimensions=5, name="vector"),
|
||||
],
|
||||
to_dict=to_dict,
|
||||
from_dict=from_dict,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def data_model_container_definition() -> object:
|
||||
def to_dict(record: dict[str, dict[str, Any]], **kwargs) -> list[dict[str, Any]]:
|
||||
return [{"id": key} | value for key, value in record.items()]
|
||||
|
||||
def from_dict(records: list[dict[str, Any]], **kwargs) -> dict[str, dict[str, Any]]:
|
||||
ret = {}
|
||||
for record in records:
|
||||
id = record.pop("id")
|
||||
ret[id] = record
|
||||
return ret
|
||||
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", dimensions=5, name="vector"),
|
||||
],
|
||||
container_mode=True,
|
||||
to_dict=to_dict,
|
||||
from_dict=from_dict,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def data_model_container_serialize_definition() -> object:
|
||||
def serialize(record: dict[str, dict[str, Any]], **kwargs) -> list[dict[str, Any]]:
|
||||
return [{"id": key} | value for key, value in record.items()]
|
||||
|
||||
def deserialize(records: list[dict[str, Any]], **kwargs) -> dict[str, dict[str, Any]]:
|
||||
ret = {}
|
||||
for record in records:
|
||||
id = record.pop("id")
|
||||
ret[id] = record
|
||||
return ret
|
||||
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", dimensions=5, name="vector"),
|
||||
],
|
||||
container_mode=True,
|
||||
serialize=serialize,
|
||||
deserialize=deserialize,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def data_model_pandas_definition() -> object:
|
||||
from pandas import DataFrame
|
||||
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
name="vector",
|
||||
index_kind="hnsw",
|
||||
dimensions=5,
|
||||
distance_function="cosine_similarity",
|
||||
type="float",
|
||||
),
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField(
|
||||
"data",
|
||||
name="content",
|
||||
type="str",
|
||||
),
|
||||
],
|
||||
container_mode=True,
|
||||
to_dict=lambda x: x.to_dict(orient="records"),
|
||||
from_dict=lambda x, **_: DataFrame(x),
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
async def pandas_vector_store_record_collection(DictVectorStoreRecordCollection, data_model_pandas_definition):
|
||||
from pandas import DataFrame
|
||||
|
||||
return DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=DataFrame,
|
||||
definition=data_model_pandas_definition,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def record_type_vanilla():
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
def __init__(
|
||||
self,
|
||||
content: Annotated[str, VectorStoreField("data")],
|
||||
id: Annotated[str, VectorStoreField("key")],
|
||||
vector: Annotated[list[float] | str | None, VectorStoreField("vector", dimensions=5)] = None,
|
||||
):
|
||||
self.content = content
|
||||
self.vector = vector
|
||||
self.id = id
|
||||
|
||||
def __eq__(self, other) -> bool:
|
||||
return self.content == other.content and self.id == other.id and self.vector == other.vector
|
||||
|
||||
return DataModelClass
|
||||
|
||||
|
||||
@fixture
|
||||
def record_type_vector_array():
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
def __init__(
|
||||
self,
|
||||
id: Annotated[str, VectorStoreField("key")],
|
||||
content: Annotated[str, VectorStoreField("data")],
|
||||
vector: Annotated[
|
||||
list[float] | str | None,
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
dimensions=5,
|
||||
),
|
||||
] = None,
|
||||
):
|
||||
self.content = content
|
||||
self.vector = vector
|
||||
self.id = id
|
||||
|
||||
def __eq__(self, other) -> bool:
|
||||
return self.content == other.content and self.id == other.id and self.vector == other.vector
|
||||
|
||||
return DataModelClass
|
||||
|
||||
|
||||
@fixture
|
||||
def record_type_vanilla_serialize():
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
def __init__(
|
||||
self,
|
||||
id: Annotated[str, VectorStoreField("key")],
|
||||
content: Annotated[str, VectorStoreField("data")],
|
||||
vector: Annotated[list[float] | str | None, VectorStoreField("vector", dimensions=5)] = None,
|
||||
):
|
||||
self.content = content
|
||||
self.vector = vector
|
||||
self.id = id
|
||||
|
||||
def serialize(self, **kwargs: Any) -> Any:
|
||||
"""Serialize the object to the format required by the data store."""
|
||||
return {"id": self.id, "content": self.content, "vector": self.vector}
|
||||
|
||||
@classmethod
|
||||
def deserialize(cls, obj: Any, **kwargs: Any):
|
||||
"""Deserialize the output of the data store to an object."""
|
||||
return cls(**obj)
|
||||
|
||||
def __eq__(self, other) -> bool:
|
||||
return self.content == other.content and self.id == other.id and self.vector == other.vector
|
||||
|
||||
return DataModelClass
|
||||
|
||||
|
||||
@fixture
|
||||
def record_type_vanilla_to_from_dict():
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
def __init__(
|
||||
self,
|
||||
id: Annotated[str, VectorStoreField("key")],
|
||||
content: Annotated[str, VectorStoreField("data")],
|
||||
vector: Annotated[str | list[float] | None, VectorStoreField("vector", dimensions=5)] = None,
|
||||
):
|
||||
self.content = content
|
||||
self.vector = vector
|
||||
self.id = id
|
||||
|
||||
def to_dict(self, **kwargs: Any) -> Any:
|
||||
"""Serialize the object to the format required by the data store."""
|
||||
return {"id": self.id, "content": self.content, "vector": self.vector}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, *args: Any, **kwargs: Any):
|
||||
"""Deserialize the output of the data store to an object."""
|
||||
return cls(**args[0])
|
||||
|
||||
def __eq__(self, other) -> bool:
|
||||
return self.content == other.content and self.id == other.id and self.vector == other.vector
|
||||
|
||||
return DataModelClass
|
||||
|
||||
|
||||
@fixture
|
||||
def record_type_pydantic():
|
||||
@vectorstoremodel
|
||||
class DataModelClass(BaseModel):
|
||||
content: Annotated[str, VectorStoreField("data")]
|
||||
id: Annotated[str, VectorStoreField("key")]
|
||||
vector: Annotated[str | list[float] | None, VectorStoreField("vector", dimensions=5)] = None
|
||||
|
||||
return DataModelClass
|
||||
|
||||
|
||||
@fixture
|
||||
def record_type_dataclass():
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class DataModelClass:
|
||||
content: Annotated[str, VectorStoreField("data")]
|
||||
id: Annotated[str, VectorStoreField("key")]
|
||||
vector: Annotated[list[float] | str | None, VectorStoreField("vector", dimensions=5)] = None
|
||||
|
||||
return DataModelClass
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def vector_store_record_collection(
|
||||
DictVectorStoreRecordCollection,
|
||||
definition,
|
||||
data_model_serialize_definition,
|
||||
data_model_to_from_dict_definition,
|
||||
data_model_container_definition,
|
||||
data_model_container_serialize_definition,
|
||||
data_model_pandas_definition,
|
||||
record_type_vanilla,
|
||||
record_type_vanilla_serialize,
|
||||
record_type_vanilla_to_from_dict,
|
||||
record_type_pydantic,
|
||||
record_type_dataclass,
|
||||
record_type_vector_array,
|
||||
request,
|
||||
) -> VectorSearch:
|
||||
item = request.param if request and hasattr(request, "param") else "definition_basic"
|
||||
defs = {
|
||||
"definition_basic": definition,
|
||||
"definition_with_serialize": data_model_serialize_definition,
|
||||
"definition_with_to_from": data_model_to_from_dict_definition,
|
||||
"definition_container": data_model_container_definition,
|
||||
"definition_container_serialize": data_model_container_serialize_definition,
|
||||
"definition_pandas": data_model_pandas_definition,
|
||||
"type_vanilla": record_type_vanilla,
|
||||
"type_vanilla_with_serialize": record_type_vanilla_serialize,
|
||||
"type_vanilla_with_to_from_dict": record_type_vanilla_to_from_dict,
|
||||
"type_pydantic": record_type_pydantic,
|
||||
"type_dataclass": record_type_dataclass,
|
||||
"type_vector_array": record_type_vector_array,
|
||||
}
|
||||
if item.endswith("pandas"):
|
||||
return DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=DataFrame,
|
||||
definition=defs[item],
|
||||
)
|
||||
if item.startswith("definition_"):
|
||||
return DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=defs[item],
|
||||
)
|
||||
return DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=defs[item],
|
||||
)
|
||||
@@ -0,0 +1,227 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from semantic_kernel import Kernel
|
||||
from semantic_kernel.data.text_search import (
|
||||
DEFAULT_DESCRIPTION,
|
||||
DEFAULT_FUNCTION_NAME,
|
||||
KernelSearchResults,
|
||||
SearchOptions,
|
||||
TextSearch,
|
||||
TextSearchResult,
|
||||
create_options,
|
||||
)
|
||||
from semantic_kernel.data.vector import VectorSearchOptions
|
||||
from semantic_kernel.exceptions import TextSearchException
|
||||
from semantic_kernel.functions import KernelArguments, KernelParameterMetadata
|
||||
from semantic_kernel.utils.list_handler import desync_list
|
||||
|
||||
|
||||
def test_text_search():
|
||||
search_base_class = TextSearch()
|
||||
assert search_base_class is not None
|
||||
assert search_base_class.options_class == SearchOptions
|
||||
|
||||
|
||||
class TestSearch(TextSearch):
|
||||
async def search(self, **kwargs) -> KernelSearchResults[Any]:
|
||||
"""Test search function."""
|
||||
output_type = kwargs.pop("output_type", str)
|
||||
|
||||
async def generator() -> AsyncGenerator[str | TextSearchResult | Any, None]:
|
||||
yield "test" if output_type is not TextSearchResult else TextSearchResult(value="test")
|
||||
|
||||
return KernelSearchResults(results=generator(), metadata=kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("output_type", [str, TextSearchResult, "Any"])
|
||||
async def test_create_kernel_function(output_type: str, kernel: Kernel):
|
||||
test_search = TestSearch()
|
||||
kernel_function = test_search._create_kernel_function(output_type=output_type)
|
||||
assert kernel_function is not None
|
||||
assert kernel_function.name == DEFAULT_FUNCTION_NAME
|
||||
assert kernel_function.description == DEFAULT_DESCRIPTION
|
||||
assert len(kernel_function.parameters) == 3
|
||||
assert kernel_function.return_parameter == KernelParameterMetadata(
|
||||
name="results",
|
||||
description="The search results.",
|
||||
type="list[str]",
|
||||
type_object=list,
|
||||
is_required=True,
|
||||
)
|
||||
results = await kernel_function.invoke(kernel, KernelArguments(query="query"))
|
||||
assert results is not None
|
||||
assert results.value == (
|
||||
['{"name":null,"value":"test","link":null}'] if output_type is TextSearchResult else ["test"]
|
||||
)
|
||||
|
||||
|
||||
def test_create_kernel_function_fail():
|
||||
test_search = TestSearch()
|
||||
with pytest.raises(TextSearchException):
|
||||
test_search.create_search_function(
|
||||
function_name="search",
|
||||
description="description",
|
||||
output_type="random",
|
||||
parameters=None,
|
||||
return_parameter=None,
|
||||
string_mapper=None,
|
||||
)
|
||||
|
||||
|
||||
async def test_create_kernel_function_inner(kernel: Kernel):
|
||||
test_search = TestSearch()
|
||||
|
||||
kernel_function = test_search._create_kernel_function(
|
||||
output_type=str,
|
||||
options=None,
|
||||
parameters=[],
|
||||
return_parameter=None,
|
||||
function_name="search",
|
||||
description="description",
|
||||
string_mapper=None,
|
||||
)
|
||||
results = await kernel_function.invoke(kernel, None)
|
||||
assert results is not None
|
||||
assert results.value == ["test"]
|
||||
|
||||
|
||||
async def test_create_kernel_function_inner_with_options(kernel: Kernel):
|
||||
test_search = TestSearch()
|
||||
|
||||
kernel_function = test_search._create_kernel_function(
|
||||
output_type=str,
|
||||
options=SearchOptions(),
|
||||
parameters=[
|
||||
KernelParameterMetadata(
|
||||
name="city",
|
||||
description="The city that you want to search for a hotel in.",
|
||||
type="str",
|
||||
is_required=False,
|
||||
type_object=str,
|
||||
)
|
||||
],
|
||||
return_parameter=None,
|
||||
function_name="search",
|
||||
description="description",
|
||||
string_mapper=None,
|
||||
)
|
||||
results = await kernel_function.invoke(kernel, KernelArguments(city="city"))
|
||||
assert results is not None
|
||||
assert results.value == ["test"]
|
||||
|
||||
|
||||
async def test_create_kernel_function_inner_with_other_options_type(kernel: Kernel):
|
||||
test_search = TestSearch()
|
||||
|
||||
kernel_function = test_search._create_kernel_function(
|
||||
output_type=str,
|
||||
options=VectorSearchOptions(),
|
||||
parameters=[
|
||||
KernelParameterMetadata(
|
||||
name="test_field",
|
||||
description="The test field.",
|
||||
type="str",
|
||||
is_required=False,
|
||||
type_object=str,
|
||||
)
|
||||
],
|
||||
return_parameter=None,
|
||||
function_name="search",
|
||||
description="description",
|
||||
string_mapper=None,
|
||||
)
|
||||
results = await kernel_function.invoke(kernel, KernelArguments(test_field="city"))
|
||||
assert results is not None
|
||||
assert results.value == ["test"]
|
||||
|
||||
|
||||
async def test_create_kernel_function_inner_no_results(kernel: Kernel):
|
||||
test_search = TestSearch()
|
||||
|
||||
kernel_function = test_search._create_kernel_function(
|
||||
output_type=str,
|
||||
options=None,
|
||||
parameters=[],
|
||||
return_parameter=None,
|
||||
function_name="search",
|
||||
description="description",
|
||||
string_mapper=None,
|
||||
)
|
||||
with (
|
||||
patch.object(test_search, "search") as mock_search,
|
||||
pytest.raises(TextSearchException),
|
||||
):
|
||||
mock_search.side_effect = Exception("fail")
|
||||
await kernel_function.invoke(kernel, None)
|
||||
|
||||
|
||||
async def test_default_map_to_string():
|
||||
test_search = TestSearch()
|
||||
assert (await test_search._map_results(results=KernelSearchResults(results=desync_list(["test"])))) == ["test"]
|
||||
|
||||
class TestClass(BaseModel):
|
||||
test: str
|
||||
|
||||
assert (
|
||||
await test_search._map_results(results=KernelSearchResults(results=desync_list([TestClass(test="test")])))
|
||||
) == ['{"test":"test"}']
|
||||
|
||||
|
||||
async def test_custom_map_to_string():
|
||||
test_search = TestSearch()
|
||||
|
||||
class TestClass(BaseModel):
|
||||
test: str
|
||||
|
||||
assert (
|
||||
await test_search._map_results(
|
||||
results=KernelSearchResults(results=desync_list([TestClass(test="test")])), string_mapper=lambda x: x.test
|
||||
)
|
||||
) == ["test"]
|
||||
|
||||
|
||||
def test_create_options():
|
||||
options = SearchOptions()
|
||||
options_class = VectorSearchOptions
|
||||
new_options = create_options(options_class, options, top=1)
|
||||
assert new_options is not None
|
||||
assert isinstance(new_options, options_class)
|
||||
assert new_options.top == 1
|
||||
|
||||
|
||||
def test_create_options_none():
|
||||
options = None
|
||||
options_class = VectorSearchOptions
|
||||
new_options = create_options(options_class, options, top=1)
|
||||
assert new_options is not None
|
||||
assert isinstance(new_options, options_class)
|
||||
assert new_options.top == 1
|
||||
|
||||
|
||||
def test_create_options_from_dict():
|
||||
options = {"skip": 1}
|
||||
options_class = SearchOptions
|
||||
new_options = create_options(options_class, options, top=1) # type: ignore
|
||||
assert new_options is not None
|
||||
assert isinstance(new_options, options_class)
|
||||
assert new_options.top == 1
|
||||
# if a non SearchOptions object is passed in, it should be ignored
|
||||
assert new_options.skip == 0
|
||||
|
||||
|
||||
def test_public_create_functions_search():
|
||||
test_search = TestSearch()
|
||||
function = test_search.create_search_function()
|
||||
assert function is not None
|
||||
assert function.name == "search"
|
||||
assert (
|
||||
function.description == "Perform a search for content related to the specified query and return string results"
|
||||
)
|
||||
assert len(function.parameters) == 3
|
||||
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.data.vector import VectorSearch, VectorSearchOptions, VectorSearchProtocol
|
||||
|
||||
|
||||
async def test_search(vector_store_record_collection: VectorSearch):
|
||||
assert isinstance(vector_store_record_collection, VectorSearchProtocol)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
await vector_store_record_collection.upsert(record)
|
||||
results = await vector_store_record_collection.search(vector=[1.0, 2.0, 3.0])
|
||||
records = [rec async for rec in results.results]
|
||||
assert records[0].record == record
|
||||
|
||||
|
||||
@pytest.mark.parametrize("include_vectors", [True, False])
|
||||
async def test_get_vector_search_results(vector_store_record_collection: VectorSearch, include_vectors: bool):
|
||||
options = VectorSearchOptions(include_vectors=include_vectors)
|
||||
results = [{"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}]
|
||||
async for result in vector_store_record_collection._get_vector_search_results_from_results(
|
||||
results=results, options=options
|
||||
):
|
||||
assert result.record == results[0] if include_vectors else {"id": "test_id", "content": "test_content"}
|
||||
break
|
||||
@@ -0,0 +1,251 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated
|
||||
|
||||
from numpy import ndarray
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from pydantic.dataclasses import dataclass as pydantic_dataclass
|
||||
from pytest import raises
|
||||
|
||||
from semantic_kernel.data.vector import VectorStoreCollectionDefinition, VectorStoreField, vectorstoremodel
|
||||
from semantic_kernel.exceptions import VectorStoreModelException
|
||||
|
||||
|
||||
def get_field(defn, name):
|
||||
return next(f for f in defn.fields if f.name == name)
|
||||
|
||||
|
||||
def test_vanilla():
|
||||
@vectorstoremodel
|
||||
class DataModelClassVanilla:
|
||||
def __init__(
|
||||
self,
|
||||
content: Annotated[str, VectorStoreField("data")],
|
||||
content2: Annotated[str, VectorStoreField("data")],
|
||||
vector: Annotated[list[float], VectorStoreField("vector", dimensions=5)],
|
||||
id: Annotated[str, VectorStoreField("key")],
|
||||
non_vector_store_content: str | None = None,
|
||||
optional_content: Annotated[str | None, VectorStoreField("data")] = None,
|
||||
annotated_content: Annotated[str | None, "description"] = None,
|
||||
):
|
||||
self.content = content
|
||||
self.content2 = content2
|
||||
self.vector = vector
|
||||
self.id = id
|
||||
self.optional_content = optional_content
|
||||
self.non_vector_store_content = non_vector_store_content
|
||||
self.annotated_content = annotated_content
|
||||
|
||||
assert hasattr(DataModelClassVanilla, "__kernel_vectorstoremodel__")
|
||||
assert hasattr(DataModelClassVanilla, "__kernel_vectorstoremodel_definition__")
|
||||
definition: VectorStoreCollectionDefinition = DataModelClassVanilla.__kernel_vectorstoremodel_definition__
|
||||
assert len(definition.fields) == 5
|
||||
assert get_field(definition, "content").name == "content"
|
||||
assert get_field(definition, "content").type_ == "str"
|
||||
assert get_field(definition, "content2").name == "content2"
|
||||
assert get_field(definition, "content2").type_ == "str"
|
||||
assert get_field(definition, "vector").name == "vector"
|
||||
assert get_field(definition, "id").name == "id"
|
||||
assert get_field(definition, "optional_content").name == "optional_content"
|
||||
assert get_field(definition, "optional_content").type_ == "str"
|
||||
assert definition.key_name == "id"
|
||||
assert definition.container_mode is False
|
||||
assert definition.vector_field_names == ["vector"]
|
||||
|
||||
|
||||
def test_vanilla_2():
|
||||
@vectorstoremodel()
|
||||
class DataModelClassVanilla2:
|
||||
def __init__(
|
||||
self,
|
||||
content: Annotated[str, VectorStoreField("data")],
|
||||
id: Annotated[str, VectorStoreField("key")],
|
||||
):
|
||||
self.content = content
|
||||
self.id = id
|
||||
|
||||
assert hasattr(DataModelClassVanilla2, "__kernel_vectorstoremodel__")
|
||||
assert hasattr(DataModelClassVanilla2, "__kernel_vectorstoremodel_definition__")
|
||||
definition: VectorStoreCollectionDefinition = DataModelClassVanilla2.__kernel_vectorstoremodel_definition__
|
||||
assert len(definition.fields) == 2
|
||||
|
||||
|
||||
def test_dataclass():
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class DataModelClassDataclass:
|
||||
content: Annotated[str, VectorStoreField("data")]
|
||||
content2: Annotated[str, VectorStoreField("data")]
|
||||
vector: Annotated[list[float], VectorStoreField("vector", dimensions=5)]
|
||||
id: Annotated[str, VectorStoreField("key")]
|
||||
non_vector_store_content: str | None = None
|
||||
optional_content: Annotated[str | None, VectorStoreField("data")] = None
|
||||
annotated_content: Annotated[str | None, "description"] = None
|
||||
|
||||
assert hasattr(DataModelClassDataclass, "__kernel_vectorstoremodel__")
|
||||
assert hasattr(DataModelClassDataclass, "__kernel_vectorstoremodel_definition__")
|
||||
definition: VectorStoreCollectionDefinition = DataModelClassDataclass.__kernel_vectorstoremodel_definition__
|
||||
assert len(definition.fields) == 5
|
||||
assert get_field(definition, "content").name == "content"
|
||||
assert get_field(definition, "content").type_ == "str"
|
||||
assert get_field(definition, "content2").name == "content2"
|
||||
assert get_field(definition, "content2").type_ == "str"
|
||||
assert get_field(definition, "vector").name == "vector"
|
||||
assert get_field(definition, "id").name == "id"
|
||||
assert get_field(definition, "optional_content").name == "optional_content"
|
||||
assert get_field(definition, "optional_content").type_ == "str"
|
||||
assert definition.key_name == "id"
|
||||
assert definition.container_mode is False
|
||||
assert definition.vector_field_names == ["vector"]
|
||||
|
||||
|
||||
def test_dataclass_inverse_fail():
|
||||
with raises(VectorStoreModelException):
|
||||
|
||||
@dataclass
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
id: Annotated[str, VectorStoreField("key")]
|
||||
content: Annotated[str, VectorStoreField("data")]
|
||||
|
||||
|
||||
def test_pydantic_base_model():
|
||||
@vectorstoremodel
|
||||
class DataModelClassPydantic(BaseModel):
|
||||
content: Annotated[str, VectorStoreField("data")]
|
||||
content2: Annotated[str, VectorStoreField("data")]
|
||||
vector: Annotated[list[float], VectorStoreField("vector", dimensions=5)]
|
||||
id: Annotated[str, VectorStoreField("key")]
|
||||
non_vector_store_content: str | None = None
|
||||
optional_content: Annotated[str | None, VectorStoreField("data")] = None
|
||||
annotated_content: Annotated[str | None, "description"] = None
|
||||
|
||||
assert hasattr(DataModelClassPydantic, "__kernel_vectorstoremodel__")
|
||||
assert hasattr(DataModelClassPydantic, "__kernel_vectorstoremodel_definition__")
|
||||
definition: VectorStoreCollectionDefinition = DataModelClassPydantic.__kernel_vectorstoremodel_definition__
|
||||
assert len(definition.fields) == 5
|
||||
assert get_field(definition, "content").name == "content"
|
||||
assert get_field(definition, "content").type_ == "str"
|
||||
assert get_field(definition, "content2").name == "content2"
|
||||
assert get_field(definition, "content2").type_ == "str"
|
||||
assert get_field(definition, "vector").name == "vector"
|
||||
assert get_field(definition, "id").name == "id"
|
||||
assert get_field(definition, "optional_content").name == "optional_content"
|
||||
assert get_field(definition, "optional_content").type_ == "str"
|
||||
assert definition.key_name == "id"
|
||||
assert definition.container_mode is False
|
||||
assert definition.vector_field_names == ["vector"]
|
||||
|
||||
|
||||
def test_pydantic_dataclass():
|
||||
@vectorstoremodel
|
||||
@pydantic_dataclass
|
||||
class DataModelClassPydanticDataclass:
|
||||
content: Annotated[str, VectorStoreField("data")]
|
||||
content2: Annotated[str, VectorStoreField("data")]
|
||||
vector: Annotated[list[float], VectorStoreField("vector", dimensions=5)]
|
||||
id: Annotated[str, VectorStoreField("key")]
|
||||
non_vector_store_content: str | None = None
|
||||
optional_content: Annotated[str | None, VectorStoreField("data")] = None
|
||||
annotated_content: Annotated[str | None, "description"] = None
|
||||
|
||||
assert hasattr(DataModelClassPydanticDataclass, "__kernel_vectorstoremodel__")
|
||||
assert hasattr(DataModelClassPydanticDataclass, "__kernel_vectorstoremodel_definition__")
|
||||
definition: VectorStoreCollectionDefinition = DataModelClassPydanticDataclass.__kernel_vectorstoremodel_definition__
|
||||
assert len(definition.fields) == 5
|
||||
assert get_field(definition, "content").name == "content"
|
||||
assert get_field(definition, "content").type_ == "str"
|
||||
assert get_field(definition, "content2").name == "content2"
|
||||
assert get_field(definition, "content2").type_ == "str"
|
||||
assert get_field(definition, "vector").name == "vector"
|
||||
assert get_field(definition, "id").name == "id"
|
||||
assert get_field(definition, "optional_content").name == "optional_content"
|
||||
assert get_field(definition, "optional_content").type_ == "str"
|
||||
assert definition.key_name == "id"
|
||||
assert definition.container_mode is False
|
||||
assert definition.vector_field_names == ["vector"]
|
||||
|
||||
|
||||
def test_empty_model():
|
||||
with raises(VectorStoreModelException):
|
||||
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_non_annotated_no_default():
|
||||
with raises(VectorStoreModelException):
|
||||
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
def __init__(self, non_vector_store_content: str):
|
||||
self.non_vector_store_content = non_vector_store_content
|
||||
|
||||
|
||||
def test_annotated_no_vsr_field_no_default():
|
||||
with raises(VectorStoreModelException):
|
||||
|
||||
@vectorstoremodel
|
||||
class DataModelClass:
|
||||
def __init__(
|
||||
self,
|
||||
annotated_content: Annotated[str, "description"],
|
||||
):
|
||||
self.annotated_content = annotated_content
|
||||
|
||||
|
||||
def test_non_vector_list_and_dict():
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class DataModelClassListDict:
|
||||
key: Annotated[str, VectorStoreField("key")]
|
||||
list1: Annotated[list[int], VectorStoreField("data")]
|
||||
list2: Annotated[list[str], VectorStoreField("data")]
|
||||
list3: Annotated[list[str] | None, VectorStoreField("data")]
|
||||
dict1: Annotated[dict[str, int], VectorStoreField("data")]
|
||||
dict2: Annotated[dict[str, str], VectorStoreField("data")]
|
||||
dict3: Annotated[dict[str, str] | None, VectorStoreField("data")]
|
||||
|
||||
assert hasattr(DataModelClassListDict, "__kernel_vectorstoremodel__")
|
||||
assert hasattr(DataModelClassListDict, "__kernel_vectorstoremodel_definition__")
|
||||
definition: VectorStoreCollectionDefinition = DataModelClassListDict.__kernel_vectorstoremodel_definition__
|
||||
assert len(definition.fields) == 7
|
||||
assert get_field(definition, "list1").name == "list1"
|
||||
assert get_field(definition, "list1").type_ == "list[int]"
|
||||
assert get_field(definition, "list2").name == "list2"
|
||||
assert get_field(definition, "list2").type_ == "list[str]"
|
||||
assert get_field(definition, "list3").name == "list3"
|
||||
assert get_field(definition, "list3").type_ == "list[str]"
|
||||
assert get_field(definition, "dict1").name == "dict1"
|
||||
assert get_field(definition, "dict1").type_ == "dict[str, int]"
|
||||
assert get_field(definition, "dict2").name == "dict2"
|
||||
assert get_field(definition, "dict2").type_ == "dict[str, str]"
|
||||
assert get_field(definition, "dict3").name == "dict3"
|
||||
assert get_field(definition, "dict3").type_ == "dict[str, str]"
|
||||
assert definition.container_mode is False
|
||||
|
||||
|
||||
def test_vector_fields_checks():
|
||||
@vectorstoremodel
|
||||
class DataModelClassVectorFields(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
id: Annotated[str, VectorStoreField("key")]
|
||||
vector_str: Annotated[str, VectorStoreField("vector", dimensions=5)]
|
||||
vector_list: Annotated[list[float], VectorStoreField("vector", dimensions=5)]
|
||||
vector_array: Annotated[
|
||||
ndarray,
|
||||
VectorStoreField("vector", dimensions=5),
|
||||
]
|
||||
|
||||
assert hasattr(DataModelClassVectorFields, "__kernel_vectorstoremodel__")
|
||||
assert hasattr(DataModelClassVectorFields, "__kernel_vectorstoremodel_definition__")
|
||||
definition: VectorStoreCollectionDefinition = DataModelClassVectorFields.__kernel_vectorstoremodel_definition__
|
||||
assert len(definition.fields) == 4
|
||||
assert get_field(definition, "id").name == "id"
|
||||
assert get_field(definition, "vector_str").type_ == "str"
|
||||
assert get_field(definition, "vector_list").type_ == "float"
|
||||
assert get_field(definition, "vector_array").type_ == "ndarray"
|
||||
@@ -0,0 +1,433 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from copy import deepcopy
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, PropertyMock, patch
|
||||
|
||||
from pandas import DataFrame
|
||||
from pytest import mark, raises
|
||||
|
||||
from semantic_kernel.data.vector import SerializeMethodProtocol, ToDictMethodProtocol
|
||||
from semantic_kernel.exceptions import (
|
||||
VectorStoreModelDeserializationException,
|
||||
VectorStoreModelSerializationException,
|
||||
VectorStoreModelValidationError,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
|
||||
|
||||
# region init
|
||||
def test_init(DictVectorStoreRecordCollection, definition):
|
||||
vsrc = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
assert vsrc.collection_name == "test"
|
||||
assert vsrc.record_type is dict
|
||||
assert vsrc._container_mode is False
|
||||
assert vsrc.definition == definition
|
||||
assert vsrc._key_field_name == "id"
|
||||
|
||||
|
||||
def test_data_model_validation(record_type_vanilla, DictVectorStoreRecordCollection):
|
||||
DictVectorStoreRecordCollection.supported_key_types = PropertyMock(return_value=["str"])
|
||||
DictVectorStoreRecordCollection.supported_vector_types = PropertyMock(return_value=["float"])
|
||||
DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=record_type_vanilla,
|
||||
)
|
||||
|
||||
|
||||
def test_data_model_validation_key_fail(record_type_vanilla, DictVectorStoreRecordCollection):
|
||||
DictVectorStoreRecordCollection.supported_key_types = PropertyMock(return_value=["int"])
|
||||
with raises(VectorStoreModelValidationError, match="Key field must be one of"):
|
||||
DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=record_type_vanilla,
|
||||
)
|
||||
|
||||
|
||||
def test_data_model_validation_vector_fail(record_type_vanilla, DictVectorStoreRecordCollection):
|
||||
DictVectorStoreRecordCollection.supported_vector_types = PropertyMock(return_value={"list[int]"})
|
||||
with raises(VectorStoreModelValidationError):
|
||||
DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=record_type_vanilla,
|
||||
)
|
||||
|
||||
|
||||
# region Collection
|
||||
async def test_collection_operations(vector_store_record_collection):
|
||||
await vector_store_record_collection.ensure_collection_exists()
|
||||
assert await vector_store_record_collection.collection_exists()
|
||||
record = {"id": "id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
await vector_store_record_collection.upsert(record)
|
||||
assert len(vector_store_record_collection.inner_storage) == 1
|
||||
await vector_store_record_collection.ensure_collection_deleted()
|
||||
assert vector_store_record_collection.inner_storage == {}
|
||||
await vector_store_record_collection.ensure_collection_exists()
|
||||
|
||||
|
||||
async def test_collection_create_if_not_exists(DictVectorStoreRecordCollection, definition):
|
||||
DictVectorStoreRecordCollection.collection_exists = AsyncMock(return_value=False)
|
||||
create_mock = AsyncMock()
|
||||
DictVectorStoreRecordCollection.ensure_collection_exists = create_mock
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
await vector_store_record_collection.ensure_collection_exists()
|
||||
create_mock.assert_called_once()
|
||||
|
||||
|
||||
# region CRUD
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"vector_store_record_collection",
|
||||
[
|
||||
"definition_basic",
|
||||
"definition_with_serialize",
|
||||
"definition_with_to_from",
|
||||
"type_vanilla",
|
||||
"type_vanilla_with_serialize",
|
||||
"type_vanilla_with_to_from_dict",
|
||||
"type_pydantic",
|
||||
"type_dataclass",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_crud_operations(vector_store_record_collection):
|
||||
ids = ["test_id_1", "test_id_2"]
|
||||
batch = [
|
||||
{"id": ids[0], "content": "test_content", "vector": [1.0, 2.0, 3.0]},
|
||||
{"id": ids[1], "content": "test_content", "vector": [1.0, 2.0, 3.0]},
|
||||
]
|
||||
if vector_store_record_collection.record_type is not dict:
|
||||
model = vector_store_record_collection.record_type
|
||||
batch = [model(**record) for record in batch]
|
||||
no_records = await vector_store_record_collection.get(ids)
|
||||
assert no_records is None
|
||||
await vector_store_record_collection.upsert(batch)
|
||||
assert len(vector_store_record_collection.inner_storage) == 2
|
||||
if vector_store_record_collection.record_type is dict:
|
||||
assert vector_store_record_collection.inner_storage[ids[0]] == batch[0]
|
||||
else:
|
||||
assert not isinstance(batch[0], dict)
|
||||
assert vector_store_record_collection.inner_storage[ids[0]]["content"] == batch[0].content
|
||||
records = await vector_store_record_collection.get(ids, include_vector=True)
|
||||
for idx, rec in enumerate(records):
|
||||
if vector_store_record_collection.record_type is dict:
|
||||
assert rec["id"] == batch[idx]["id"]
|
||||
assert rec["content"] == batch[idx]["content"]
|
||||
else:
|
||||
assert rec.id == batch[idx].id
|
||||
assert rec.content == batch[idx].content
|
||||
await vector_store_record_collection.delete(ids)
|
||||
assert len(vector_store_record_collection.inner_storage) == 0
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"vector_store_record_collection",
|
||||
["definition_container", "definition_container_serialize"],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_crud_operations_container(vector_store_record_collection):
|
||||
id = "test_id"
|
||||
record = {id: {"content": "test_content", "vector": [1.0, 2.0, 3.0]}}
|
||||
no_records = await vector_store_record_collection.get(id)
|
||||
assert no_records is None
|
||||
await vector_store_record_collection.upsert(record)
|
||||
assert len(vector_store_record_collection.inner_storage) == 1
|
||||
assert vector_store_record_collection.inner_storage[id]["content"] == record[id]["content"]
|
||||
assert vector_store_record_collection.inner_storage[id]["vector"] == record[id]["vector"]
|
||||
record_2 = await vector_store_record_collection.get(id)
|
||||
record_2["id"] = id
|
||||
record_2["content"] = record_2[id]["content"]
|
||||
await vector_store_record_collection.delete(id)
|
||||
assert len(vector_store_record_collection.inner_storage) == 0
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"vector_store_record_collection",
|
||||
["definition_container", "definition_container_serialize"],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_crud_batch_operations_container(vector_store_record_collection):
|
||||
ids = ["test_id_1", "test_id_2"]
|
||||
batch = {
|
||||
ids[0]: {"content": "test_content", "vector": [1.0, 2.0, 3.0]},
|
||||
ids[1]: {"content": "test_content", "vector": [1.0, 2.0, 3.0]},
|
||||
}
|
||||
no_records = await vector_store_record_collection.get(ids)
|
||||
assert no_records is None
|
||||
await vector_store_record_collection.upsert(batch)
|
||||
assert len(vector_store_record_collection.inner_storage) == 2
|
||||
assert vector_store_record_collection.inner_storage[ids[0]]["content"] == batch[ids[0]]["content"]
|
||||
assert vector_store_record_collection.inner_storage[ids[0]]["vector"] == batch[ids[0]]["vector"]
|
||||
records = await vector_store_record_collection.get(ids)
|
||||
assert records == batch
|
||||
await vector_store_record_collection.delete(ids)
|
||||
assert len(vector_store_record_collection.inner_storage) == 0
|
||||
|
||||
|
||||
async def test_crud_operations_pandas(pandas_vector_store_record_collection):
|
||||
id = "test_id"
|
||||
record = DataFrame([{"id": id, "content": "test_content", "vector": [1.0, 2.0, 3.0]}])
|
||||
no_records = await pandas_vector_store_record_collection.get(id)
|
||||
assert no_records is None
|
||||
await pandas_vector_store_record_collection.upsert(record)
|
||||
assert len(pandas_vector_store_record_collection.inner_storage) == 1
|
||||
|
||||
assert pandas_vector_store_record_collection.inner_storage[id]["content"] == record["content"].values[0]
|
||||
assert pandas_vector_store_record_collection.inner_storage[id]["vector"] == record["vector"].values[0]
|
||||
record_2 = await pandas_vector_store_record_collection.get(id)
|
||||
assert record_2.equals(record)
|
||||
await pandas_vector_store_record_collection.delete(id)
|
||||
assert len(pandas_vector_store_record_collection.inner_storage) == 0
|
||||
|
||||
|
||||
async def test_crud_batch_operations_pandas(pandas_vector_store_record_collection):
|
||||
ids = ["test_id_1", "test_id_2"]
|
||||
|
||||
batch = DataFrame([{"id": id, "content": "test_content", "vector": [1.0, 2.0, 3.0]} for id in ids])
|
||||
no_records = await pandas_vector_store_record_collection.get(ids)
|
||||
assert no_records is None
|
||||
await pandas_vector_store_record_collection.upsert(batch)
|
||||
assert len(pandas_vector_store_record_collection.inner_storage) == 2
|
||||
assert pandas_vector_store_record_collection.inner_storage[ids[0]]["content"] == batch["content"].values[0]
|
||||
assert pandas_vector_store_record_collection.inner_storage[ids[0]]["vector"] == batch["vector"].values[0]
|
||||
records = await pandas_vector_store_record_collection.get(ids)
|
||||
assert records.equals(batch)
|
||||
await pandas_vector_store_record_collection.delete(ids)
|
||||
assert len(pandas_vector_store_record_collection.inner_storage) == 0
|
||||
|
||||
|
||||
# region Fails
|
||||
async def test_upsert_fail(DictVectorStoreRecordCollection, definition):
|
||||
DictVectorStoreRecordCollection._inner_upsert = MagicMock(side_effect=Exception)
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
with raises(VectorStoreOperationException, match="Error upserting record"):
|
||||
await vector_store_record_collection.upsert(record)
|
||||
with raises(VectorStoreOperationException, match="Error upserting record"):
|
||||
await vector_store_record_collection.upsert([record])
|
||||
with raises(VectorStoreOperationException, match="Error upserting record"):
|
||||
await vector_store_record_collection.upsert([record])
|
||||
assert len(vector_store_record_collection.inner_storage) == 0
|
||||
|
||||
|
||||
async def test_get_fail(DictVectorStoreRecordCollection, definition):
|
||||
DictVectorStoreRecordCollection._inner_get = MagicMock(side_effect=Exception)
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
await vector_store_record_collection.upsert(record)
|
||||
assert len(vector_store_record_collection.inner_storage) == 1
|
||||
with raises(VectorStoreOperationException, match="Error getting record"):
|
||||
await vector_store_record_collection.get("test_id")
|
||||
with raises(VectorStoreOperationException, match="Error getting record"):
|
||||
await vector_store_record_collection.get(["test_id"])
|
||||
with raises(VectorStoreOperationException, match="Error getting record"):
|
||||
await vector_store_record_collection.get(["test_id"])
|
||||
|
||||
|
||||
async def test_deserialize_in_get_fail(DictVectorStoreRecordCollection, definition):
|
||||
definition.deserialize = MagicMock(side_effect=Exception)
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
await vector_store_record_collection.upsert(record)
|
||||
assert len(vector_store_record_collection.inner_storage) == 1
|
||||
with raises(VectorStoreModelDeserializationException, match="Error deserializing records:"):
|
||||
await vector_store_record_collection.get("test_id")
|
||||
with raises(VectorStoreModelDeserializationException, match="Error deserializing records:"):
|
||||
await vector_store_record_collection.get(["test_id"])
|
||||
|
||||
|
||||
async def test_get_fail_multiple(DictVectorStoreRecordCollection, definition):
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
await vector_store_record_collection.upsert(record)
|
||||
assert len(vector_store_record_collection.inner_storage) == 1
|
||||
with (
|
||||
patch("semantic_kernel.data.vector.VectorStoreCollection.deserialize") as deserialize_mock,
|
||||
raises(
|
||||
VectorStoreModelDeserializationException, match="Error deserializing record, multiple records returned:"
|
||||
),
|
||||
):
|
||||
deserialize_mock.return_value = [
|
||||
{"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]},
|
||||
{"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]},
|
||||
]
|
||||
await vector_store_record_collection.get("test_id")
|
||||
|
||||
|
||||
async def test_delete_fail(DictVectorStoreRecordCollection, definition):
|
||||
DictVectorStoreRecordCollection._inner_delete = MagicMock(side_effect=Exception)
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
await vector_store_record_collection.upsert(record)
|
||||
assert len(vector_store_record_collection.inner_storage) == 1
|
||||
with raises(VectorStoreOperationException, match="Error deleting record"):
|
||||
await vector_store_record_collection.delete("test_id")
|
||||
with raises(VectorStoreOperationException, match="Error deleting record"):
|
||||
await vector_store_record_collection.delete(["test_id"])
|
||||
with raises(VectorStoreOperationException, match="Error deleting record"):
|
||||
await vector_store_record_collection.delete(["test_id"])
|
||||
assert len(vector_store_record_collection.inner_storage) == 1
|
||||
|
||||
|
||||
# region Serialize
|
||||
async def test_serialize_in_upsert_fail(DictVectorStoreRecordCollection, definition):
|
||||
DictVectorStoreRecordCollection.serialize = MagicMock(side_effect=VectorStoreModelSerializationException)
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
with raises(VectorStoreModelSerializationException):
|
||||
await vector_store_record_collection.upsert(record)
|
||||
with raises(VectorStoreModelSerializationException):
|
||||
await vector_store_record_collection.upsert([record])
|
||||
|
||||
|
||||
async def test_serialize_record_type_serialize_fail(DictVectorStoreRecordCollection, record_type_vanilla_serialize):
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=record_type_vanilla_serialize,
|
||||
)
|
||||
record = MagicMock(spec=SerializeMethodProtocol)
|
||||
record.serialize = MagicMock(side_effect=Exception)
|
||||
with raises(VectorStoreModelSerializationException, match="Error serializing record"):
|
||||
await vector_store_record_collection.serialize(record)
|
||||
|
||||
|
||||
def test_serialize_data_model_to_dict_fail_object(DictVectorStoreRecordCollection, record_type_vanilla):
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=record_type_vanilla,
|
||||
)
|
||||
record = Mock(spec=record_type_vanilla)
|
||||
with raises(AttributeError):
|
||||
vector_store_record_collection._serialize_data_model_to_dict(record)
|
||||
|
||||
|
||||
@mark.parametrize("vector_store_record_collection", ["type_pydantic"], indirect=True)
|
||||
async def test_pydantic_serialize_fail(vector_store_record_collection):
|
||||
id = "test_id"
|
||||
model = deepcopy(vector_store_record_collection.record_type)
|
||||
model.model_dump = MagicMock(side_effect=Exception)
|
||||
vector_store_record_collection.record_type = model
|
||||
dict_record = {"id": id, "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
record = model(**dict_record)
|
||||
with raises(VectorStoreModelSerializationException, match="Error serializing record"):
|
||||
await vector_store_record_collection.serialize(record)
|
||||
|
||||
|
||||
@mark.parametrize("vector_store_record_collection", ["type_vanilla_with_to_from_dict"], indirect=True)
|
||||
async def test_to_dict_fail(vector_store_record_collection):
|
||||
record = MagicMock(spec=ToDictMethodProtocol)
|
||||
record.to_dict = MagicMock(side_effect=Exception)
|
||||
with raises(VectorStoreModelSerializationException, match="Error serializing record"):
|
||||
await vector_store_record_collection.serialize(record)
|
||||
|
||||
|
||||
# region Deserialize
|
||||
|
||||
|
||||
async def test_deserialize_definition_fail(DictVectorStoreRecordCollection, definition):
|
||||
definition.deserialize = MagicMock(side_effect=Exception)
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
vector_store_record_collection.inner_storage["test_id"] = record
|
||||
with raises(VectorStoreModelDeserializationException, match="Error deserializing record"):
|
||||
await vector_store_record_collection.get("test_id")
|
||||
with raises(VectorStoreModelDeserializationException, match="Error deserializing record"):
|
||||
await vector_store_record_collection.get(["test_id"])
|
||||
|
||||
|
||||
async def test_deserialize_definition_none(DictVectorStoreRecordCollection, definition):
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
assert vector_store_record_collection.deserialize([]) is None
|
||||
assert vector_store_record_collection.deserialize({}) is None
|
||||
|
||||
|
||||
def test_deserialize_type_fail(DictVectorStoreRecordCollection, record_type_vanilla_serialize):
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=record_type_vanilla_serialize,
|
||||
)
|
||||
record = {"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
vector_store_record_collection.record_type.deserialize = MagicMock(side_effect=Exception)
|
||||
with raises(VectorStoreModelDeserializationException, match="Error deserializing record"):
|
||||
vector_store_record_collection.deserialize(record)
|
||||
|
||||
|
||||
def test_deserialize_dict_data_model_fail_sequence(DictVectorStoreRecordCollection, record_type_vanilla):
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=record_type_vanilla,
|
||||
)
|
||||
with raises(VectorStoreModelDeserializationException, match="Cannot deserialize multiple records"):
|
||||
vector_store_record_collection._deserialize_dict_to_data_model([{}, {}])
|
||||
|
||||
|
||||
def test_deserialize_dict_data_model_shortcut(DictVectorStoreRecordCollection, definition):
|
||||
vector_store_record_collection = DictVectorStoreRecordCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
record = vector_store_record_collection._deserialize_dict_to_data_model([
|
||||
{"id": "test_id", "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
])
|
||||
assert record == {"id": "test_id", "content": "test_content"}
|
||||
|
||||
|
||||
@mark.parametrize("vector_store_record_collection", ["type_pydantic"], indirect=True)
|
||||
async def test_pydantic_deserialize_fail(vector_store_record_collection):
|
||||
id = "test_id"
|
||||
dict_record = {"id": id, "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
vector_store_record_collection.record_type.model_validate = MagicMock(side_effect=Exception)
|
||||
with raises(VectorStoreModelDeserializationException, match="Error deserializing record"):
|
||||
vector_store_record_collection.deserialize(dict_record)
|
||||
|
||||
|
||||
@mark.parametrize("vector_store_record_collection", ["type_vanilla_with_to_from_dict"], indirect=True)
|
||||
def test_from_dict_fail(vector_store_record_collection):
|
||||
id = "test_id"
|
||||
model = deepcopy(vector_store_record_collection.record_type)
|
||||
dict_record = {"id": id, "content": "test_content", "vector": [1.0, 2.0, 3.0]}
|
||||
model.from_dict = MagicMock(side_effect=Exception)
|
||||
vector_store_record_collection.record_type = model
|
||||
with raises(VectorStoreModelDeserializationException, match="Error deserializing record"):
|
||||
vector_store_record_collection.deserialize(dict_record)
|
||||
@@ -0,0 +1,106 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from pytest import raises
|
||||
|
||||
from semantic_kernel.data.vector import VectorStoreCollectionDefinition, VectorStoreField
|
||||
from semantic_kernel.exceptions import VectorStoreModelException
|
||||
|
||||
|
||||
def test_vector_store_record_definition():
|
||||
id_field = VectorStoreField("key", name="id")
|
||||
vsrd = VectorStoreCollectionDefinition(fields=[id_field])
|
||||
assert vsrd.fields == [VectorStoreField("key", name="id")]
|
||||
assert vsrd.key_name == "id"
|
||||
assert vsrd.key_field == id_field
|
||||
assert vsrd.names == ["id"]
|
||||
assert vsrd.vector_field_names == []
|
||||
assert vsrd.container_mode is False
|
||||
assert vsrd.to_dict is None
|
||||
assert vsrd.from_dict is None
|
||||
assert vsrd.serialize is None
|
||||
assert vsrd.deserialize is None
|
||||
|
||||
|
||||
def test_no_fields_fail():
|
||||
with raises(VectorStoreModelException):
|
||||
VectorStoreCollectionDefinition(fields=[])
|
||||
|
||||
|
||||
def test_no_name_fields_fail():
|
||||
with raises(VectorStoreModelException):
|
||||
VectorStoreCollectionDefinition(fields=[VectorStoreField("key", name=None)])
|
||||
with raises(VectorStoreModelException):
|
||||
VectorStoreCollectionDefinition(fields=[VectorStoreField("key", name="")])
|
||||
|
||||
|
||||
def test_no_key_field_fail():
|
||||
with raises(VectorStoreModelException):
|
||||
VectorStoreCollectionDefinition(fields=[VectorStoreField("data", name="content")])
|
||||
|
||||
|
||||
def test_multiple_key_field_fail():
|
||||
with raises(VectorStoreModelException):
|
||||
VectorStoreCollectionDefinition(
|
||||
fields=[VectorStoreField("key", name="key1"), VectorStoreField("key", name="key2")]
|
||||
)
|
||||
|
||||
|
||||
def test_vector_and_non_vector_field_names():
|
||||
definition = VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", name="vector", dimensions=5),
|
||||
]
|
||||
)
|
||||
assert definition.vector_field_names == ["vector"]
|
||||
assert definition.data_field_names == ["content"]
|
||||
|
||||
|
||||
def test_try_get_vector_field():
|
||||
definition = VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", name="vector", dimensions=5),
|
||||
]
|
||||
)
|
||||
assert definition.try_get_vector_field() == definition.fields[2]
|
||||
assert definition.try_get_vector_field("vector") == definition.fields[2]
|
||||
|
||||
|
||||
def test_try_get_vector_field_none():
|
||||
definition = VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
]
|
||||
)
|
||||
assert definition.try_get_vector_field() is None
|
||||
with raises(VectorStoreModelException, match="Field vector not found."):
|
||||
definition.try_get_vector_field("vector")
|
||||
|
||||
|
||||
def test_try_get_vector_field_wrong_name_fail():
|
||||
definition = VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
]
|
||||
)
|
||||
with raises(VectorStoreModelException, match="Field content is not a vector field."):
|
||||
definition.try_get_vector_field("content")
|
||||
|
||||
|
||||
def test_get_field_names():
|
||||
definition = VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", name="vector", dimensions=5),
|
||||
]
|
||||
)
|
||||
assert definition.get_names() == ["id", "content", "vector"]
|
||||
assert definition.get_names(include_vector_fields=False) == ["id", "content"]
|
||||
assert definition.get_names(include_key_field=False) == ["content", "vector"]
|
||||
assert definition.get_names(include_vector_fields=False, include_key_field=False) == ["content"]
|
||||
Reference in New Issue
Block a user