diff --git a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py index 436861a45..1efcd47b3 100644 --- a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py +++ b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py @@ -3,6 +3,7 @@ from uuid import UUID from typing import List, Dict, Optional from qdrant_client import AsyncQdrantClient, models +from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult from cognee.infrastructure.engine import DataPoint from ..vector_db_interface import VectorDBInterface from ..embeddings.EmbeddingEngine import EmbeddingEngine diff --git a/cognee/modules/engine/utils/__init__.py b/cognee/modules/engine/utils/__init__.py index 9cc2bc573..4d4ab02e7 100644 --- a/cognee/modules/engine/utils/__init__.py +++ b/cognee/modules/engine/utils/__init__.py @@ -1,2 +1,3 @@ from .generate_node_id import generate_node_id from .generate_node_name import generate_node_name +from .generate_edge_name import generate_edge_name diff --git a/cognee/modules/engine/utils/generate_node_name.py b/cognee/modules/engine/utils/generate_node_name.py index 84b266198..a2871875b 100644 --- a/cognee/modules/engine/utils/generate_node_name.py +++ b/cognee/modules/engine/utils/generate_node_name.py @@ -1,2 +1,2 @@ def generate_node_name(name: str) -> str: - return name.lower().replace(" ", "_").replace("'", "") + return name.lower().replace("'", "") diff --git a/cognee/tasks/graph/extract_graph_from_data.py b/cognee/tasks/graph/extract_graph_from_data.py index 36cc3e2fc..9e6edcabd 100644 --- a/cognee/tasks/graph/extract_graph_from_data.py +++ b/cognee/tasks/graph/extract_graph_from_data.py @@ -5,7 +5,7 @@ from cognee.infrastructure.databases.graph import get_graph_engine from cognee.modules.data.extraction.knowledge_graph import extract_content_graph from cognee.modules.chunking.models.DocumentChunk import DocumentChunk from cognee.modules.engine.models import EntityType, Entity -from cognee.modules.engine.utils import generate_node_id, generate_node_name +from cognee.modules.engine.utils import generate_edge_name, generate_node_id, generate_node_name from cognee.tasks.storage import add_data_points async def extract_graph_from_data(data_chunks: list[DocumentChunk], graph_model: Type[BaseModel]): @@ -95,7 +95,7 @@ async def extract_graph_from_data(data_chunks: list[DocumentChunk], graph_model: for edge in graph.edges: source_node_id = generate_node_id(edge.source_node_id) target_node_id = generate_node_id(edge.target_node_id) - relationship_name = generate_node_name(edge.relationship_name) + relationship_name = generate_edge_name(edge.relationship_name) edge_key = str(source_node_id) + str(target_node_id) + relationship_name @@ -105,7 +105,7 @@ async def extract_graph_from_data(data_chunks: list[DocumentChunk], graph_model: target_node_id, edge.relationship_name, dict( - relationship_name = generate_node_name(edge.relationship_name), + relationship_name = generate_edge_name(edge.relationship_name), source_node_id = source_node_id, target_node_id = target_node_id, ), diff --git a/cognee/tests/test_qdrant.py b/cognee/tests/test_qdrant.py index faa2cbcf4..784b3f27a 100644 --- a/cognee/tests/test_qdrant.py +++ b/cognee/tests/test_qdrant.py @@ -37,7 +37,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", "AI"))[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) diff --git a/cognee/tests/test_weaviate.py b/cognee/tests/test_weaviate.py index 121c1749e..3f853f63e 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", "AI"))[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)