diff --git a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py index 3543418fc..2f9d42b06 100644 --- a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py +++ b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py @@ -62,10 +62,12 @@ class Neo4jAdapter(GraphDBInterface): async def add_node(self, node: DataPoint): serialized_properties = self.serialize_properties(node.model_dump()) - query = dedent("""MERGE (node {id: $node_id}) + query = dedent( + """MERGE (node {id: $node_id}) ON CREATE SET node += $properties, node.updated_at = timestamp() ON MATCH SET node += $properties, node.updated_at = timestamp() - RETURN ID(node) AS internal_id, node.id AS nodeId""") + RETURN ID(node) AS internal_id, node.id AS nodeId""" + ) params = { "node_id": str(node.id), @@ -182,13 +184,15 @@ class Neo4jAdapter(GraphDBInterface): ): serialized_properties = self.serialize_properties(edge_properties) - query = dedent("""MATCH (from_node {id: $from_node}), + query = dedent( + """MATCH (from_node {id: $from_node}), (to_node {id: $to_node}) MERGE (from_node)-[r]->(to_node) ON CREATE SET r += $properties, r.updated_at = timestamp(), r.type = $relationship_name ON MATCH SET r += $properties, r.updated_at = timestamp() RETURN r - """) + """ + ) params = { "from_node": str(from_node), diff --git a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py index 77bfd74e6..dd8934c34 100644 --- a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py +++ b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py @@ -88,23 +88,27 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface): } ) - return dedent(f""" + return dedent( + f""" MERGE (node:{node_label} {{id: '{str(data_point.id)}'}}) ON CREATE SET node += ({{{node_properties}}}), node.updated_at = timestamp() ON MATCH SET node += ({{{node_properties}}}), node.updated_at = timestamp() - """).strip() + """ + ).strip() async def create_edge_query(self, edge: tuple[str, str, str, dict]) -> str: properties = await self.stringify_properties(edge[3]) properties = f"{{{properties}}}" - return dedent(f""" + return dedent( + f""" MERGE (source {{id:'{edge[0]}'}}) MERGE (target {{id: '{edge[1]}'}}) MERGE (source)-[edge:{edge[2]} {properties}]->(target) ON MATCH SET edge.updated_at = timestamp() ON CREATE SET edge.updated_at = timestamp() - """).strip() + """ + ).strip() async def create_collection(self, collection_name: str): pass @@ -195,12 +199,14 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface): self.query(query) async def has_edges(self, edges): - query = dedent(""" + query = dedent( + """ UNWIND $edges AS edge MATCH (a)-[r]->(b) WHERE id(a) = edge.from_node AND id(b) = edge.to_node AND type(r) = edge.relationship_name RETURN edge.from_node AS from_node, edge.to_node AS to_node, edge.relationship_name AS relationship_name, count(r) > 0 AS edge_exists - """).strip() + """ + ).strip() params = { "edges": [ @@ -279,14 +285,16 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface): [label, attribute_name] = collection_name.split(".") - query = dedent(f""" + query = dedent( + f""" CALL db.idx.vector.queryNodes( '{label}', '{attribute_name}', {limit}, vecf32({query_vector}) ) YIELD node, score - """).strip() + """ + ).strip() result = self.query(query) diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py index 6c3c5029d..68561979d 100644 --- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py +++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py @@ -93,10 +93,12 @@ class SQLAlchemyAdapter: if self.engine.dialect.name == "postgresql": async with self.engine.begin() as connection: result = await connection.execute( - text(""" + text( + """ SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('pg_catalog', 'pg_toast', 'information_schema'); - """) + """ + ) ) return [schema[0] for schema in result.fetchall()] return [] diff --git a/cognee/shared/utils.py b/cognee/shared/utils.py index f11b77ae6..6b1ca7f8f 100644 --- a/cognee/shared/utils.py +++ b/cognee/shared/utils.py @@ -445,7 +445,7 @@ async def create_cognee_style_network_with_logo( logging.info(f"Saving visualization to {output_filename}...") html_content = file_html(p, CDN, title) - home_dir = os.path.expanduser('~') + home_dir = os.path.expanduser("~") # Construct the final output file path output_filepath = os.path.join(home_dir, output_filename)