mirror of
https://github.com/microsoft/semantic-kernel.git
synced 2026-10-12 02:19:06 +00:00
Add support for weaviate as a memory store (Python) (#804)
Add support for weaviate as a memory store see #803 Co-authored-by: Abby Harrison <[email protected]> Co-authored-by: Devis Lucato <[email protected]> Co-authored-by: Abby Harrison <[email protected]>
This commit is contained in:
Shukri
co-authored by
Abby Harrison
Devis Lucato
Abby Harrison
Abby Harrison
Devis Lucato
Abby Harrison
parent
a5144ba6ac
commit
18cd0cd042
Generated
+481
-393
File diff suppressed because it is too large
Load Diff
@@ -29,6 +29,10 @@ torch = "2.0.0"
|
||||
[tool.poetry.group.chromadb.dependencies]
|
||||
chromadb = "^0.3.23"
|
||||
|
||||
|
||||
[tool.poetry.group.weaviate.dependencies]
|
||||
weaviate-client = "^3.18.0"
|
||||
|
||||
[tool.isort]
|
||||
profile = "black"
|
||||
|
||||
|
||||
@@ -0,0 +1,318 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from logging import Logger
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import weaviate
|
||||
from weaviate.embedded import EmbeddedOptions
|
||||
|
||||
from semantic_kernel.memory.memory_record import MemoryRecord
|
||||
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
|
||||
from semantic_kernel.utils.null_logger import NullLogger
|
||||
|
||||
SCHEMA = {
|
||||
"class": "MemoryRecord",
|
||||
"description": "A document from semantic kernel.",
|
||||
"properties": [
|
||||
{
|
||||
"name": "key",
|
||||
"description": "The key of the record.",
|
||||
"dataType": ["string"],
|
||||
},
|
||||
{
|
||||
"name": "timestamp",
|
||||
"description": "The timestamp of the record.",
|
||||
"dataType": ["date"],
|
||||
},
|
||||
{
|
||||
"name": "isReference",
|
||||
"description": "Whether the record is a reference record.",
|
||||
"dataType": ["boolean"],
|
||||
},
|
||||
{
|
||||
"name": "externalSourceName",
|
||||
"description": "The name of the external source.",
|
||||
"dataType": ["string"],
|
||||
},
|
||||
{
|
||||
"name": "skId",
|
||||
"description": "A unique identifier for the record.",
|
||||
"dataType": ["string"],
|
||||
},
|
||||
{
|
||||
"name": "description",
|
||||
"description": "The description of the record.",
|
||||
"dataType": ["text"],
|
||||
},
|
||||
{
|
||||
"name": "text",
|
||||
"description": "The text of the record.",
|
||||
"dataType": ["text"],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
ALL_PROPERTIES = [property["name"] for property in SCHEMA["properties"]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class WeaviateConfig:
|
||||
use_embed: bool = False
|
||||
url: str = None
|
||||
api_key: str = None
|
||||
|
||||
|
||||
class WeaviateMemoryStore(MemoryStoreBase):
|
||||
class FieldMapper:
|
||||
"""
|
||||
This inner class is responsible for mapping attribute names between
|
||||
the SK's memory record and weaviate's schema. It provides methods
|
||||
for converting between the two naming conventions.
|
||||
"""
|
||||
|
||||
SK_TO_WEAVIATE_MAPPING = {
|
||||
"_key": "key",
|
||||
"_timestamp": "timestamp",
|
||||
"_is_reference": "isReference",
|
||||
"_external_source_name": "externalSourceName",
|
||||
"_id": "skId",
|
||||
"_description": "description",
|
||||
"_text": "text",
|
||||
"_embedding": "vector",
|
||||
}
|
||||
|
||||
WEAVIATE_TO_SK_MAPPING = {v: k for k, v in SK_TO_WEAVIATE_MAPPING.items()}
|
||||
|
||||
@classmethod
|
||||
def sk_to_weaviate(cls, sk_dict):
|
||||
return {
|
||||
cls.SK_TO_WEAVIATE_MAPPING.get(k, k): v
|
||||
for k, v in sk_dict.items()
|
||||
if k in cls.SK_TO_WEAVIATE_MAPPING
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def weaviate_to_sk(cls, weaviate_dict):
|
||||
return {
|
||||
cls.WEAVIATE_TO_SK_MAPPING.get(k, k): v
|
||||
for k, v in weaviate_dict.items()
|
||||
if k in cls.WEAVIATE_TO_SK_MAPPING
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def remove_underscore_prefix(cls, sk_dict):
|
||||
"""
|
||||
Used to initialize a MemoryRecord from a SK's dict of private attribute-values.
|
||||
"""
|
||||
return {key.lstrip("_"): value for key, value in sk_dict.items()}
|
||||
|
||||
def __init__(self, config: WeaviateConfig, logger: Optional[Logger] = None):
|
||||
self._logger = logger or NullLogger()
|
||||
self.config = config
|
||||
self.client = self._initialize_client()
|
||||
|
||||
def _initialize_client(self):
|
||||
if self.config.use_embed:
|
||||
return weaviate.Client(embedded_options=EmbeddedOptions())
|
||||
elif self.config.url:
|
||||
if self.config.api_key:
|
||||
return weaviate.Client(
|
||||
url=self.config.url,
|
||||
auth_client_secret=weaviate.auth.AuthApiKey(
|
||||
api_key=self.config.api_key
|
||||
),
|
||||
)
|
||||
else:
|
||||
return weaviate.Client(url=self.config.url)
|
||||
else:
|
||||
raise ValueError("Weaviate config must have either url or use_embed set")
|
||||
|
||||
async def create_collection_async(self, collection_name: str) -> None:
|
||||
schema = SCHEMA.copy()
|
||||
schema["class"] = collection_name
|
||||
await asyncio.to_thread(self.client.schema.create_class, schema)
|
||||
|
||||
async def get_collections_async(self) -> List[str]:
|
||||
schemas = await asyncio.to_thread(self.client.schema.get)
|
||||
return [schema["class"] for schema in schemas["classes"]]
|
||||
|
||||
async def delete_collection_async(self, collection_name: str) -> bool:
|
||||
await asyncio.to_thread(self.client.schema.delete_class, collection_name)
|
||||
|
||||
async def does_collection_exist_async(self, collection_name: str) -> bool:
|
||||
collections = await self.get_collections_async()
|
||||
return collection_name in collections
|
||||
|
||||
async def upsert_async(self, collection_name: str, record: MemoryRecord) -> str:
|
||||
weaviate_record = self.FieldMapper.sk_to_weaviate(vars(record))
|
||||
|
||||
vector = weaviate_record.pop("vector", None)
|
||||
weaviate_id = weaviate.util.generate_uuid5(weaviate_record, collection_name)
|
||||
|
||||
return await asyncio.to_thread(
|
||||
self.client.data_object.create,
|
||||
data_object=weaviate_record,
|
||||
uuid=weaviate_id,
|
||||
vector=vector,
|
||||
class_name=collection_name,
|
||||
)
|
||||
|
||||
async def upsert_batch_async(
|
||||
self, collection_name: str, records: List[MemoryRecord]
|
||||
) -> List[str]:
|
||||
def _upsert_batch_inner():
|
||||
results = []
|
||||
with self.client.batch as batch:
|
||||
for record in records:
|
||||
weaviate_record = self.FieldMapper.sk_to_weaviate(vars(record))
|
||||
vector = weaviate_record.pop("vector", None)
|
||||
weaviate_id = weaviate.util.generate_uuid5(
|
||||
weaviate_record, collection_name
|
||||
)
|
||||
batch.add_data_object(
|
||||
data_object=weaviate_record,
|
||||
uuid=weaviate_id,
|
||||
vector=vector,
|
||||
class_name=collection_name,
|
||||
)
|
||||
results.append(weaviate_id)
|
||||
|
||||
return results
|
||||
|
||||
return await asyncio.to_thread(_upsert_batch_inner)
|
||||
|
||||
async def get_async(
|
||||
self, collection_name: str, key: str, with_embedding: bool
|
||||
) -> MemoryRecord:
|
||||
# Call the batched version with a single key
|
||||
results = await self.get_batch_async(collection_name, [key], with_embedding)
|
||||
return results[0] if results else None
|
||||
|
||||
async def get_batch_async(
|
||||
self, collection_name: str, keys: List[str], with_embedding: bool
|
||||
) -> List[MemoryRecord]:
|
||||
queries = self._build_multi_get_query(collection_name, keys, with_embedding)
|
||||
|
||||
results = await asyncio.to_thread(self.client.query.multi_get(queries).do)
|
||||
|
||||
get_dict = results.get("data", {}).get("Get", {})
|
||||
|
||||
memory_records = [
|
||||
self._convert_weaviate_doc_to_memory_record(doc)
|
||||
for docs in get_dict.values()
|
||||
for doc in docs
|
||||
]
|
||||
|
||||
return memory_records
|
||||
|
||||
def _build_multi_get_query(
|
||||
self, collection_name: str, keys: List[str], with_embedding: bool
|
||||
):
|
||||
queries = []
|
||||
for i, key in enumerate(keys):
|
||||
query = self.client.query.get(collection_name, ALL_PROPERTIES).with_where(
|
||||
{
|
||||
"path": ["key"],
|
||||
"operator": "Equal",
|
||||
"valueString": key,
|
||||
}
|
||||
)
|
||||
if with_embedding:
|
||||
query = query.with_additional("vector")
|
||||
|
||||
query = query.with_alias(f"query_{i}")
|
||||
|
||||
queries.append(query)
|
||||
|
||||
return queries
|
||||
|
||||
def _convert_weaviate_doc_to_memory_record(
|
||||
self, weaviate_doc: dict
|
||||
) -> MemoryRecord:
|
||||
weaviate_doc_copy = weaviate_doc.copy()
|
||||
vector = weaviate_doc_copy.pop("_additional", {}).get("vector")
|
||||
weaviate_doc_copy["vector"] = np.array(vector) if vector else None
|
||||
sk_doc = self.FieldMapper.weaviate_to_sk(weaviate_doc_copy)
|
||||
mem_vals = self.FieldMapper.remove_underscore_prefix(sk_doc)
|
||||
return MemoryRecord(**mem_vals)
|
||||
|
||||
async def remove_async(self, collection_name: str, key: str) -> None:
|
||||
await self.remove_batch_async(collection_name, [key])
|
||||
|
||||
async def remove_batch_async(self, collection_name: str, keys: List[str]) -> None:
|
||||
# TODO: Use In operator when it's available
|
||||
# (https://github.com/weaviate/weaviate/issues/2387)
|
||||
# and handle max delete objects
|
||||
# (https://weaviate.io/developers/weaviate/api/rest/batch#maximum-number-of-deletes-per-query)
|
||||
for key in keys:
|
||||
where = {
|
||||
"path": ["key"],
|
||||
"operator": "Equal",
|
||||
"valueString": key,
|
||||
}
|
||||
|
||||
await asyncio.to_thread(
|
||||
self.client.batch.delete_objects,
|
||||
class_name=collection_name,
|
||||
where=where,
|
||||
)
|
||||
|
||||
async def get_nearest_matches_async(
|
||||
self,
|
||||
collection_name: str,
|
||||
embedding: np.ndarray,
|
||||
limit: int,
|
||||
min_relevance_score: float,
|
||||
with_embeddings: bool,
|
||||
) -> List[Tuple[MemoryRecord, float]]:
|
||||
nearVector = {
|
||||
"vector": embedding,
|
||||
"certainty": min_relevance_score,
|
||||
}
|
||||
|
||||
additional = ["certainty"]
|
||||
if with_embeddings:
|
||||
additional.append("vector")
|
||||
|
||||
query = (
|
||||
self.client.query.get(collection_name, ALL_PROPERTIES)
|
||||
.with_near_vector(nearVector)
|
||||
.with_additional(additional)
|
||||
.with_limit(limit)
|
||||
)
|
||||
|
||||
results = await asyncio.to_thread(query.do)
|
||||
|
||||
get_dict = results.get("data", {}).get("Get", {})
|
||||
|
||||
memory_records_and_scores = [
|
||||
(
|
||||
self._convert_weaviate_doc_to_memory_record(doc),
|
||||
item["_additional"]["certainty"],
|
||||
)
|
||||
for items in get_dict.values()
|
||||
for item in items
|
||||
for doc in [item]
|
||||
]
|
||||
|
||||
return memory_records_and_scores
|
||||
|
||||
async def get_nearest_match_async(
|
||||
self,
|
||||
collection_name: str,
|
||||
embedding: np.ndarray,
|
||||
min_relevance_score: float,
|
||||
with_embedding: bool,
|
||||
) -> Tuple[MemoryRecord, float]:
|
||||
results = await self.get_nearest_matches_async(
|
||||
collection_name,
|
||||
embedding,
|
||||
limit=1,
|
||||
min_relevance_score=min_relevance_score,
|
||||
with_embeddings=with_embedding,
|
||||
)
|
||||
|
||||
return results[0]
|
||||
@@ -1,7 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Awaitable, Callable, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import numpy.testing as npt
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.memory import weaviate_memory_store
|
||||
from semantic_kernel.memory.memory_record import MemoryRecord
|
||||
|
||||
if not sys.platform.startswith("linux"):
|
||||
pytest.skip(
|
||||
"test_weaviate_memory_store uses embedded weaviate which only runs on Linux at the moment",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def documents():
|
||||
records = []
|
||||
|
||||
records.append(
|
||||
MemoryRecord.local_record(
|
||||
"1",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"A classic pangram.",
|
||||
np.array([0.1, 0.1]),
|
||||
)
|
||||
)
|
||||
records.append(
|
||||
MemoryRecord.local_record(
|
||||
"2",
|
||||
"The five boxing wizards jump quickly.",
|
||||
"Another popular pangram.",
|
||||
np.array([0.1, 0.11]),
|
||||
)
|
||||
)
|
||||
records.append(
|
||||
MemoryRecord.local_record(
|
||||
"3",
|
||||
"Pack my box with five dozen liquor jugs.",
|
||||
"A useful pangram.",
|
||||
np.array([0.11, 0.1]),
|
||||
)
|
||||
)
|
||||
|
||||
records.append(
|
||||
MemoryRecord.local_record(
|
||||
"4",
|
||||
"Lorem ipsum dolor sit amet.",
|
||||
"A common placeholder text.",
|
||||
np.array([-10, -10]),
|
||||
)
|
||||
)
|
||||
records.append(
|
||||
MemoryRecord.local_record(
|
||||
"5",
|
||||
"Etiam faucibus orci vitae lacus pellentesque.",
|
||||
"A Latin text.",
|
||||
np.array([-10.1, -10.2]),
|
||||
)
|
||||
)
|
||||
|
||||
yield records
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_store():
|
||||
config = weaviate_memory_store.WeaviateConfig(use_embed=True)
|
||||
store = weaviate_memory_store.WeaviateMemoryStore(config)
|
||||
store.client.schema.delete_all()
|
||||
yield store
|
||||
store.client.schema.delete_all()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_store_with_empty_collection(memory_store, event_loop):
|
||||
collection_name = "MindRepository"
|
||||
event_loop.run_until_complete(memory_store.create_collection_async(collection_name))
|
||||
return collection_name, memory_store
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_store_with_collection(memory_store, event_loop, documents):
|
||||
collection_name = "BigMemory"
|
||||
event_loop.run_until_complete(memory_store.create_collection_async(collection_name))
|
||||
|
||||
keys = ["Alpha", "Beta", "Gamma", "Delta", "Epsilon"]
|
||||
for document, key in zip(documents, keys):
|
||||
document._key = key
|
||||
event_loop.run_until_complete(
|
||||
memory_store.upsert_async(collection_name, document)
|
||||
)
|
||||
|
||||
return collection_name, memory_store
|
||||
|
||||
|
||||
def test_embedded_weaviate():
|
||||
config = weaviate_memory_store.WeaviateConfig(use_embed=True)
|
||||
memory_store = weaviate_memory_store.WeaviateMemoryStore(config=config)
|
||||
|
||||
assert memory_store.client._connection.embedded_db
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_collection(memory_store):
|
||||
collection_name = "MemoryVault"
|
||||
await memory_store.create_collection_async(collection_name)
|
||||
|
||||
assert memory_store.client.schema.get(collection_name)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_collections(memory_store):
|
||||
collection_names = ["MemoryVault", "ThoughtArchive"]
|
||||
|
||||
for collection_name in collection_names:
|
||||
await memory_store.create_collection_async(collection_name)
|
||||
|
||||
results = await memory_store.get_collections_async()
|
||||
|
||||
assert set(results) == set(collection_names)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_collection(memory_store_with_empty_collection):
|
||||
collection_name, memory_store = memory_store_with_empty_collection
|
||||
|
||||
schemas = memory_store.client.schema.get()["classes"]
|
||||
assert len(schemas) == 1
|
||||
|
||||
await memory_store.delete_collection_async(collection_name)
|
||||
|
||||
schemas = memory_store.client.schema.get()["classes"]
|
||||
assert len(schemas) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_collection_exists(memory_store_with_empty_collection):
|
||||
collection_name, memory_store = memory_store_with_empty_collection
|
||||
|
||||
memory_store.client.schema.get()["classes"]
|
||||
|
||||
assert await memory_store.does_collection_exist_async(collection_name)
|
||||
assert not await memory_store.does_collection_exist_async("NotACollection")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert(memory_store_with_empty_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_empty_collection
|
||||
|
||||
for doc in documents[:2]:
|
||||
await memory_store.upsert_async(collection_name, doc)
|
||||
|
||||
total_docs = memory_store.client.data_object.get(class_name=collection_name)[
|
||||
"totalResults"
|
||||
]
|
||||
assert total_docs == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_batch(memory_store_with_empty_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_empty_collection
|
||||
|
||||
await memory_store.upsert_batch_async(collection_name, documents)
|
||||
|
||||
total_docs = memory_store.client.data_object.get(class_name=collection_name)[
|
||||
"totalResults"
|
||||
]
|
||||
assert total_docs == len(documents)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get(memory_store_with_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_collection
|
||||
|
||||
key = "Alpha"
|
||||
|
||||
expected_result = [doc for doc in documents if doc._key == key][0]
|
||||
actual_result = await memory_store.get_async(
|
||||
collection_name, key, with_embedding=True
|
||||
)
|
||||
npt.assert_equal(expected_result.__dict__, actual_result.__dict__)
|
||||
|
||||
actual_result = await memory_store.get_async(
|
||||
collection_name, key, with_embedding=False
|
||||
)
|
||||
expected_result.__dict__["_embedding"] = None
|
||||
npt.assert_equal(expected_result.__dict__, actual_result.__dict__)
|
||||
|
||||
key = "NotInCollection"
|
||||
actual_result = await memory_store.get_async(
|
||||
collection_name, key, with_embedding=True
|
||||
)
|
||||
|
||||
assert actual_result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_batch(memory_store_with_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_collection
|
||||
|
||||
keys = ["Alpha", "Beta", "Gamma"]
|
||||
|
||||
expected_results = [doc for doc in documents if doc._key in keys]
|
||||
|
||||
actual_results = await memory_store.get_batch_async(
|
||||
collection_name, keys, with_embedding=True
|
||||
)
|
||||
|
||||
for expected, actual in zip(expected_results, actual_results):
|
||||
npt.assert_equal(expected.__dict__, actual.__dict__)
|
||||
|
||||
actual_results = await memory_store.get_batch_async(
|
||||
collection_name, keys, with_embedding=False
|
||||
)
|
||||
|
||||
for expected, actual in zip(expected_results, actual_results):
|
||||
expected.__dict__["_embedding"] = None
|
||||
npt.assert_equal(expected.__dict__, actual.__dict__)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_batch(memory_store_with_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_collection
|
||||
|
||||
keys = ["Alpha", "Beta", "Gamma"]
|
||||
|
||||
await memory_store.remove_batch_async(collection_name, keys)
|
||||
|
||||
remaining_docs = memory_store.client.data_object.get(class_name=collection_name)[
|
||||
"totalResults"
|
||||
]
|
||||
assert remaining_docs == len(documents) - len(keys)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove(memory_store_with_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_collection
|
||||
|
||||
key = "Alpha"
|
||||
|
||||
await memory_store.remove_async(collection_name, key)
|
||||
|
||||
remaining_docs = memory_store.client.data_object.get(class_name=collection_name)[
|
||||
"totalResults"
|
||||
]
|
||||
assert remaining_docs == len(documents) - 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nearest_matches(memory_store_with_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_collection
|
||||
|
||||
search_query = np.array([-10, -10])
|
||||
min_relevance_score = 0.9
|
||||
limit = 4
|
||||
|
||||
expected_result = [documents[3], documents[4]]
|
||||
actual_result = await memory_store.get_nearest_matches_async(
|
||||
collection_name, search_query, limit, min_relevance_score, with_embeddings=True
|
||||
)
|
||||
actual_docss, _ = list(zip(*actual_result))
|
||||
|
||||
assert len(actual_result) == len(expected_result)
|
||||
for expected, actual in zip(expected_result, actual_docss):
|
||||
npt.assert_equal(expected.__dict__, actual.__dict__)
|
||||
|
||||
actual_result = await memory_store.get_nearest_matches_async(
|
||||
collection_name, search_query, limit, min_relevance_score, with_embeddings=False
|
||||
)
|
||||
actual_docss, _ = list(zip(*actual_result))
|
||||
|
||||
assert len(actual_result) == len(expected_result)
|
||||
for expected, actual in zip(expected_result, actual_docss):
|
||||
expected.__dict__["_embedding"] = None
|
||||
npt.assert_equal(expected.__dict__, actual.__dict__)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nearest_match(memory_store_with_collection, documents):
|
||||
collection_name, memory_store = memory_store_with_collection
|
||||
|
||||
search_query = np.array([-10, -10])
|
||||
min_relevance_score = 0.9
|
||||
|
||||
expected_result = documents[3]
|
||||
actual_result = await memory_store.get_nearest_match_async(
|
||||
collection_name, search_query, min_relevance_score, with_embedding=True
|
||||
)
|
||||
|
||||
npt.assert_equal(expected_result.__dict__, actual_result[0].__dict__)
|
||||
|
||||
actual_result = await memory_store.get_nearest_match_async(
|
||||
collection_name, search_query, min_relevance_score, with_embedding=False
|
||||
)
|
||||
|
||||
expected_result.__dict__["_embedding"] = None
|
||||
npt.assert_equal(expected_result.__dict__, actual_result[0].__dict__)
|
||||
Reference in New Issue
Block a user