Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 28 additions & 3 deletions codewiki/cli/commands/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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}")
Expand Down Expand Up @@ -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"),
Expand Down
78 changes: 78 additions & 0 deletions tests/test_model_flags.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
"""`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",
)
Loading