Cluster the potential core components into modules. Args: completer: optional ``(prompt: str) -> str`` callable. When provided, clustering calls go through this completer instead of the legacy ``call_llm``. This is how the LLMBackend abstraction injects
(
leaf_nodes: List[str],
components: Dict[str, Node],
config: Config,
current_module_tree: dict[str, Any] = {},
current_module_name: str = None,
current_module_path: List[str] = [],
completer: Optional[Completer] = None,
)
| 54 | |
| 55 | |
| 56 | def get_clustering_input_token_count( |
| 57 | leaf_nodes: List[str], components: Dict[str, Node] |
| 58 | ) -> int: |
| 59 | """Count the tokens used to decide whether a module needs clustering.""" |
| 60 | _, potential_core_components_with_code = format_potential_core_components( |
| 61 | leaf_nodes, components |
| 62 | ) |
| 63 | return count_tokens(potential_core_components_with_code) |
| 64 | |
| 65 | |
| 66 | def _cluster_batch_fits(node_ids: List[str], config: Config) -> bool: |
| 67 | """Whether a single LLM clustering call can handle these nodes. |
| 68 | |
| 69 | The clustering response must re-emit every component ID verbatim, so the |
| 70 | joined ID list is a direct proxy for output size; keep 2x headroom under |
| 71 | max_tokens for dict syntax, module names/paths, and preamble. |
| 72 | """ |
| 73 | max_nodes = getattr( |
| 74 | config, "max_leaf_nodes_per_cluster", DEFAULT_MAX_LEAF_NODES_PER_CLUSTER |
| 75 | ) |
| 76 | if len(node_ids) > max_nodes: |
| 77 | return False |
| 78 | output_budget = max(2048, config.max_tokens // 2) |
| 79 | return count_tokens("\n".join(node_ids)) <= output_budget |
| 80 | |
| 81 | |
| 82 | def partition_leaf_nodes_by_structure( |
| 83 | leaf_nodes: List[str], |
| 84 | components: Dict[str, Node], |
| 85 | fits: Callable[[List[str]], bool], |
| 86 | ) -> List[List[str]]: |
| 87 | """Partition leaf nodes into batches that each satisfy ``fits``. |
| 88 | |
| 89 | Splits along the directory structure of the nodes' relative paths, then |
| 90 | greedily coalesces path-adjacent small groups so root-level files and tiny |
| 91 | directories don't become their own batches. Deterministic: nodes are |
| 92 | processed in (relative_path, id) order. Returns the input as a single |
| 93 | batch when it already fits. |
| 94 | """ |
| 95 | valid = [] |
| 96 | for leaf_node in leaf_nodes: |
| 97 | if leaf_node in components: |
| 98 | valid.append(leaf_node) |
| 99 | else: |
| 100 | logger.warning(f"Skipping invalid leaf node '{leaf_node}' - not found in components") |
| 101 | valid.sort(key=lambda node: (components[node].relative_path, node)) |
| 102 | if not valid or fits(valid): |
| 103 | return [valid] |
| 104 | |
| 105 | def path_parts(node: str) -> List[str]: |
| 106 | return components[node].relative_path.strip("/").split("/") |
| 107 | |
| 108 | def chunk(nodes: List[str]) -> List[List[str]]: |
| 109 | # Nodes that share one directory/file and still don't fit can only be |
| 110 | # cut into fixed-size slices. |
| 111 | size = len(nodes) |
| 112 | while size > 1 and not fits(nodes[:size]): |
| 113 | size = (size + 1) // 2 |
no test coverage detected