fix: Fixes ruff formatting

This commit is contained in:
hajdul88 2025-01-09 12:00:49 +01:00
parent fe57eb69e7
commit 341f30fcdc
3 changed files with 8 additions and 6 deletions

View file

@ -67,7 +67,7 @@ async def specific_search(query_type: SearchType, query: str, user) -> list:
SearchType.INSIGHTS: query_graph_connections,
SearchType.CHUNKS: query_chunks,
SearchType.COMPLETION: query_completion,
SearchType.GRAPH_COMPLETION: graph_query_completion
SearchType.GRAPH_COMPLETION: graph_query_completion,
}
search_task = search_tasks.get(query_type)

View file

@ -8,13 +8,14 @@ from cognee.modules.retrieval.brute_force_triplet_search import brute_force_trip
def retrieved_edges_to_string(retrieved_edges):
edge_strings = []
for edge in retrieved_edges:
node1_string = edge.node1.attributes['text'] or edge.node1.attributes.get('name')
node2_string = edge.node2.attributes['text'] or edge.node2.attributes.get('name')
edge_string = edge.attributes['relationship_type']
node1_string = edge.node1.attributes["text"] or edge.node1.attributes.get("name")
node2_string = edge.node2.attributes["text"] or edge.node2.attributes.get("name")
edge_string = edge.attributes["relationship_type"]
edge_str = f"{node1_string} -- {edge_string} -- {node2_string}"
edge_strings.append(edge_str)
return "\n---\n".join(edge_strings)
async def graph_query_completion(query: str) -> list:
"""
Parameters:

View file

@ -186,11 +186,12 @@ async def main(enable_steps):
# Step 4: Query insights
if enable_steps.get("retriever"):
search_results = await cognee.search(
SearchType.GRAPH_COMPLETION, query_text='Who has experience in design tools?'
SearchType.GRAPH_COMPLETION, query_text="Who has experience in design tools?"
)
print(search_results)
if __name__ == '__main__':
if __name__ == "__main__":
setup_logging(logging.INFO)
rebuild_kg = True