From 4ad30369d20ee64bd91e3499bb8243e38edcf58e Mon Sep 17 00:00:00 2001 From: Perceval <36658095+perce7al@users.noreply.github.com> Date: Sat, 26 Sep 2026 14:31:19 +0300 Subject: [PATCH 1/3] feat: add max_retries config for PydanticAI agent retries Add a configurable `max_retries` parameter (default: 3) that controls the number of retries the PydanticAI agent attempts on failed tool calls via `Agent(retries=...)`. The parameter flows through Config, Configuration serialization, CLI (`config set`/`config show`/`generate`), and MCP server. --- .gitignore | 1 + codewiki/cli/adapters/doc_generator.py | 1 + codewiki/cli/commands/config.py | 17 +++- codewiki/cli/commands/generate.py | 9 ++ codewiki/cli/config_manager.py | 4 + codewiki/cli/models/config.py | 5 + codewiki/mcp/server.py | 1 + .../generate_sub_module_documentations.py | 2 + codewiki/src/be/pydantic_ai_backend.py | 3 + codewiki/src/config.py | 6 ++ tests/test_max_retries.py | 97 +++++++++++++++++++ 11 files changed, 143 insertions(+), 3 deletions(-) create mode 100644 tests/test_max_retries.py diff --git a/.gitignore b/.gitignore index d36f2b71..9ee8b9c5 100644 --- a/.gitignore +++ b/.gitignore @@ -47,6 +47,7 @@ tests/* !tests/test_branding.py !tests/test_artifact_analyzer.py !tests/test_scala_analyzer.py +!tests/test_max_retries.py # Jupyter *.ipynb diff --git a/codewiki/cli/adapters/doc_generator.py b/codewiki/cli/adapters/doc_generator.py index 92c066dd..86d94133 100644 --- a/codewiki/cli/adapters/doc_generator.py +++ b/codewiki/cli/adapters/doc_generator.py @@ -140,6 +140,7 @@ def generate(self) -> DocumentationJob: fallback_model=self.config.get("fallback_model"), provider=self.config.get("provider", "openai-compatible"), aws_region=self.config.get("aws_region", "us-east-1"), + max_retries=self.config.get("max_retries", 3), 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), diff --git a/codewiki/cli/commands/config.py b/codewiki/cli/commands/config.py index 17bed6cd..a3b9439e 100644 --- a/codewiki/cli/commands/config.py +++ b/codewiki/cli/commands/config.py @@ -39,6 +39,7 @@ def config_group(): @click.option("--main-model", type=str, help="Primary model for documentation generation") @click.option("--cluster-model", type=str, help="Model for module clustering (recommend top-tier)") @click.option("--fallback-model", type=str, help="Fallback model for documentation generation") +@click.option("--max-retries", type=int, help="Maximum number of retries for failed agent tool calls (default: 3)") @click.option("--max-tokens", type=int, help="Maximum tokens for LLM response (default: 32768)") @click.option( "--max-token-per-module", @@ -93,6 +94,7 @@ def config_set( main_model: Optional[str], cluster_model: Optional[str], fallback_model: Optional[str], + max_retries: Optional[int], max_tokens: Optional[int], max_token_per_module: Optional[int], max_token_per_leaf_module: Optional[int], @@ -106,12 +108,12 @@ def config_set( ): """ Set configuration values for CodeWiki. - + API keys are stored securely in your system keychain: • macOS: Keychain Access - • Windows: Credential Manager + • Windows: Credential Manager • Linux: Secret Service (GNOME Keyring, KWallet) - + Examples: \b @@ -162,6 +164,7 @@ def config_set( main_model, cluster_model, fallback_model, + max_retries, max_tokens, max_token_per_module, max_token_per_leaf_module, @@ -206,6 +209,11 @@ def config_set( if fallback_model: validated_data["fallback_model"] = validate_model_name(fallback_model) + if max_retries is not None: + if max_retries < 0: + raise ConfigurationError("max_retries must be a non-negative integer") + validated_data["max_retries"] = max_retries + if max_tokens is not None: if max_tokens < 1: raise ConfigurationError("max_tokens must be a positive integer") @@ -254,6 +262,7 @@ def config_set( main_model=validated_data.get("main_model"), cluster_model=validated_data.get("cluster_model"), fallback_model=validated_data.get("fallback_model"), + max_retries=validated_data.get("max_retries"), max_tokens=validated_data.get("max_tokens"), max_token_per_module=validated_data.get("max_token_per_module"), max_token_per_leaf_module=validated_data.get("max_token_per_leaf_module"), @@ -380,6 +389,7 @@ def config_show(output_json: bool): "cluster_model": config.cluster_model if config else "", "fallback_model": config.fallback_model if config else "glm-4p5", "default_output": config.default_output if config else "docs", + "max_retries": config.max_retries if config else 3, "max_tokens": config.max_tokens if config else 32768, "max_token_per_module": config.max_token_per_module if config else 36369, "max_token_per_leaf_module": config.max_token_per_leaf_module if config else 16000, @@ -441,6 +451,7 @@ def config_show(output_json: bool): click.echo() click.secho("Token Settings", fg="cyan", bold=True) if config: + click.echo(f" Max Retries: {config.max_retries}") click.echo(f" Max Tokens: {config.max_tokens}") click.echo(f" Max Token/Module: {config.max_token_per_module}") click.echo(f" Max Token/Leaf Module: {config.max_token_per_leaf_module}") diff --git a/codewiki/cli/commands/generate.py b/codewiki/cli/commands/generate.py index 2e476c86..fe385309 100644 --- a/codewiki/cli/commands/generate.py +++ b/codewiki/cli/commands/generate.py @@ -272,6 +272,12 @@ def _find_affected(tree, parent_names=None): is_flag=True, help="Show detailed progress and debug information", ) +@click.option( + "--max-retries", + type=int, + default=None, + help="Maximum retries for failed LLM requests (overrides config)", +) @click.option( "--max-tokens", type=int, @@ -394,6 +400,7 @@ def generate_command( instructions: str | None, use_gitignore: bool | None, verbose: bool, + max_retries: int | None, max_tokens: int | None, max_token_per_module: int | None, max_token_per_leaf_module: int | None, @@ -707,6 +714,8 @@ def generate_command( "provider": getattr(config, "provider", "openai-compatible"), "aws_region": getattr(config, "aws_region", "us-east-1"), "agent_instructions": agent_instructions_dict, + # Max retries setting (runtime override takes precedence) + "max_retries": max_retries if max_retries is not None else config.max_retries, # Max token settings (runtime overrides take precedence) "max_tokens": max_tokens if max_tokens is not None else config.max_tokens, "max_token_per_module": max_token_per_module diff --git a/codewiki/cli/config_manager.py b/codewiki/cli/config_manager.py index 3cbadcff..d421ee82 100644 --- a/codewiki/cli/config_manager.py +++ b/codewiki/cli/config_manager.py @@ -132,6 +132,7 @@ def save( cluster_model: Optional[str] = None, fallback_model: Optional[str] = None, default_output: Optional[str] = None, + max_retries: Optional[int] = None, max_tokens: Optional[int] = None, max_token_per_module: Optional[int] = None, max_token_per_leaf_module: Optional[int] = None, @@ -153,6 +154,7 @@ def save( cluster_model: Clustering model fallback_model: Fallback model default_output: Default output directory + max_retries: Maximum number of retries for failed agent tool calls max_tokens: Maximum tokens for LLM response max_token_per_module: Maximum tokens per module for clustering max_token_per_leaf_module: Maximum tokens per leaf module @@ -197,6 +199,8 @@ def save( self._config.fallback_model = fallback_model if default_output is not None: self._config.default_output = default_output + if max_retries is not None: + self._config.max_retries = max_retries if max_tokens is not None: self._config.max_tokens = max_tokens if max_token_per_module is not None: diff --git a/codewiki/cli/models/config.py b/codewiki/cli/models/config.py index c1cabe07..ed0884ca 100644 --- a/codewiki/cli/models/config.py +++ b/codewiki/cli/models/config.py @@ -126,6 +126,7 @@ class Configuration: aws_region: AWS region for Bedrock provider api_version: Azure OpenAI API version azure_deployment: Azure OpenAI deployment name + max_retries: Maximum number of retries for failed agent tool calls (default: 3) max_tokens: Maximum tokens for LLM response (default: 32768) max_token_per_module: Maximum tokens per module for clustering (default: 36369) max_token_per_leaf_module: Maximum tokens per leaf module (default: 16000) @@ -144,6 +145,7 @@ class Configuration: aws_region: str = "us-east-1" api_version: str = "2024-12-01-preview" azure_deployment: str = "" + max_retries: int = 3 max_tokens: int = 32768 max_token_per_module: int = 36369 max_token_per_leaf_module: int = 16000 @@ -183,6 +185,7 @@ def to_dict(self) -> dict: "aws_region": self.aws_region, "api_version": self.api_version, "azure_deployment": self.azure_deployment, + "max_retries": self.max_retries, "max_tokens": self.max_tokens, "max_token_per_module": self.max_token_per_module, "max_token_per_leaf_module": self.max_token_per_leaf_module, @@ -220,6 +223,7 @@ def from_dict(cls, data: dict) -> "Configuration": aws_region=data.get("aws_region", "us-east-1"), api_version=data.get("api_version", "2024-12-01-preview"), azure_deployment=data.get("azure_deployment", ""), + max_retries=data.get("max_retries", 3), max_tokens=data.get("max_tokens", 32768), max_token_per_module=data.get("max_token_per_module", 36369), max_token_per_leaf_module=data.get("max_token_per_leaf_module", 16000), @@ -296,6 +300,7 @@ def to_backend_config( aws_region=self.aws_region, api_version=self.api_version, azure_deployment=self.azure_deployment, + max_retries=self.max_retries, max_tokens=self.max_tokens, max_token_per_module=self.max_token_per_module, max_token_per_leaf_module=self.max_token_per_leaf_module, diff --git a/codewiki/mcp/server.py b/codewiki/mcp/server.py index bfcb810a..39254dfe 100644 --- a/codewiki/mcp/server.py +++ b/codewiki/mcp/server.py @@ -534,6 +534,7 @@ async def _legacy_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: fallback_model=config.fallback_model, provider=getattr(config, "provider", "openai-compatible"), aws_region=getattr(config, "aws_region", "us-east-1"), + max_retries=config.max_retries, max_tokens=config.max_tokens, agent_instructions=agent_instructions or None, use_gitignore=arguments.get("use_gitignore", True), diff --git a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py index 938a34cd..be3e4bf6 100644 --- a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py +++ b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py @@ -91,6 +91,7 @@ async def generate_sub_module_documentation( str_replace_editor_tool, generate_sub_module_documentation_tool, ], + retries=ctx.deps.config.max_retries, ) else: sub_agent = Agent( @@ -101,6 +102,7 @@ async def generate_sub_module_documentation( module_name=sub_module_name, custom_instructions=ctx.deps.custom_instructions ), tools=[read_code_components_tool, str_replace_editor_tool], + retries=ctx.deps.config.max_retries, ) deps.current_module_name = sub_module_name diff --git a/codewiki/src/be/pydantic_ai_backend.py b/codewiki/src/be/pydantic_ai_backend.py index 23ad08e9..599e7aae 100644 --- a/codewiki/src/be/pydantic_ai_backend.py +++ b/codewiki/src/be/pydantic_ai_backend.py @@ -81,6 +81,7 @@ async def run_update_agent( deps_type=CodeWikiDeps, tools=[read_code_components_tool, str_replace_editor_tool], system_prompt=system_prompt, + retries=self._config.max_retries, ) started = time.time() result = await agent.run(user_prompt, deps=deps) @@ -127,6 +128,7 @@ async def run_module_agent( generate_sub_module_documentation_tool, ], system_prompt=format_system_prompt(module_name, self._custom_instructions), + retries=config.max_retries, ) else: agent = Agent( @@ -135,6 +137,7 @@ async def run_module_agent( deps_type=CodeWikiDeps, tools=[read_code_components_tool, str_replace_editor_tool], system_prompt=format_leaf_system_prompt(module_name, self._custom_instructions), + retries=config.max_retries, ) deps = CodeWikiDeps( diff --git a/codewiki/src/config.py b/codewiki/src/config.py index b98208ed..d2afa277 100644 --- a/codewiki/src/config.py +++ b/codewiki/src/config.py @@ -85,6 +85,9 @@ class Config: aws_region: str = "us-east-1" api_version: str = "2024-12-01-preview" # Azure OpenAI API version azure_deployment: str = "" # Azure OpenAI deployment name + # Number of times an agent may retry a tool call whose arguments fail + # validation before giving up (default: 3) + max_retries: int = 3 # Max token settings max_tokens: int = DEFAULT_MAX_TOKENS max_token_per_module: int = DEFAULT_MAX_TOKEN_PER_MODULE @@ -211,6 +214,7 @@ def from_cli( aws_region: str = "us-east-1", api_version: str = "2024-12-01-preview", azure_deployment: str = "", + max_retries: int = 3, max_tokens: int = DEFAULT_MAX_TOKENS, max_token_per_module: int = DEFAULT_MAX_TOKEN_PER_MODULE, max_token_per_leaf_module: int = DEFAULT_MAX_TOKEN_PER_LEAF_MODULE, @@ -239,6 +243,7 @@ def from_cli( aws_region: AWS region for Bedrock provider api_version: Azure OpenAI API version azure_deployment: Azure OpenAI deployment name + max_retries: Maximum number of retries for failed agent tool calls (default: 3) max_tokens: Maximum tokens for LLM response max_token_per_module: Maximum tokens per module for clustering max_token_per_leaf_module: Maximum tokens per leaf module @@ -276,6 +281,7 @@ def from_cli( aws_region=aws_region, api_version=api_version, azure_deployment=azure_deployment, + max_retries=max_retries, max_tokens=max_tokens, max_token_per_module=max_token_per_module, max_token_per_leaf_module=max_token_per_leaf_module, diff --git a/tests/test_max_retries.py b/tests/test_max_retries.py new file mode 100644 index 00000000..f8ca101a --- /dev/null +++ b/tests/test_max_retries.py @@ -0,0 +1,97 @@ +"""Tests for max_retries PydanticAI agent configuration.""" + +from unittest.mock import patch + +from codewiki.cli.models.config import Configuration +from codewiki.src.config import Config + +from pydantic_ai import Agent, ModelRetry +from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart +from pydantic_ai.models.function import FunctionModel + + +def _make_config(max_retries: int | None = None) -> Config: + """Build a minimal Config for testing.""" + kwargs = { + "repo_path": "/tmp/repo", + "output_dir": "/tmp/out", + "llm_base_url": "http://localhost:1/v1", + "llm_api_key": "test-key", + "main_model": "test-model", + "cluster_model": "test-model", + } + + if max_retries is not None: + kwargs["max_retries"] = max_retries + + return Config.from_cli(**kwargs) + + +def test_config_max_retries_defaults_to_three(): + """Config defaults max_retries to 3.""" + config = _make_config() + assert config.max_retries == 3 + + +def test_config_accepts_max_retries(): + """Config accepts a custom max_retries value.""" + config = _make_config(max_retries=5) + assert config.max_retries == 5 + + +def test_configuration_round_trip_preserves_max_retries(): + """Configuration serialization preserves max_retries.""" + config = Configuration( + base_url="http://localhost:1/v1", + main_model="test-model", + cluster_model="test-model", + max_retries=7, + ) + + restored = Configuration.from_dict(config.to_dict()) + assert restored.max_retries == 7 + + +def test_agent_retries_failed_tool_call(): + """Agent retries a tool call after ModelRetry.""" + tool_attempts = 0 + model_calls = 0 + + def model_function(messages, info): + nonlocal model_calls + model_calls += 1 + + if model_calls <= 2: + return ModelResponse( + parts=[ + ToolCallPart( + tool_name="test_tool", + args={}, + ) + ] + ) + + return ModelResponse( + parts=[TextPart("success")], + ) + + agent = Agent( + FunctionModel(model_function), + retries=1, + ) + + @agent.tool_plain + def test_tool() -> str: + nonlocal tool_attempts + tool_attempts += 1 + + if tool_attempts == 1: + raise ModelRetry("Intentional failure") + + return "success" + + result = agent.run_sync("Call the test tool") + + assert result.output == "success" + assert tool_attempts == 2 + assert model_calls == 3 From 595e4d5c2e50e3d7dcd72ea37f80093dae799939 Mon Sep 17 00:00:00 2001 From: Perceval <36658095+perce7al@users.noreply.github.com> Date: Sun, 27 Sep 2026 10:19:03 +0300 Subject: [PATCH 2/3] test: add max_retries to sub-module dedupe test mock config --- tests/test_sub_module_dedupe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sub_module_dedupe.py b/tests/test_sub_module_dedupe.py index 83a78a3c..2816b381 100644 --- a/tests/test_sub_module_dedupe.py +++ b/tests/test_sub_module_dedupe.py @@ -53,7 +53,7 @@ def _deps(tmp_path) -> CodeWikiDeps: module_tree={}, max_depth=2, current_depth=1, - config=SimpleNamespace(max_token_per_leaf_module=4000), + config=SimpleNamespace(max_token_per_leaf_module=4000, max_retries=3), custom_instructions="", ) From 4afc6c9b7815ffedeed2faa36314c6f77d548c6b Mon Sep 17 00:00:00 2001 From: Perceval <36658095+perce7al@users.noreply.github.com> Date: Sun, 27 Sep 2026 10:20:41 +0300 Subject: [PATCH 3/3] chore: fix linting issues --- tests/test_max_retries.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_max_retries.py b/tests/test_max_retries.py index f8ca101a..6c8e7da0 100644 --- a/tests/test_max_retries.py +++ b/tests/test_max_retries.py @@ -1,7 +1,5 @@ """Tests for max_retries PydanticAI agent configuration.""" -from unittest.mock import patch - from codewiki.cli.models.config import Configuration from codewiki.src.config import Config