Skip to content

Commit 458a7b1

Browse files
committed
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.
1 parent 2b99143 commit 458a7b1

2 files changed

Lines changed: 86 additions & 3 deletions

File tree

‎codewiki/cli/commands/generate.py‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,21 @@ def _find_affected(tree, parent_names=None):
338338
help="Add prompt-cache breakpoints to agentic LLM calls; auto-falls back to "
339339
"normal calls if the provider rejects them (default: enabled)",
340340
)
341+
@click.option(
342+
"--main-model",
343+
default=None,
344+
help="Model for documentation agents (overrides config)",
345+
)
346+
@click.option(
347+
"--cluster-model",
348+
default=None,
349+
help="Model for module clustering (overrides config)",
350+
)
351+
@click.option(
352+
"--fallback-model",
353+
default=None,
354+
help="Model used when the main model fails (overrides config)",
355+
)
341356
@click.option(
342357
"--artifacts/--no-artifacts",
343358
default=True,
@@ -454,6 +469,9 @@ def generate_command(
454469
request_limit: int | None,
455470
agent_retries: int | None,
456471
prompt_caching: bool | None,
472+
main_model: str | None = None,
473+
cluster_model: str | None = None,
474+
fallback_model: str | None = None,
457475
artifacts: bool = True,
458476
artifact_token_budget: int = 200_000,
459477
with_prose: bool = False,
@@ -518,6 +536,10 @@ def generate_command(
518536
\b
519537
# Override max depth for hierarchical decomposition
520538
$ codewiki generate --max-depth 3
539+
540+
\b
541+
# Use other models for this generation
542+
$ codewiki generate --main-model claude-haiku-5-5 --cluster-model claude-haiku-5-5
521543
"""
522544
print_banner()
523545
logger = create_logger(verbose=verbose)
@@ -724,6 +746,8 @@ def generate_command(
724746
effective_prompt_caching = (
725747
prompt_caching if prompt_caching is not None else config.prompt_caching
726748
)
749+
logger.debug(f"Main model: {main_model or config.main_model}")
750+
logger.debug(f"Cluster model: {cluster_model or config.cluster_model}")
727751
logger.debug(f"Max tokens: {effective_max_tokens}")
728752
logger.debug(f"Max token/module: {effective_max_token_per_module}")
729753
logger.debug(f"Max token/leaf module: {effective_max_token_per_leaf}")
@@ -761,9 +785,10 @@ def generate_command(
761785
repo_path=repo_path,
762786
output_dir=output_dir,
763787
config={
764-
"main_model": config.main_model,
765-
"cluster_model": config.cluster_model,
766-
"fallback_model": config.fallback_model,
788+
# Models (runtime overrides take precedence)
789+
"main_model": main_model or config.main_model,
790+
"cluster_model": cluster_model or config.cluster_model,
791+
"fallback_model": fallback_model or config.fallback_model,
767792
"base_url": config.base_url,
768793
"api_key": api_key,
769794
"provider": getattr(config, "provider", "openai-compatible"),

‎tests/test_model_flags.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
"""`codewiki generate --main-model/--cluster-model/--fallback-model` override the saved config."""
2+
3+
from types import SimpleNamespace
4+
5+
import pytest
6+
from click.testing import CliRunner
7+
8+
import codewiki.cli.commands.generate as gen
9+
10+
11+
class Stop(Exception):
12+
pass
13+
14+
15+
@pytest.fixture
16+
def captured(tmp_path, monkeypatch):
17+
seen = {}
18+
cfg = SimpleNamespace(
19+
main_model="saved-main", cluster_model="saved-cluster", fallback_model="saved-fallback",
20+
base_url="http://x", provider="openai-compatible", aws_region="us-east-1",
21+
agent_instructions=None, max_tokens=1, max_token_per_module=1, max_token_per_leaf_module=1,
22+
max_depth=1, request_limit=1, agent_retries=1, use_gitignore=True, prompt_caching=False,
23+
)
24+
25+
class FakeManager:
26+
def load(self):
27+
return True
28+
29+
def is_configured(self):
30+
return True
31+
32+
def get_config(self):
33+
return cfg
34+
35+
def get_api_key(self):
36+
return "k"
37+
38+
def fake_generator(**kwargs):
39+
seen.update(kwargs["config"])
40+
raise Stop
41+
42+
monkeypatch.setattr(gen, "ConfigManager", FakeManager)
43+
monkeypatch.setattr(gen, "CLIDocumentationGenerator", fake_generator)
44+
(tmp_path / "a.py").write_text("def f():\n return 1\n")
45+
monkeypatch.chdir(tmp_path)
46+
return seen
47+
48+
49+
def test_saved_models_by_default(captured):
50+
CliRunner().invoke(gen.generate_command, ["-o", "out"])
51+
assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == (
52+
"saved-main", "saved-cluster", "saved-fallback")
53+
54+
55+
def test_flags_override_saved_models(captured):
56+
CliRunner().invoke(gen.generate_command, ["-o", "out", "--main-model", "m", "--cluster-model", "c",
57+
"--fallback-model", "f"])
58+
assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == ("m", "c", "f")

0 commit comments

Comments
 (0)