diff --git a/blarify/code_hierarchy/languages/go_definitions.py b/blarify/code_hierarchy/languages/go_definitions.py index b69b2b44..b21578af 100644 --- a/blarify/code_hierarchy/languages/go_definitions.py +++ b/blarify/code_hierarchy/languages/go_definitions.py @@ -24,23 +24,13 @@ def get_parsers_for_extensions() -> Dict[str, Parser]: def should_create_node(node: Node) -> bool: return LanguageDefinitions._should_create_node_base_implementation( node, - ["type_declaration", "method_declaration", "function_declaration"], + ["type_spec", "type_alias", "method_declaration", "function_declaration"], ) def get_identifier_node(node: Node) -> Node: - if node.type == "type_declaration": - for child in node.named_children: - if child.type == "type_spec" or "type_alias": - node = child - return LanguageDefinitions._get_identifier_node_base_implementation(node) def get_body_node(node: Node) -> Node: - if node.type == "type_declaration": - for child in node.named_children: - if child.type == "type_spec": - node = child - return LanguageDefinitions._get_body_node_base_implementation(node) def get_relationship_type(node: GraphNode, node_in_point_reference: Node) -> Optional[FoundRelationshipScope]: @@ -51,7 +41,8 @@ def get_relationship_type(node: GraphNode, node_in_point_reference: Node) -> Opt def get_node_label_from_type(type: str) -> NodeLabels: return { - "type_declaration": NodeLabels.CLASS, + "type_spec": NodeLabels.CLASS, + "type_alias": NodeLabels.CLASS, "method_declaration": NodeLabels.FUNCTION, "function_declaration": NodeLabels.FUNCTION, }[type] diff --git a/blarify/db_managers/db_manager.py b/blarify/db_managers/db_manager.py new file mode 100644 index 00000000..9c85e852 --- /dev/null +++ b/blarify/db_managers/db_manager.py @@ -0,0 +1,20 @@ +class AbstractDbManager: + def close(self): + """Close the connection to the database.""" + raise NotImplementedError + + def save_graph(self, nodes, edges): + """Save nodes and edges to the database.""" + raise NotImplementedError + + def create_nodes(self, nodeList): + """Create nodes in the database.""" + raise NotImplementedError + + def create_edges(self, edgesList): + """Create edges between nodes in the database.""" + raise NotImplementedError + + def detatch_delete_nodes_with_path(self, path): + """Detach and delete nodes matching the given path.""" + raise NotImplementedError diff --git a/blarify/db_managers/falkordb_manager.py b/blarify/db_managers/falkordb_manager.py new file mode 100644 index 00000000..6b65eebc --- /dev/null +++ b/blarify/db_managers/falkordb_manager.py @@ -0,0 +1,81 @@ +import os +import time +from typing import Any, List +import logging + +from dotenv import load_dotenv +from falkordb import FalkorDB, exceptions + +logger = logging.getLogger(__name__) + +load_dotenv() + + +class FalkorDBManager: + entity_id: str + repo_id: str + db: FalkorDB + + def __init__( + self, + repo_id: str = None, + entity_id: str = None, + uri: str = None, + user: str = None, + password: str = None, + ): + host = uri or os.getenv("FALKORDB_URI", "localhost") + port = int(os.getenv("FALKORDB_PORT", 6379)) + user = user or os.getenv("FALKORDB_USERNAME") + password = password or os.getenv("FALKORDB_PASSWORD") + + self.db = FalkorDB(host=host, port=port, username=user, password=password) + + self.repo_id = repo_id if repo_id is not None else "default_repo" + self.entity_id = entity_id if entity_id is not None else "default_user" + + def close(self): + # Close the connection to the database + self.db.close() + + def save_graph(self, nodes: List[Any], edges: List[Any]): + self.create_nodes(nodes) + self.create_edges(edges) + + def create_nodes(self, nodeList: List[Any]): + # Function to create nodes in the FalkorDB database + graph = self.db.select_graph(self.repo_id) + for node in nodeList: + labels = ":".join(node.get("extra_labels", []) + [node["type"], "NODE"]) + attributes = node.get("attributes", {}) + attributes.update({"repoId": self.repo_id, "entityId": self.entity_id}) + # Construct parameterized query + cypher_query = f"CREATE (n:{labels} $props)" + graph.query(cypher_query, params={"props": attributes}) + + def create_edges(self, edgesList: List[Any]): + # Function to create edges between nodes in the FalkorDB database + graph = self.db.select_graph(self.repo_id) + for edge in edgesList: + # Construct parameterized query + cypher_query = ( + "MATCH (a:NODE {node_id: $sourceId}), " + "(b:NODE {node_id: $targetId}) " + "CREATE (a)-[r:$type {scopeText: $scopeText}]->(b)" + ) + graph.query( + cypher_query, + params={ + "sourceId": edge["sourceId"], + "targetId": edge["targetId"], + "type": edge["type"], + "scopeText": edge["scopeText"], + }, + ) + + def detach_delete_nodes_with_path(self, path: str): + graph = self.db.select_graph(self.repo_id) + # Construct parameterized query + cypher_query = "MATCH (n {path: $path}) DETACH DELETE n" + result = graph.query(cypher_query, params={"path": path}) + return result.result_set diff --git a/blarify/examples/graph_builder.py b/blarify/examples/graph_builder.py index e2c026a8..4d93c47a 100644 --- a/blarify/examples/graph_builder.py +++ b/blarify/examples/graph_builder.py @@ -1,5 +1,6 @@ from blarify.prebuilt.graph_builder import GraphBuilder from blarify.db_managers.neo4j_manager import Neo4jManager +from blarify.db_managers.falkordb_manager import FalkorDBManager import dotenv import os @@ -12,7 +13,8 @@ def build(root_path: str = None): relationships = graph.get_relationships_as_objects() nodes = graph.get_nodes_as_objects() - save_to_neo4j(relationships, nodes) + # save_to_neo4j(relationships, nodes) + save_to_falkordb(relationships, nodes) def save_to_neo4j(relationships, nodes): @@ -23,6 +25,14 @@ def save_to_neo4j(relationships, nodes): graph_manager.close() +def save_to_falkordb(relationships, nodes): + graph_manager = FalkorDBManager(repo_id="repo", entity_id="organization") + + print(f"Saving graph with {len(nodes)} nodes and {len(relationships)} relationships") + graph_manager.save_graph(nodes, relationships) + graph_manager.close() + + if __name__ == "__main__": import logging diff --git a/poetry.lock b/poetry.lock index 989caa97..8d1c3879 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,5 +1,18 @@ # This file is automatically @generated by Poetry 2.1.1 and should not be changed by hand. +[[package]] +name = "async-timeout" +version = "5.0.1" +description = "Timeout context manager for asyncio programs" +optional = false +python-versions = ">=3.8" +groups = ["main"] +markers = "python_full_version < \"3.11.3\"" +files = [ + {file = "async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c"}, + {file = "async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3"}, +] + [[package]] name = "asyncio" version = "3.4.3" @@ -203,6 +216,20 @@ files = [ [package.extras] test = ["pytest (>=6)"] +[[package]] +name = "falkordb" +version = "1.0.10" +description = "Python client for interacting with FalkorDB database" +optional = false +python-versions = "<4.0,>=3.8" +groups = ["main"] +files = [ + {file = "falkordb-1.0.10.tar.gz", hash = "sha256:f3db65e4e1f85f1b68a8cda97b1f7d756af4942a88712696fb2cfa53537025b8"}, +] + +[package.dependencies] +redis = ">=5.0.1,<6.0.0" + [[package]] name = "idna" version = "3.10" @@ -292,7 +319,7 @@ requests = "2.32.3" type = "git" url = "https://github.com/blarApp/multilspy.git" reference = "HEAD" -resolved_reference = "2d5a3906bcdea76b70799cf79739394b7eb27e91" +resolved_reference = "048bbeca841b6ae2e8891e6addfcc2328fc56219" [[package]] name = "neo4j" @@ -443,6 +470,25 @@ files = [ {file = "pytz-2025.1.tar.gz", hash = "sha256:c2db42be2a2518b28e65f9207c4d05e6ff547d1efa4086469ef855e4ab70178e"}, ] +[[package]] +name = "redis" +version = "5.2.1" +description = "Python client for Redis database and key-value store" +optional = false +python-versions = ">=3.8" +groups = ["main"] +files = [ + {file = "redis-5.2.1-py3-none-any.whl", hash = "sha256:ee7e1056b9aea0f04c6c2ed59452947f34c4940ee025f5dd83e6a6418b6989e4"}, + {file = "redis-5.2.1.tar.gz", hash = "sha256:16f2e22dff21d5125e8481515e386711a34cbec50f0e44413dd7d9c060a54e0f"}, +] + +[package.dependencies] +async-timeout = {version = ">=4.0.3", markers = "python_full_version < \"3.11.3\""} + +[package.extras] +hiredis = ["hiredis (>=3.0.0)"] +ocsp = ["cryptography (>=36.0.1)", "pyopenssl (==23.2.1)", "requests (>=2.31.0)"] + [[package]] name = "requests" version = "2.32.3" @@ -775,4 +821,4 @@ files = [ [metadata] lock-version = "2.1" python-versions = ">=3.10,<=3.14" -content-hash = "5a7403ca78dbdb36d7185e498e2223008abf732669128bc7498c7f510f155dd5" +content-hash = "06565b2c28596ebb4359a4c55aae0ec1a9d06478d8399bc12164fe9ab8b832c8" diff --git a/pyproject.toml b/pyproject.toml index 55061322..fa9b5a5c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ tree-sitter-typescript = "^0.23.2" tree-sitter-c-sharp = "^0.23.1" multilspy = {git = "https://github.com/blarApp/multilspy.git"} tree-sitter-go = "^0.23.4" +falkordb = "^1.0.10" [build-system]