gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: convert qdrant search results to ScoredPoint

Browse source
This commit is contained in:
Boris Arzentar 2024-11-11 14:38:59 +01:00 • committed by Leon Luithlen
parent 9fe1b6c5fa
commit b1b6b79ca4
6 changed files with 8 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py
View file

@ -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

1
cognee/modules/engine/utils/__init__.py
View file

@ -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

2
cognee/modules/engine/utils/generate_node_name.py
View file

@ -1,2 +1,2 @@
def generate_node_name(name: str) -> str:
return name.lower().replace(" ", "_").replace("'", "")
return name.lower().replace("'", "")

6
cognee/tasks/graph/extract_graph_from_data.py
View file

@ -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,
),

2
cognee/tests/test_qdrant.py
View file

@ -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)

2
cognee/tests/test_weaviate.py
View file

@ -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)