Skip to content

Commit bbcac24

Browse files
anhnh2002claude
andcommitted
Trim module tree in module/cluster prompts when over codex input cap
format_user_prompt embeds the whole module tree with every component listed per file. On large repos (wazuh: 822 modules, ~823k chars) the prompt exceeds codex's hard 1,048,576-char turn input limit (input_too_large, code -32602), failing every remaining module. Same bug family as the overview prompts fixed earlier on this branch. Build the prompt unchanged first; only if it exceeds a 900k-char cap, rebuild with a names-only tree (wazuh: 26k chars) plus a note pointing the agent at its reading tools, and as a last resort truncate inlined file contents. format_cluster_prompt gets the same guard. Prompts under the cap are byte-identical to before. The duplicated tree-formatting closures are extracted into _format_module_tree_str. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 003b876 commit bbcac24

1 file changed

Lines changed: 142 additions & 77 deletions

File tree

‎codewiki/src/be/prompt_template.py‎

Lines changed: 142 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -257,10 +257,32 @@
257257
Reasoning at first, then return the list of relative paths in JSON format.
258258
"""
259259

260+
import logging
261+
from collections import defaultdict
260262
from typing import Any
261263

262264
from codewiki.src.utils import file_manager
263265

266+
logger = logging.getLogger(__name__)
267+
268+
# codex rejects any turn whose total input exceeds 1,048,576 characters
269+
# (input_too_large, code -32602 — server-side, not configurable). Cap the
270+
# user prompt below that, leaving headroom for the system prompt, tool
271+
# schemas and protocol overhead.
272+
MAX_USER_PROMPT_CHARS = 900_000
273+
274+
MODULE_TREE_TRIMMED_NOTE = (
275+
"NOTE: per-module component listings were omitted because the full module "
276+
"tree exceeds the model input limit. Module names and hierarchy are "
277+
"complete; read the referenced modules' documentation files or use your "
278+
"code-reading tools when you need component-level detail."
279+
)
280+
281+
CODE_TRUNCATED_NOTE = (
282+
"\n... [file contents truncated to fit the model input limit — use your "
283+
"file-reading tools to read the full files]"
284+
)
285+
264286
EXTENSION_TO_LANGUAGE = {
265287
".py": "python",
266288
".md": "markdown",
@@ -289,6 +311,54 @@
289311
}
290312

291313

314+
def _format_module_tree_str(
315+
module_tree: dict[str, Any],
316+
current_module_name: str | None = None,
317+
include_components: bool = True,
318+
) -> str:
319+
"""
320+
Render a module tree as an indented text outline.
321+
322+
With include_components=False only module names and hierarchy are
323+
emitted, which keeps the outline small enough for huge trees that would
324+
otherwise blow past MAX_USER_PROMPT_CHARS.
325+
"""
326+
lines: list[str] = []
327+
328+
def _walk(tree: dict[str, Any], indent: int = 0) -> None:
329+
for key, value in tree.items():
330+
if key == current_module_name:
331+
lines.append(f"{' ' * indent}{key} (current module)")
332+
else:
333+
lines.append(f"{' ' * indent}{key}")
334+
335+
if include_components:
336+
# Group components by file
337+
by_file = defaultdict(list)
338+
for c in value["components"]:
339+
if "::" in c:
340+
fpath, name = c.split("::", 1)
341+
by_file[fpath].append(name)
342+
else:
343+
by_file[""].append(c)
344+
for fpath, names in by_file.items():
345+
if fpath:
346+
lines.append(f"{' ' * (indent + 1)} {fpath}: {', '.join(names)}")
347+
else:
348+
lines.append(f"{' ' * (indent + 1)} {', '.join(names)}")
349+
350+
if (
351+
("children" in value)
352+
and isinstance(value["children"], dict)
353+
and len(value["children"]) > 0
354+
):
355+
lines.append(f"{' ' * (indent + 1)} Children:")
356+
_walk(value["children"], indent + 2)
357+
358+
_walk(module_tree, 0)
359+
return "\n".join(lines)
360+
361+
292362
def format_user_prompt(
293363
module_name: str,
294364
core_component_ids: list[str],
@@ -306,41 +376,7 @@ def format_user_prompt(
306376
Returns:
307377
Formatted user prompt string
308378
"""
309-
310-
# format module tree
311-
lines = []
312-
313-
def _format_module_tree(module_tree: dict[str, any], indent: int = 0):
314-
for key, value in module_tree.items():
315-
if key == module_name:
316-
lines.append(f"{' ' * indent}{key} (current module)")
317-
else:
318-
lines.append(f"{' ' * indent}{key}")
319-
320-
# Group components by file
321-
from collections import defaultdict
322-
323-
by_file = defaultdict(list)
324-
for c in value["components"]:
325-
if "::" in c:
326-
fpath, name = c.split("::", 1)
327-
by_file[fpath].append(name)
328-
else:
329-
by_file[""].append(c)
330-
for fpath, names in by_file.items():
331-
if fpath:
332-
lines.append(f"{' ' * (indent + 1)} {fpath}: {', '.join(names)}")
333-
else:
334-
lines.append(f"{' ' * (indent + 1)} {', '.join(names)}")
335-
336-
if isinstance(value["children"], dict) and len(value["children"]) > 0:
337-
lines.append(f"{' ' * (indent + 1)} Children:")
338-
_format_module_tree(value["children"], indent + 2)
339-
340-
_format_module_tree(module_tree, 0)
341-
formatted_module_tree = "\n".join(lines)
342-
343-
# print(f"Formatted module tree:\n{formatted_module_tree}")
379+
formatted_module_tree = _format_module_tree_str(module_tree, module_name)
344380

