Merge pull request #137 from FSoft-AI4Code/feat/generate-model-flags

generate: add --main-model, --cluster-model and --fallback-model flags

Author: anhnh2002

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"),

tests/test_model_flags.py

@@ -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",
+    )