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.INSIGHTS: query_graph_connections,
SearchType.CHUNKS: query_chunks, SearchType.CHUNKS: query_chunks,
SearchType.COMPLETION: query_completion, SearchType.COMPLETION: query_completion,
SearchType.GRAPH_COMPLETION: graph_query_completion SearchType.GRAPH_COMPLETION: graph_query_completion,
} }
search_task = search_tasks.get(query_type) 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): def retrieved_edges_to_string(retrieved_edges):
edge_strings = [] edge_strings = []
for edge in retrieved_edges: for edge in retrieved_edges:
node1_string = edge.node1.attributes['text'] or edge.node1.attributes.get('name') node1_string = edge.node1.attributes["text"] or edge.node1.attributes.get("name")
node2_string = edge.node2.attributes['text'] or edge.node2.attributes.get('name') node2_string = edge.node2.attributes["text"] or edge.node2.attributes.get("name")
edge_string = edge.attributes['relationship_type'] edge_string = edge.attributes["relationship_type"]
edge_str = f"{node1_string} -- {edge_string} -- {node2_string}" edge_str = f"{node1_string} -- {edge_string} -- {node2_string}"
edge_strings.append(edge_str) edge_strings.append(edge_str)
return "\n---\n".join(edge_strings) return "\n---\n".join(edge_strings)
async def graph_query_completion(query: str) -> list: async def graph_query_completion(query: str) -> list:
""" """
Parameters: Parameters:

View file

@ -186,11 +186,12 @@ async def main(enable_steps):
# Step 4: Query insights # Step 4: Query insights
if enable_steps.get("retriever"): if enable_steps.get("retriever"):
search_results = await cognee.search( 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) print(search_results)
if __name__ == '__main__':
if __name__ == "__main__":
setup_logging(logging.INFO) setup_logging(logging.INFO)
rebuild_kg = True rebuild_kg = True