Partition huge clustering inputs by repo structure; survive empty LLM responses

The clustering response must re-emit every leaf-node ID verbatim, so repos
with many leaf nodes (electron: 1085) blow past max_tokens; the proxy then
returns content: null and cluster_modules crashed on the None (twice: the
tag check and the except handler's response[:200]) instead of falling back.

- cluster_modules now partitions leaf nodes along directory structure into
  batches capped by max_leaf_nodes_per_cluster (default 600) and an output
  token budget (max_tokens // 2), clusters each batch separately, and merges
  the results into one flat top level for super_group_modules to organize.
  A failed batch degrades to a directory-named fallback module.
- call_llm logs finish_reason=length truncation and may return None; all
  consumers (clustering, super-grouping, overview generation) now treat an
  empty response as a graceful fallback instead of a TypeError.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: anhnh2002

codewiki/cli/adapters/doc_generator.py

@@ -145,6 +145,7 @@ def generate(self) -> DocumentationJob:
                 max_tokens=self.config.get('max_tokens', 32768),
                 max_token_per_module=self.config.get('max_token_per_module', 36369),
                 max_token_per_leaf_module=self.config.get('max_token_per_leaf_module', 16000),
+                max_leaf_nodes_per_cluster=self.config.get('max_leaf_nodes_per_cluster', 600),
                 max_depth=self.config.get('max_depth', 2),
                 agent_instructions=self.config.get('agent_instructions'),
                 use_gitignore=self.config.get('use_gitignore', True),

codewiki/src/be/cluster_modules.py

@@ -7,11 +7,16 @@
 
 from codewiki.src.be.dependency_analyzer.models.core import Node
 from codewiki.src.be.llm_services import call_llm
+from codewiki.src.be.module_naming import resolve_unique_name, sanitize_module_name
 from codewiki.src.be.utils import count_tokens
-from codewiki.src.config import Config, DEFAULT_MIN_MODULES_FOR_SUPER_GROUPING
+from codewiki.src.config import (
+    Config,
+    DEFAULT_MAX_LEAF_NODES_PER_CLUSTER,
+    DEFAULT_MIN_MODULES_FOR_SUPER_GROUPING,
+)
 from codewiki.src.be.prompt_template import format_cluster_prompt, format_super_group_prompt
 
-Completer = Callable[[str], str]
+Completer = Callable[[str], Optional[str]]
 
 # When whole-repo mode is chosen but leaf entry points touch fewer than this
 # fraction of parsed files, warn that coverage depends on agent exploration.
@@ -58,6 +63,198 @@ def get_clustering_input_token_count(
     return count_tokens(potential_core_components_with_code)
 
 
+def _cluster_batch_fits(node_ids: List[str], config: Config) -> bool:
+    """Whether a single LLM clustering call can handle these nodes.
+
+    The clustering response must re-emit every component ID verbatim, so the
+    joined ID list is a direct proxy for output size; keep 2x headroom under
+    max_tokens for dict syntax, module names/paths, and preamble.
+    """
+    max_nodes = getattr(
+        config, "max_leaf_nodes_per_cluster", DEFAULT_MAX_LEAF_NODES_PER_CLUSTER
+    )
+    if len(node_ids) > max_nodes:
+        return False
+    output_budget = max(2048, config.max_tokens // 2)
+    return count_tokens("\n".join(node_ids)) <= output_budget
+
+
+def partition_leaf_nodes_by_structure(
+    leaf_nodes: List[str],
+    components: Dict[str, Node],
+    fits: Callable[[List[str]], bool],
+) -> List[List[str]]:
+    """Partition leaf nodes into batches that each satisfy ``fits``.
+
+    Splits along the directory structure of the nodes' relative paths, then
+    greedily coalesces path-adjacent small groups so root-level files and tiny
+    directories don't become their own batches. Deterministic: nodes are
+    processed in (relative_path, id) order. Returns the input as a single
+    batch when it already fits.
+    """
+    valid = []
+    for leaf_node in leaf_nodes:
+        if leaf_node in components:
+            valid.append(leaf_node)
+        else:
+            logger.warning(f"Skipping invalid leaf node '{leaf_node}' - not found in components")
+    valid.sort(key=lambda node: (components[node].relative_path, node))
+    if not valid or fits(valid):
+        return [valid]
+
+    def path_parts(node: str) -> List[str]:
+        return components[node].relative_path.strip("/").split("/")
+
+    def chunk(nodes: List[str]) -> List[List[str]]:
+        # Nodes that share one directory/file and still don't fit can only be
+        # cut into fixed-size slices.
+        size = len(nodes)
+        while size > 1 and not fits(nodes[:size]):
+            size = (size + 1) // 2
+        logger.warning(
+            "Splitting %d leaf nodes that share one path into slices of at "
+            "most %d; directory structure can't divide them further.",
+            len(nodes),
+            size,
+        )
+        return [nodes[i:i + size] for i in range(0, len(nodes), size)]
+
+    def split(nodes: List[str], depth: int) -> List[List[str]]:
+        by_prefix = defaultdict(list)
+        for node in nodes:
+            by_prefix["/".join(path_parts(node)[:depth])].append(node)
+        groups: List[List[str]] = []
+        for prefix in sorted(by_prefix):
+            sub = by_prefix[prefix]
+            if fits(sub):
+                groups.append(sub)
+            elif any(len(path_parts(node)) > depth for node in sub):
+                groups.extend(split(sub, depth + 1))
+            else:
+                groups.extend(chunk(sub))
+        return groups
+
+    batches: List[List[str]] = []
+    current: List[str] = []
+    for group in split(valid, 1):
+        if not current:
+            current = group
+        elif fits(current + group):
+            current.extend(group)
+        else:
+            batches.append(current)
+            current = group
+    if current:
+        batches.append(current)
+    return batches
+
+
+def _cluster_via_llm(
+    leaf_nodes: List[str],
+    components: Dict[str, Node],
+    config: Config,
+    current_module_tree: Dict[str, Any],
+    current_module_name: Optional[str],
+    module_label: str,
+    completer: Optional[Completer],
+) -> Dict[str, Any]:
+    """Run one clustering LLM call over these nodes.
+
+    Returns {} for any empty, malformed, or non-dict response instead of
+    raising, so callers can fall back gracefully.
+    """
+    potential_core_components, _ = format_potential_core_components(leaf_nodes, components)
+    prompt = format_cluster_prompt(potential_core_components, current_module_tree, current_module_name)
+    if completer is not None:
+        response = completer(prompt)
+    else:
+        response = call_llm(prompt, config, model=config.cluster_model)
+
+    if not response:
+        logger.warning(
+            "Empty LLM clustering response for %s (provider returned no "
+            "content, likely output truncation); falling back.",
+            module_label,
+        )
+        return {}
+
+    try:
+        if "<GROUPED_COMPONENTS>" not in response or "</GROUPED_COMPONENTS>" not in response:
+            logger.warning(
+                "Invalid LLM clustering response for %s: missing <GROUPED_COMPONENTS> "
+                "tags; falling back. Response preview: %s...",
+                module_label,
+                str(response)[:200],
+            )
+            return {}
+
+        response_content = response.split("<GROUPED_COMPONENTS>")[1].split("</GROUPED_COMPONENTS>")[0]
+        module_tree = eval(response_content)
+
+        if not isinstance(module_tree, dict):
+            logger.error(f"Invalid module tree format - expected dict, got {type(module_tree)}")
+            return {}
+
+    except Exception as e:
+        logger.warning(
+            "Failed to parse LLM clustering response for %s; falling back. "
+            "Error: %s. Response preview: %s...",
+            module_label,
+            e,
+            str(response)[:200],
+        )
+        logger.error(f"Traceback: {traceback.format_exc()}")
+        return {}
+
+    return module_tree
+
+
+def _merge_module_trees(target: Dict[str, Any], addition: Dict[str, Any]) -> None:
+    """Merge one batch's module dict into the accumulated tree in place.
+
+    Batches cluster independently, so two of them may propose the same module
+    name for the same concept; merge their components rather than renaming
+    (dedupe_module_tree_names still guards anything unexpected later).
+    """
+    for name, info in addition.items():
+        if name not in target:
+            target[name] = info
+            continue
+        existing = target[name]
+        if not isinstance(existing, dict) or not isinstance(info, dict):
+            logger.warning(
+                "Cannot merge module '%s' from another batch; keeping the first version.",
+                name,
+            )
+            continue
+        seen = set(existing.get("components", []))
+        merged = list(existing.get("components", []))
+        for component in info.get("components", []):
+            if component not in seen:
+                seen.add(component)
+                merged.append(component)
+        existing["components"] = merged
+        paths = [existing.get("path", ""), info.get("path", "")]
+        existing["path"] = _common_path_prefix(paths) if all(paths) else ""
+        logger.info(
+            "Module '%s' was produced by multiple clustering batches; merged "
+            "into %d components.",
+            name,
+            len(merged),
+        )
+
+
+def _batch_fallback_name(
+    batch: List[str], components: Dict[str, Node], existing: Dict[str, Any]
+) -> str:
+    """Directory-derived module name for a batch whose LLM clustering failed."""
+    prefix = _common_path_prefix(
+        [components[node].relative_path for node in batch if node in components]
+    )
+    name = sanitize_module_name(prefix.replace("/", "_")) if prefix else "root"
+    return resolve_unique_name(name, None, set(existing))
+
+
 def cluster_modules(
     leaf_nodes: List[str],
     components: Dict[str, Node],
@@ -77,7 +274,7 @@ def cluster_modules(
             subscription-mode (caw) routing.  If ``None``, falls back to
             ``call_llm`` for backward compatibility with direct callers.
     """
-    potential_core_components, potential_core_components_with_code = (
+    _, potential_core_components_with_code = (
         format_potential_core_components(leaf_nodes, components)
     )
     input_tokens = count_tokens(potential_core_components_with_code)
@@ -118,46 +315,68 @@ def cluster_modules(
                 )
         return {}
 
-    prompt = format_cluster_prompt(potential_core_components, current_module_tree, current_module_name)
     logger.info(
         "Requesting LLM module clustering for %s because %d tokens exceed the %d-token threshold.",
         module_label,
         input_tokens,
         threshold,
     )
-    if completer is not None:
-        response = completer(prompt)
-    else:
-        response = call_llm(prompt, config, model=config.cluster_model)
 
-    #parse the response
-    try:
-        if "<GROUPED_COMPONENTS>" not in response or "</GROUPED_COMPONENTS>" not in response:
-            logger.warning(
-                "Invalid LLM clustering response for %s: missing <GROUPED_COMPONENTS> "
-                "tags; falling back to whole-module documentation. Response preview: %s...",
-                module_label,
-                response[:200],
-            )
-            return {}
-        
-        response_content = response.split("<GROUPED_COMPONENTS>")[1].split("</GROUPED_COMPONENTS>")[0]
-        module_tree = eval(response_content)
-        
-        if not isinstance(module_tree, dict):
-            logger.error(f"Invalid module tree format - expected dict, got {type(module_tree)}")
+    batches = partition_leaf_nodes_by_structure(
+        leaf_nodes, components, lambda ids: _cluster_batch_fits(ids, config)
+    )
+
+    if len(batches) == 1:
+        module_tree = _cluster_via_llm(
+            batches[0],
+            components,
+            config,
+            current_module_tree,
+            current_module_name,
+            module_label,
+            completer,
+        )
+        if not module_tree:
             return {}
-            
-    except Exception as e:
-        logger.warning(
-            "Failed to parse LLM clustering response for %s; falling back to "
-            "whole-module documentation. Error: %s. Response preview: %s...",
+    else:
+        logger.info(
+            "Partitioned %d leaf nodes for %s into %d structure-based batches for clustering.",
+            len(leaf_nodes),
             module_label,
-            e,
-            response[:200],
+            len(batches),
         )
-        logger.error(f"Traceback: {traceback.format_exc()}")
-        return {}
+        module_tree = {}
+        for i, batch in enumerate(batches, 1):
+            batch_label = f"{module_label} (batch {i}/{len(batches)})"
+            # Each batch gets the *unmodified* current_module_tree: passing the
+            # accumulating merge would flip format_cluster_prompt into its
+            # module-level variant mid-partition.
+            partial = _cluster_via_llm(
+                batch,
+                components,
+                config,
+                current_module_tree,
+                current_module_name,
+                batch_label,
+                completer,
+            )
+            if not partial:
+                name = _batch_fallback_name(batch, components, module_tree)
+                partial = {
+                    name: {
+                        "path": _common_path_prefix(
+                            [components[n].relative_path for n in batch if n in components]
+                        ),
+                        "components": list(batch),
+                    }
+                }
+                logger.warning(
+                    "Clustering failed for %s; keeping its %d components as fallback module '%s'.",
+                    batch_label,
+                    len(batch),
+                    name,
+                )
+            _merge_module_trees(module_tree, partial)
 
     # check if the module tree is valid
     if len(module_tree) <= 1:
@@ -182,7 +401,7 @@ def cluster_modules(
         for key in current_module_path:
             value = value[key]["children"]
         for module_name, module_info in module_tree.items():
-            del module_info["path"]
+            module_info.pop("path", None)
             value[module_name] = module_info
 
     for module_name, module_info in module_tree.items():
@@ -226,12 +445,18 @@ def _common_path_prefix(paths: List[str]) -> str:
     return "/".join(common)
 
 
-def _parse_super_group_response(response: str) -> Optional[Dict[str, Any]]:
+def _parse_super_group_response(response: Optional[str]) -> Optional[Dict[str, Any]]:
+    if not response:
+        logger.warning(
+            "Empty super-grouping response (provider returned no content, "
+            "likely output truncation); keeping the flat module tree."
+        )
+        return None
     if "<GROUPED_MODULES>" not in response or "</GROUPED_MODULES>" not in response:
         logger.warning(
             "Invalid super-grouping response: missing <GROUPED_MODULES> tags; "
             "keeping the flat module tree. Response preview: %s...",
-            response[:200],
+            str(response)[:200],
         )
         return None
     try:
@@ -243,7 +468,7 @@ def _parse_super_group_response(response: str) -> Optional[Dict[str, Any]]:
             "Failed to parse super-grouping response; keeping the flat module "
             "tree. Error: %s. Response preview: %s...",
             e,
-            response[:200],
+            str(response)[:200],
         )
         return None
     if not isinstance(grouping, dict):

codewiki/src/be/documentation_generator.py

@@ -287,6 +287,11 @@ async def generate_parent_module_docs(self, module_path: List[str],
 
         try:
             parent_docs = self.backend.complete(prompt)
+            if not parent_docs:
+                raise RuntimeError(
+                    f"LLM returned empty content for {module_name} overview "
+                    f"(possible output truncation at max_tokens)"
+                )
 
             # Parse and save parent documentation. Subscription-CLI backends
             # (claude-code / codex) sometimes ignore the <OVERVIEW> wrapper and