From 599151ed24c7315417dc920f38732ae55bb1ffaa Mon Sep 17 00:00:00 2001 From: berrazuriz1 Date: Fri, 26 Dec 2025 17:57:51 -0300 Subject: [PATCH 1/4] Reset processing status when merging nodes --- blarify/repositories/graph_db_manager/neo4j_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/blarify/repositories/graph_db_manager/neo4j_manager.py b/blarify/repositories/graph_db_manager/neo4j_manager.py index 4842f5cb..49684438 100644 --- a/blarify/repositories/graph_db_manager/neo4j_manager.py +++ b/blarify/repositories/graph_db_manager/neo4j_manager.py @@ -125,8 +125,8 @@ def _create_nodes_txn( "CALL apoc.merge.node( node.extra_labels + [node.type, 'NODE'], {hashed_id: node.attributes.hashed_id, repoId: $repoId, entityId: $entityId, environment: $environment, diff_identifier: node.attributes.diff_identifier}, - node.attributes, - node.attributes + apoc.map.merge(node.attributes, {processing_status: null, processing_run_id: null}), + apoc.map.merge(node.attributes, {processing_status: null, processing_run_id: null}) ) YIELD node as n RETURN count(n) as count", {batchSize: $batchSize, parallel: false, iterateList: true, params: {nodeList: $nodeList, repoId: $repoId, entityId: $entityId, environment: $environment}} From 9ab3c8373d5a6e28f9c952669d30289930eacd65 Mon Sep 17 00:00:00 2001 From: berrazuriz1 Date: Fri, 26 Dec 2025 17:57:59 -0300 Subject: [PATCH 2/4] Add queries for incremental documentation updates --- blarify/documentation/queries/__init__.py | 22 ++++-- .../queries/batch_processing_queries.py | 76 +++++++++++++++++++ 2 files changed, 91 insertions(+), 7 deletions(-) diff --git a/blarify/documentation/queries/__init__.py b/blarify/documentation/queries/__init__.py index c53d031e..25408803 100644 --- a/blarify/documentation/queries/__init__.py +++ b/blarify/documentation/queries/__init__.py @@ -2,20 +2,28 @@ from .batch_processing_queries import ( check_pending_nodes_query, + get_child_descriptions_query, + get_direct_callers_of_nodes_in_files_query, + get_latest_processing_run_id_query, get_leaf_nodes_batch_query, - get_processable_nodes_with_descriptions_query, - mark_nodes_completed_query, get_leaf_nodes_under_node_query, - get_child_descriptions_query, + get_parent_folders_for_files_query, + get_processable_nodes_with_descriptions_query, get_remaining_pending_functions_query, + mark_nodes_completed_query, + reset_processing_status_for_nodes_query, ) __all__ = [ - "get_leaf_nodes_batch_query", - "get_processable_nodes_with_descriptions_query", "check_pending_nodes_query", - "mark_nodes_completed_query", - "get_leaf_nodes_under_node_query", "get_child_descriptions_query", + "get_direct_callers_of_nodes_in_files_query", + "get_latest_processing_run_id_query", + "get_leaf_nodes_batch_query", + "get_leaf_nodes_under_node_query", + "get_parent_folders_for_files_query", + "get_processable_nodes_with_descriptions_query", "get_remaining_pending_functions_query", + "mark_nodes_completed_query", + "reset_processing_status_for_nodes_query", ] diff --git a/blarify/documentation/queries/batch_processing_queries.py b/blarify/documentation/queries/batch_processing_queries.py index 0831fbcd..dafc51cd 100644 --- a/blarify/documentation/queries/batch_processing_queries.py +++ b/blarify/documentation/queries/batch_processing_queries.py @@ -324,3 +324,79 @@ def get_hierarchical_parents_query() -> LiteralString: info.node.end_line as end_line ORDER BY info.sortKey ASC, info.depth ASC """ + + +def get_direct_callers_of_nodes_in_files_query() -> LiteralString: + """ + Find nodes that directly CALL functions within the specified file paths (1 level). + + This query finds all FUNCTION nodes that call any function defined within + the specified files. Used to identify which callers need re-documentation + when a file changes. + + Expected params: $entity_id, $repo_ids (optional), $file_paths + """ + return """ + MATCH (file:FILE {entityId: $entity_id}) + WHERE ($repo_ids IS NULL OR file.repoId IN $repo_ids) + AND file.node_path IN $file_paths + MATCH (file)-[:FUNCTION_DEFINITION|CLASS_DEFINITION*1..]->(target:FUNCTION) + MATCH (caller:FUNCTION)-[:CALLS]->(target) + WHERE NOT caller.node_path STARTS WITH file.node_path + RETURN DISTINCT caller.node_id as id, caller.node_path as path + """ + + +def reset_processing_status_for_nodes_query() -> LiteralString: + """ + Reset processing_status to NULL for specified nodes in a single batch. + + Used during incremental documentation to mark affected nodes (callers, parent folders) + for re-processing without affecting unrelated nodes. + + Expected params: $entity_id, $repo_ids (optional), $node_ids + """ + return """ + MATCH (n:NODE {entityId: $entity_id}) + WHERE ($repo_ids IS NULL OR n.repoId IN $repo_ids) + AND n.node_id IN $node_ids + SET n.processing_status = NULL, n.processing_run_id = NULL + RETURN count(n) as reset_count + """ + + +def get_parent_folders_for_files_query() -> LiteralString: + """ + Get all parent folder nodes for the specified file paths. + + Traverses up the CONTAINS relationship to find all ancestor folders. + Used to identify which folders need re-documentation when a file changes. + + Expected params: $entity_id, $repo_ids (optional), $file_paths + """ + return """ + MATCH (file:FILE {entityId: $entity_id}) + WHERE ($repo_ids IS NULL OR file.repoId IN $repo_ids) + AND file.node_path IN $file_paths + MATCH (parent:FOLDER)-[:CONTAINS*1..]->(file) + RETURN DISTINCT parent.node_id as id, parent.node_path as path + """ + + +def get_latest_processing_run_id_query() -> LiteralString: + """ + Get the most recent processing_run_id from any completed node. + + Used to retrieve the run_id from the previous documentation run, + so incremental updates can reuse it and skip already-processed nodes. + + Expected params: $entity_id, $repo_ids (optional) + """ + return """ + MATCH (n:NODE {entityId: $entity_id}) + WHERE ($repo_ids IS NULL OR n.repoId IN $repo_ids) + AND n.processing_run_id IS NOT NULL + AND n.processing_status = 'completed' + RETURN n.processing_run_id as run_id + LIMIT 1 + """ From 7cfe5eae53228e47278e5e56a1c5a492d9b1738d Mon Sep 17 00:00:00 2001 From: berrazuriz1 Date: Fri, 26 Dec 2025 17:58:06 -0300 Subject: [PATCH 3/4] Simplify targeted documentation for incremental updates --- .../documentation/documentation_creator.py | 151 ++++++++++-------- .../utils/bottom_up_batch_processor.py | 68 +------- blarify/prebuilt/graph_builder.py | 3 + 3 files changed, 95 insertions(+), 127 deletions(-) diff --git a/blarify/documentation/documentation_creator.py b/blarify/documentation/documentation_creator.py index 0b2d33fe..d4a684ba 100644 --- a/blarify/documentation/documentation_creator.py +++ b/blarify/documentation/documentation_creator.py @@ -8,6 +8,7 @@ import time import logging +import uuid from typing import List, Optional, TYPE_CHECKING, Dict, Any if TYPE_CHECKING: @@ -24,6 +25,12 @@ update_documentation_embeddings_query, create_vector_index_query, ) +from .queries import ( + get_direct_callers_of_nodes_in_files_query, + get_latest_processing_run_id_query, + get_parent_folders_for_files_query, + reset_processing_status_for_nodes_query, +) from .queries.workflow_queries import cleanup_orphaned_documentation_query from ..graph.graph_environment import GraphEnvironment from ..graph.relationship.relationship_creator import RelationshipCreator @@ -209,11 +216,16 @@ def _create_targeted_documentation( generate_embeddings: bool = False, ) -> DocumentationResult: """ - Create documentation for specific paths - optimized for SWE benchmarks. + Create documentation for specific paths - for incremental updates. + + This method: + 1. Gets the previous run_id from the last full documentation run + 2. Resets processing_status for affected nodes (callers + parent folders) + 3. Processes from ROOT with the old run_id (only NULL-status nodes get processed) Args: - target_paths: List of specific paths to document - framework_info: Framework detection results + target_paths: List of specific file paths to document + generate_embeddings: Whether to generate embeddings for documentation nodes Returns: DocumentationResult with targeted documentation @@ -221,84 +233,97 @@ def _create_targeted_documentation( try: logger.info(f"Creating targeted documentation for {len(target_paths)} paths") - all_information_nodes = [] - all_documentation_nodes = [] - all_source_nodes = [] - analyzed_nodes = [] - warnings = [] - - try: - entry_points_paths = self._discover_entry_points(file_paths=target_paths) - - # Use BottomUpBatchProcessor for each target path - # Phase 1 - Process execution stacks from entry points - processor = BottomUpBatchProcessor( - db_manager=self.db_manager, - agent_caller=self.agent_caller, - graph_environment=self.graph_environment, - max_workers=self.max_workers, - overwrite_documentation=self.overwrite_documentation, - generate_embeddings=generate_embeddings, - ) - for entry_point in entry_points_paths: - processor_result = processor.process_node(entry_point) - - if processor_result.error: - warnings.append( - f"Error processing path {target_paths} entry point {entry_point}: {processor_result.error}" - ) - continue - - # Collect results - all_documentation_nodes.extend(processor_result.documentation_nodes) - all_source_nodes.extend(processor_result.source_nodes) - analyzed_nodes.append( - { - "path": entry_point, - "node_count": len(processor_result.information_nodes), - "hierarchical_analysis": processor_result.hierarchical_analysis, - } - ) + # Step 1: Get the old run_id from previous full documentation + previous_run_id = self._get_previous_run_id() + if not previous_run_id: + logger.warning("No previous run_id found, generating new one") + previous_run_id = str(uuid.uuid4()) - logger.info(f"Processed {target_paths}: {len(processor_result.documentation_nodes)} nodes") + # Step 2: Reset processing status for affected nodes + self._reset_affected_nodes_status(target_paths) - # Phase 2: Process upstream definition dependencies - for target_path in target_paths: - processor_result = processor.process_upstream_definitions(target_path) + # Step 3: Get root path and process from root with old run_id + root_path = get_root_path(db_manager=self.db_manager) + if not root_path: + return DocumentationResult(error="No root path found") - if processor_result.error: - warnings.append( - f"Error processing upstream definitions for path {target_paths}: {processor_result.error}" - ) - continue - all_information_nodes.extend(processor_result.information_nodes) - all_documentation_nodes.extend(processor_result.documentation_nodes) - all_source_nodes.extend(processor_result.source_nodes) + processor = BottomUpBatchProcessor( + db_manager=self.db_manager, + agent_caller=self.agent_caller, + graph_environment=self.graph_environment, + max_workers=self.max_workers, + overwrite_documentation=self.overwrite_documentation, + generate_embeddings=generate_embeddings, + processing_run_id=previous_run_id, + ) - except Exception as e: - error_msg = f"Error processing path {target_paths}: {str(e)}" - logger.exception(error_msg) - warnings.append(error_msg) + result = processor.process_node(root_path) # Cleanup orphaned documentation nodes after incremental update orphans_deleted = self.cleanup_orphaned_documentation() if orphans_deleted > 0: logger.info(f"Cleaned up {orphans_deleted} orphaned documentation nodes") - logger.info(f"Targeted documentation completed: {len(all_documentation_nodes)} total nodes") + logger.info(f"Targeted documentation completed: {len(result.documentation_nodes)} total nodes") return DocumentationResult( - information_nodes=all_information_nodes, - documentation_nodes=all_documentation_nodes, - source_nodes=all_source_nodes, - analyzed_nodes=analyzed_nodes, - warnings=warnings, + information_nodes=result.information_nodes, + documentation_nodes=result.documentation_nodes, + source_nodes=result.source_nodes, + total_nodes_processed=result.total_nodes_processed, + warnings=[result.error] if result.error else [], ) except Exception as e: logger.exception(f"Error in targeted documentation creation: {e}") return DocumentationResult(error=str(e)) + def _get_previous_run_id(self) -> Optional[str]: + """Get the run_id from the previous documentation run.""" + try: + result = self.db_manager.query( + cypher_query=get_latest_processing_run_id_query(), + parameters={}, + ) + if result and result[0].get("run_id"): + return str(result[0]["run_id"]) + return None + except Exception as e: + logger.warning(f"Error getting previous run_id: {e}") + return None + + def _reset_affected_nodes_status(self, file_paths: List[str]) -> None: + """Reset processing_status for nodes that need re-documentation.""" + try: + nodes_to_reset: List[str] = [] + + # 1. Get direct callers (1 level) of ANY node in changed files + callers = self.db_manager.query( + cypher_query=get_direct_callers_of_nodes_in_files_query(), + parameters={"file_paths": file_paths}, + ) + if callers: + nodes_to_reset.extend([c["id"] for c in callers if c.get("id")]) + + # 2. Get parent folders + parents = self.db_manager.query( + cypher_query=get_parent_folders_for_files_query(), + parameters={"file_paths": file_paths}, + ) + if parents: + nodes_to_reset.extend([p["id"] for p in parents if p.get("id")]) + + # 3. Reset all in one batch + if nodes_to_reset: + self.db_manager.query( + cypher_query=reset_processing_status_for_nodes_query(), + parameters={"node_ids": nodes_to_reset}, + ) + logger.info(f"Reset processing status for {len(nodes_to_reset)} affected nodes") + + except Exception as e: + logger.warning(f"Error resetting affected nodes: {e}") + def _create_full_documentation(self, generate_embeddings: bool = False) -> DocumentationResult: """ Create documentation for the entire codebase. diff --git a/blarify/documentation/utils/bottom_up_batch_processor.py b/blarify/documentation/utils/bottom_up_batch_processor.py index 1da4d576..3a5d0f73 100644 --- a/blarify/documentation/utils/bottom_up_batch_processor.py +++ b/blarify/documentation/utils/bottom_up_batch_processor.py @@ -23,7 +23,6 @@ ) from blarify.documentation.queries.batch_processing_queries import ( get_child_descriptions_query, - get_hierarchical_parents_query, get_leaf_nodes_under_node_query, get_remaining_pending_functions_query, ) @@ -103,6 +102,7 @@ def __init__( overwrite_documentation: bool = False, batch_size: int = 1000, generate_embeddings: bool = False, + processing_run_id: Optional[str] = None, ): """ Initialize the query-based batch processor. @@ -117,6 +117,7 @@ def __init__( root_node: Optional root node to start processing from overwrite_documentation: Whether to overwrite existing documentation batch_size: Number of nodes to process in each batch + processing_run_id: Optional run ID for incremental updates (reuses previous run's ID) """ self.db_manager = db_manager self.agent_caller = agent_caller @@ -127,8 +128,8 @@ def __init__( self.batch_size = batch_size self.generate_embeddings = generate_embeddings - # Unique ID for this processing run - self.processing_run_id = str(uuid.uuid4()) + # Use provided run_id or generate new one + self.processing_run_id = processing_run_id or str(uuid.uuid4()) # Initialize embedding service if needed self.embedding_service = None @@ -141,44 +142,6 @@ def __init__( self.all_documentation_nodes: List[DocumentationNode] = [] self.all_source_nodes: List[NodeWithContentDto] = [] - def process_upstream_definitions(self, node_path: str) -> ProcessingResult: - """ - Process upstream definition dependencies for a given node path. - - Args: - node_path: Path to the node (folder or file) to process - Returns: - ProcessingResult with processing statistics - """ - try: - # Reset tracking lists for this processing run - self.all_documentation_nodes = [] - self.all_source_nodes = [] - - # Get root node - root_node = self.root_node - if not root_node: - root_node = get_node_by_path(self.db_manager, node_path) - if not root_node: - return ProcessingResult(node_path=node_path, error=f"Node not found: {node_path}") - - # Process using queries - total_processed = self._process_upstream_definitions(root_node) - - return ProcessingResult( - node_path=node_path, - hierarchical_analysis={"complete": True}, - total_nodes_processed=total_processed, - error=None, - information_nodes=[], - documentation_nodes=self.all_documentation_nodes, - source_nodes=self.all_source_nodes, - ) - - except Exception as e: - logger.exception(f"Error in upstream definition processing: {e}") - return ProcessingResult(node_path=node_path, error=str(e)) - def process_node(self, node_path: str) -> ProcessingResult: """ Entry point - process using database queries only. @@ -218,29 +181,6 @@ def process_node(self, node_path: str) -> ProcessingResult: logger.exception(f"Error in query-based processing: {e}") return ProcessingResult(node_path=node_path, error=str(e)) - def _process_upstream_definitions(self, root_node: NodeWithContentDto) -> int: - """Process upstream definitions using database queries.""" - - hierarchical_parents = self.db_manager.query( - cypher_query=get_hierarchical_parents_query(), - parameters={"node_id": root_node.id}, - ) - - for parent in hierarchical_parents: - parent_node = NodeWithContentDto( - id=parent["id"], - name=parent["name"], - labels=parent["labels"], - path=parent["path"], - start_line=parent.get("start_line"), - end_line=parent.get("end_line"), - content=parent.get("content", ""), - ) - child_descriptions = self.__get_child_descriptions(root_node) - self._process_parent_node(parent_node, child_descriptions=child_descriptions) - - return len(hierarchical_parents) - def _process_node_query_based(self, root_node: NodeWithContentDto) -> int: """Process using database queries without memory storage.""" diff --git a/blarify/prebuilt/graph_builder.py b/blarify/prebuilt/graph_builder.py index 9e49d935..83cab9de 100644 --- a/blarify/prebuilt/graph_builder.py +++ b/blarify/prebuilt/graph_builder.py @@ -139,6 +139,7 @@ def incremental_update( save_to_db: bool = True, create_workflows: bool = False, create_documentation: bool = False, + max_workers: int = 75, ) -> Graph: """Incrementally update the code graph for specific files with optional persistence and documentation/workflow generation. @@ -150,6 +151,7 @@ def incremental_update( save_to_db: Whether to save the graph to database (requires db_manager) create_workflows: Whether to discover and create workflows (requires db_manager and save_to_db) create_documentation: Whether to generate documentation (requires db_manager and save_to_db) + max_workers: Maximum number of parallel workers for documentation creation Returns: Graph object containing updated code nodes @@ -214,6 +216,7 @@ def incremental_update( db_manager=self.db_manager, agent_caller=agent_caller, graph_environment=self.graph_environment, + max_workers=max_workers, ) doc_creator.create_documentation(target_paths=node_paths, generate_embeddings=self.generate_embeddings) From f1aa9eabe8be70194e8a50e577fa8cf35a4b8cf4 Mon Sep 17 00:00:00 2001 From: berrazuriz1 Date: Fri, 26 Dec 2025 17:58:13 -0300 Subject: [PATCH 4/4] Add test for incremental documentation affected nodes --- .../test_graphbuilder_incremental_update.py | 117 ++++++++++++++++++ 1 file changed, 117 insertions(+) diff --git a/tests/integration/test_graphbuilder_incremental_update.py b/tests/integration/test_graphbuilder_incremental_update.py index 2e06ab5c..99be0d0e 100644 --- a/tests/integration/test_graphbuilder_incremental_update.py +++ b/tests/integration/test_graphbuilder_incremental_update.py @@ -986,3 +986,120 @@ async def test_incremental_update_cleans_up_orphaned_documentation( assert docs_with_describes_final[0]["count"] == len(final_docs), "All docs should have DESCRIBES after cleanup" db_manager.close() + + async def test_incremental_update_redocuments_affected_nodes( + self, + docker_check: Any, + test_data_isolation: Dict[str, Any], + temp_project_dir: Path, + test_code_examples_path: Path, + graph_assertions: GraphAssertions, + ) -> None: + """ + Test that incremental_update re-documents affected nodes: + - All nodes inside changed file (hierarchy) + - Direct callers of changed nodes (1 level only) + - Parent folders + + And does NOT re-document nodes 2+ levels away. + + Uses simple_linear workflow: main() → process_data() → transform_value() + """ + import shutil + + # Copy simple_linear workflow to temp dir + src_dir = test_code_examples_path / "workflows" / "simple_linear" + for file_name in ["main.py", "processor.py", "utils.py", "__init__.py"]: + shutil.copy(src_dir / file_name, temp_project_dir / file_name) + + response_prefix: Dict[str, str] = {"value": "Initial"} + + def node_specific_doc(input_dict: Dict[str, Any]) -> str: + node_name = input_dict.get("node_name") or input_dict.get("node_path") or "unknown" + return f"{response_prefix['value']} doc for {node_name}" + + llm_provider = make_llm_mock(dumb_response=node_specific_doc) + + db_manager = Neo4jManager( + uri=test_data_isolation["uri"], + user="neo4j", + password=test_data_isolation["password"], + repo_id=test_data_isolation["repo_id"], + entity_id=test_data_isolation["entity_id"], + ) + + try: + with ( + patch("blarify.prebuilt.graph_builder.LLMProvider", return_value=llm_provider), + patch( + "blarify.documentation.documentation_creator.LLMProvider", + return_value=llm_provider, + ), + ): + builder = GraphBuilder( + root_path=str(temp_project_dir), + db_manager=db_manager, + generate_embeddings=False, + ) + + builder.build( + save_to_db=True, + create_workflows=False, + create_documentation=True, + ) + + # Verify initial docs were created for all nodes + async def get_doc_content(node_name: str) -> str: + docs = await test_data_isolation["container"].execute_cypher( + """ + MATCH (doc:DOCUMENTATION)-[:DESCRIBES]->(n:NODE) + WHERE n.name = $name + RETURN doc.content as content + """, + parameters={"name": node_name}, + ) + return docs[0]["content"] if docs else "" + + # Verify initial state - all should have "Initial" + for node_name in ["transform_value", "process_data", "main"]: + content = await get_doc_content(node_name) + assert "Initial" in content, f"{node_name} should have Initial doc before update" + + # Change mock to return "Updated" + response_prefix["value"] = "Updated" + + # Modify utils.py (leaf node in call chain) + utils_file = temp_project_dir / "utils.py" + utils_file.write_text( + '"""Utility functions for the workflow."""\n\n' + "def transform_value(value: int) -> int:\n" + ' """Transform a value by tripling it."""\n' + " return value * 3\n" + ) + + # Run incremental update targeting utils.py + updated_files = [UpdatedFile(path=str(utils_file))] + builder.incremental_update( + updated_files, + save_to_db=True, + create_workflows=False, + create_documentation=True, + ) + + # ASSERT: Changed file nodes → "Updated" + transform_content = await get_doc_content("transform_value") + assert "Updated" in transform_content, "transform_value (changed file) should have Updated doc" + + # ASSERT: Direct caller (1 level) → "Updated" + process_data_content = await get_doc_content("process_data") + assert "Updated" in process_data_content, "process_data (direct caller) should have Updated doc" + + # ASSERT: 2 levels away → still "Initial" (NOT updated) + main_content = await get_doc_content("main") + assert "Initial" in main_content, ( + "main() (2 levels away) should NOT be re-documented. " + f"Got: {main_content}" + ) + + finally: + db_manager.close()