wehub-resource-sync
27 lines
1.2 KiB
Python
27 lines
1.2 KiB
Python
# 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
|