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

feat: adds graph completion retriever fix (#676)

Browse source 
<!-- .github/pull_request_template.md -->

## Description
Adds graph completion retriever fix

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin
This commit is contained in:
hajdul88 2025-03-28 17:39:56 +01:00 • committed by GitHub
parent f6747128a2
commit 119fa1eb73
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 7 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
cognee/modules/retrieval/graph_completion_retriever.py
View file

@ -60,9 +60,13 @@ class GraphCompletionRetriever(BaseRetriever):
vector_index_collections = []
for subclass in subclasses:
index_fields = subclass.model_fields["metadata"].default.get("index_fields", [])
for field_name in index_fields:
vector_index_collections.append(f"{subclass.__name__}_{field_name}")
if "metadata" in subclass.model_fields:
metadata_field = subclass.model_fields["metadata"]
if hasattr(metadata_field, "default") and metadata_field.default is not None:
if isinstance(metadata_field.default, dict):
index_fields = metadata_field.default.get("index_fields", [])
for field_name in index_fields:
vector_index_collections.append(f"{subclass.__name__}_{field_name}")
found_triplets = await brute_force_triplet_search(
query, top_k=self.top_k, collections=vector_index_collections or None