From eb5f30fcd1adad2b35a5ac433bdc205f03f56a75 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Mon, 11 Nov 2024 15:56:09 +0100 Subject: [PATCH] fix: fix single data point addition to weaiate --- .../vector/weaviate_db/WeaviateAdapter.py | 17 +++++++++-------- cognee/tests/test_weaviate.py | 2 +- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py index dd7539118..4ebae3b29 100644 --- a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py +++ b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py @@ -11,7 +11,6 @@ from ..embeddings.EmbeddingEngine import EmbeddingEngine logger = logging.getLogger("WeaviateAdapter") class IndexSchema(DataPoint): - uuid: str text: str _metadata: dict = { @@ -89,8 +88,10 @@ class WeaviateAdapter(VectorDBInterface): def convert_to_weaviate_data_points(data_point: DataPoint): vector = data_vectors[data_points.index(data_point)] properties = data_point.model_dump() - properties["uuid"] = properties["id"] - del properties["id"] + + if "id" in properties: + properties["uuid"] = str(data_point.id) + del properties["id"] return DataObject( uuid = data_point.id, @@ -114,7 +115,7 @@ class WeaviateAdapter(VectorDBInterface): ) else: data_point: DataObject = data_points[0] - return collection.data.update( + return collection.data.insert( uuid = data_point.uuid, vector = data_point.vector, properties = data_point.properties, @@ -130,8 +131,8 @@ class WeaviateAdapter(VectorDBInterface): async def index_data_points(self, index_name: str, index_property_name: str, data_points: list[DataPoint]): await self.create_data_points(f"{index_name}_{index_property_name}", [ IndexSchema( - uuid = str(data_point.id), - text = getattr(data_point, data_point._metadata["index_fields"][0]), + id = data_point.id, + text = data_point.get_embeddable_data(), ) for data_point in data_points ]) @@ -178,9 +179,9 @@ class WeaviateAdapter(VectorDBInterface): return [ ScoredResult( - id = UUID(result.uuid), + id = UUID(str(result.uuid)), payload = result.properties, - score = float(result.metadata.score) + score = 1 - float(result.metadata.score) ) for result in search_result.objects ] diff --git a/cognee/tests/test_weaviate.py b/cognee/tests/test_weaviate.py index 3f853f63e..f788f9973 100644 --- a/cognee/tests/test_weaviate.py +++ b/cognee/tests/test_weaviate.py @@ -35,7 +35,7 @@ async def main(): from cognee.infrastructure.databases.vector import get_vector_engine vector_engine = get_vector_engine() - random_node = (await vector_engine.search("Entity_name", "quantum computer"))[0] + random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] random_node_name = random_node.payload["text"] search_results = await cognee.search(SearchType.INSIGHTS, query = random_node_name)