345381
# Group core component IDs by their file path
346382
grouped_components: dict[str, list[str]] = {}
@@ -375,12 +411,56 @@ def _format_module_tree(module_tree: dict[str, any], indent: int = 0):
375411

376412
core_component_codes += "```\n\n"
377413

378-
return USER_PROMPT.format(
414+
prompt = USER_PROMPT.format(
379415
module_name=module_name,
380416
formatted_core_component_codes=core_component_codes,
381417
module_tree=formatted_module_tree,
382418
)
383419

420+
if len(prompt) > MAX_USER_PROMPT_CHARS:
421+
full_len = len(prompt)
422+
formatted_module_tree = (
423+
MODULE_TREE_TRIMMED_NOTE
424+
+ "\n\n"
425+
+ _format_module_tree_str(module_tree, module_name, include_components=False)
426+
)
427+
prompt = USER_PROMPT.format(
428+
module_name=module_name,
429+
formatted_core_component_codes=core_component_codes,
430+
module_tree=formatted_module_tree,
431+
)
432+
logger.warning(
433+
"Module %s: user prompt (%d chars) exceeds %d; "
434+
"module tree trimmed to names only (%d chars)",
435+
module_name,
436+
full_len,
437+
MAX_USER_PROMPT_CHARS,
438+
len(prompt),
439+
)
440+
441+
if len(prompt) > MAX_USER_PROMPT_CHARS:
442+
# Even the slim tree was not enough — the inlined file contents
443+
# dominate. Truncation is recoverable: the agent has file-reading
444+
# tools (read_code_components / str_replace_editor).
445+
excess = len(prompt) - MAX_USER_PROMPT_CHARS + len(CODE_TRUNCATED_NOTE)
446+
core_component_codes = (
447+
core_component_codes[: max(0, len(core_component_codes) - excess)] + CODE_TRUNCATED_NOTE
448+
)
449+
prompt = USER_PROMPT.format(
450+
module_name=module_name,
451+
formatted_core_component_codes=core_component_codes,
452+
module_tree=formatted_module_tree,
453+
)
454+
logger.warning(
455+
"Module %s: user prompt still over %d chars after tree trim; "
456+
"truncated inlined file contents (now %d chars)",
457+
module_name,
458+
MAX_USER_PROMPT_CHARS,
459+
len(prompt),
460+
)
461+
462+
return prompt
463+
384464

385465
def format_cluster_prompt(
386466
potential_core_components: str,
@@ -393,53 +473,38 @@ def format_cluster_prompt(
393473
if module_tree is None:
394474
module_tree = {}
395475

396-
# format module tree
397-
lines = []
398-
399-
# print(f"Module tree:\n{json.dumps(module_tree, indent=2)}")
400-
401-
def _format_module_tree(module_tree: dict[str, any], indent: int = 0):
402-
for key, value in module_tree.items():
403-
if key == module_name:
404-
lines.append(f"{' ' * indent}{key} (current module)")
405-
else:
406-
lines.append(f"{' ' * indent}{key}")
407-
408-
# Group components by file
409-
from collections import defaultdict
410-
411-
by_file = defaultdict(list)
412-
for c in value["components"]:
413-
if "::" in c:
414-
fpath, name = c.split("::", 1)
415-
by_file[fpath].append(name)
416-
else:
417-
by_file[""].append(c)
418-
for fpath, names in by_file.items():
419-
if fpath:
420-
lines.append(f"{' ' * (indent + 1)} {fpath}: {', '.join(names)}")
421-
else:
422-
lines.append(f"{' ' * (indent + 1)} {', '.join(names)}")
423-
424-
if (
425-
("children" in value)
426-
and isinstance(value["children"], dict)
427-
and len(value["children"]) > 0
428-
):
429-
lines.append(f"{' ' * (indent + 1)} Children:")
430-
_format_module_tree(value["children"], indent + 2)
431-
432-
_format_module_tree(module_tree, 0)
433-
formatted_module_tree = "\n".join(lines)
434-
435476
if module_tree == {}:
436477
return CLUSTER_REPO_PROMPT.format(potential_core_components=potential_core_components)
437-
else:
438-
return CLUSTER_MODULE_PROMPT.format(
478+
479+
formatted_module_tree = _format_module_tree_str(module_tree, module_name)
480+
prompt = CLUSTER_MODULE_PROMPT.format(
481+
potential_core_components=potential_core_components,
482+
module_tree=formatted_module_tree,
483+
module_name=module_name,
484+
)
485+
486+
if len(prompt) > MAX_USER_PROMPT_CHARS:
487+
full_len = len(prompt)
488+
formatted_module_tree = (
489+
MODULE_TREE_TRIMMED_NOTE
490+
+ "\n\n"
491+
+ _format_module_tree_str(module_tree, module_name, include_components=False)
492+
)
493+
prompt = CLUSTER_MODULE_PROMPT.format(
439494
potential_core_components=potential_core_components,
440495
module_tree=formatted_module_tree,
441496
module_name=module_name,
442497
)
498+
logger.warning(
499+
"Module %s: cluster prompt (%d chars) exceeds %d; "
500+
"module tree trimmed to names only (%d chars)",
501+
module_name,
502+
full_len,
503+
MAX_USER_PROMPT_CHARS,
504+
len(prompt),
505+
)
506+
507+
return prompt
443508

444509

445510
def format_super_group_prompt(module_tree: dict[str, Any]) -> str:

0 commit comments

Comments
 (0)