fix: Fixes ruff formatting
This commit is contained in:
parent
fe57eb69e7
commit
341f30fcdc
3 changed files with 8 additions and 6 deletions
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue