From 458a7b1d3fcf9386953a0426d063de67a763f63d Mon Sep 17 00:00:00 2001 From: anhnh2002 Date: Thu, 15 Oct 2026 19:58:19 +0700 Subject: [PATCH 1/2] generate: add --main-model, --cluster-model and --fallback-model flags Callers that run codewiki for one repo (e.g. a migration pipeline) can now pick the models per run instead of editing the global ~/.codewiki/config.json. The flags override the saved config for this run only. --- codewiki/cli/commands/generate.py | 31 +++++++++++++++-- tests/test_model_flags.py | 58 +++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 3 deletions(-) create mode 100644 tests/test_model_flags.py diff --git a/codewiki/cli/commands/generate.py b/codewiki/cli/commands/generate.py index d5d057f6..d2122edc 100644 --- a/codewiki/cli/commands/generate.py +++ b/codewiki/cli/commands/generate.py @@ -338,6 +338,21 @@ def _find_affected(tree, parent_names=None): help="Add prompt-cache breakpoints to agentic LLM calls; auto-falls back to " "normal calls if the provider rejects them (default: enabled)", ) +@click.option( + "--main-model", + default=None, + help="Model for documentation agents (overrides config)", +) +@click.option( + "--cluster-model", + default=None, + help="Model for module clustering (overrides config)", +) +@click.option( + "--fallback-model", + default=None, + help="Model used when the main model fails (overrides config)", +) @click.option( "--artifacts/--no-artifacts", default=True, @@ -454,6 +469,9 @@ def generate_command( request_limit: int | None, agent_retries: int | None, prompt_caching: bool | None, + main_model: str | None = None, + cluster_model: str | None = None, + fallback_model: str | None = None, artifacts: bool = True, artifact_token_budget: int = 200_000, with_prose: bool = False, @@ -518,6 +536,10 @@ def generate_command( \b # Override max depth for hierarchical decomposition $ codewiki generate --max-depth 3 + + \b + # Use other models for this generation + $ codewiki generate --main-model claude-haiku-5-5 --cluster-model claude-haiku-5-5 """ print_banner() logger = create_logger(verbose=verbose) @@ -724,6 +746,8 @@ def generate_command( effective_prompt_caching = ( prompt_caching if prompt_caching is not None else config.prompt_caching ) + logger.debug(f"Main model: {main_model or config.main_model}") + logger.debug(f"Cluster model: {cluster_model or config.cluster_model}") logger.debug(f"Max tokens: {effective_max_tokens}") logger.debug(f"Max token/module: {effective_max_token_per_module}") logger.debug(f"Max token/leaf module: {effective_max_token_per_leaf}") @@ -761,9 +785,10 @@ def generate_command( repo_path=repo_path, output_dir=output_dir, config={ - "main_model": config.main_model, - "cluster_model": config.cluster_model, - "fallback_model": config.fallback_model, + # Models (runtime overrides take precedence) + "main_model": main_model or config.main_model, + "cluster_model": cluster_model or config.cluster_model, + "fallback_model": fallback_model or config.fallback_model, "base_url": config.base_url, "api_key": api_key, "provider": getattr(config, "provider", "openai-compatible"), diff --git a/tests/test_model_flags.py b/tests/test_model_flags.py new file mode 100644 index 00000000..ba958aee --- /dev/null +++ b/tests/test_model_flags.py @@ -0,0 +1,58 @@ +"""`codewiki generate --main-model/--cluster-model/--fallback-model` override the saved config.""" + +from types import SimpleNamespace + +import pytest +from click.testing import CliRunner + +import codewiki.cli.commands.generate as gen + + +class Stop(Exception): + pass + + +@pytest.fixture +def captured(tmp_path, monkeypatch): + seen = {} + cfg = SimpleNamespace( + main_model="saved-main", cluster_model="saved-cluster", fallback_model="saved-fallback", + base_url="http://x", provider="openai-compatible", aws_region="us-east-1", + agent_instructions=None, max_tokens=1, max_token_per_module=1, max_token_per_leaf_module=1, + max_depth=1, request_limit=1, agent_retries=1, use_gitignore=True, prompt_caching=False, + ) + + class FakeManager: + def load(self): + return True + + def is_configured(self): + return True + + def get_config(self): + return cfg + + def get_api_key(self): + return "k" + + def fake_generator(**kwargs): + seen.update(kwargs["config"]) + raise Stop + + monkeypatch.setattr(gen, "ConfigManager", FakeManager) + monkeypatch.setattr(gen, "CLIDocumentationGenerator", fake_generator) + (tmp_path / "a.py").write_text("def f():\n return 1\n") + monkeypatch.chdir(tmp_path) + return seen + + +def test_saved_models_by_default(captured): + CliRunner().invoke(gen.generate_command, ["-o", "out"]) + assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == ( + "saved-main", "saved-cluster", "saved-fallback") + + +def test_flags_override_saved_models(captured): + CliRunner().invoke(gen.generate_command, ["-o", "out", "--main-model", "m", "--cluster-model", "c", + "--fallback-model", "f"]) + assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == ("m", "c", "f") From df96a09585ce27b7201b59cbb30c1d2f8bdeaedb Mon Sep 17 00:00:00 2001 From: anhnh2002 Date: Thu, 15 Oct 2026 20:20:36 +0700 Subject: [PATCH 2/2] tests: ruff format test_model_flags.py --- tests/test_model_flags.py | 36 ++++++++++++++++++++++++++++-------- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/tests/test_model_flags.py b/tests/test_model_flags.py index ba958aee..40aeec37 100644 --- a/tests/test_model_flags.py +++ b/tests/test_model_flags.py @@ -16,10 +16,21 @@ class Stop(Exception): def captured(tmp_path, monkeypatch): seen = {} cfg = SimpleNamespace( - main_model="saved-main", cluster_model="saved-cluster", fallback_model="saved-fallback", - base_url="http://x", provider="openai-compatible", aws_region="us-east-1", - agent_instructions=None, max_tokens=1, max_token_per_module=1, max_token_per_leaf_module=1, - max_depth=1, request_limit=1, agent_retries=1, use_gitignore=True, prompt_caching=False, + main_model="saved-main", + cluster_model="saved-cluster", + fallback_model="saved-fallback", + base_url="http://x", + provider="openai-compatible", + aws_region="us-east-1", + agent_instructions=None, + max_tokens=1, + max_token_per_module=1, + max_token_per_leaf_module=1, + max_depth=1, + request_limit=1, + agent_retries=1, + use_gitignore=True, + prompt_caching=False, ) class FakeManager: @@ -49,10 +60,19 @@ def fake_generator(**kwargs): def test_saved_models_by_default(captured): CliRunner().invoke(gen.generate_command, ["-o", "out"]) assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == ( - "saved-main", "saved-cluster", "saved-fallback") + "saved-main", + "saved-cluster", + "saved-fallback", + ) def test_flags_override_saved_models(captured): - CliRunner().invoke(gen.generate_command, ["-o", "out", "--main-model", "m", "--cluster-model", "c", - "--fallback-model", "f"]) - assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == ("m", "c", "f") + CliRunner().invoke( + gen.generate_command, + ["-o", "out", "--main-model", "m", "--cluster-model", "c", "--fallback-model", "f"], + ) + assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == ( + "m", + "c", + "f", + )