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],
|
||||
)
|
||||
Reference in New Issue
Block a user