diff --git a/graphify/__main__.py b/graphify/__main__.py index 9daa5f2f..758278ba 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -2141,8 +2141,10 @@ def main() -> None: print(" --graph path to graph.json (default /graphify-out/graph.json)") print(" --no-label keep 'Community N' placeholders (skip LLM community naming)") print(" --backend= backend to use for community naming (default: auto-detect)") + print(" --model= model to use for community naming") print(" label (re)name communities with the configured LLM backend, regenerate report") print(" --backend= backend to use (default: auto-detect from API keys)") + print(" --model= model to use for community naming") print(" query \"\" BFS traversal of graph.json for a question") print(" --dfs use depth-first instead of breadth-first") print(" --context C explicit edge-context filter (repeatable)") @@ -3131,6 +3133,8 @@ def main() -> None: no_label = "--no-label" in sys.argv _backend_arg = next((a for a in sys.argv if a.startswith("--backend=")), None) label_backend = _backend_arg.split("=", 1)[1] if _backend_arg else None + _model_arg = next((a for a in sys.argv if a.startswith("--model=")), None) + label_model = _model_arg.split("=", 1)[1] if _model_arg else None _min_cs_arg = next((a for a in sys.argv if a.startswith("--min-community-size=")), None) min_community_size = int(_min_cs_arg.split("=")[1]) if _min_cs_arg else 3 args = sys.argv[2:] @@ -3143,6 +3147,14 @@ def main() -> None: a = args[i_arg] if a == "--graph" and i_arg + 1 < len(args): graph_override = Path(args[i_arg + 1]); i_arg += 2 + elif a == "--backend" and i_arg + 1 < len(args): + label_backend = args[i_arg + 1]; i_arg += 2 + elif a.startswith("--backend="): + label_backend = a.split("=", 1)[1]; i_arg += 1 + elif a == "--model" and i_arg + 1 < len(args): + label_model = args[i_arg + 1]; i_arg += 2 + elif a.startswith("--model="): + label_model = a.split("=", 1)[1]; i_arg += 1 elif a == "--resolution" and i_arg + 1 < len(args): co_resolution = float(args[i_arg + 1]); i_arg += 2 elif a.startswith("--resolution="): @@ -3240,7 +3252,7 @@ def main() -> None: # The final labels (LLM or placeholder fallback) are persisted to # .graphify_labels.json by the unconditional write below. labels, _ = generate_community_labels( - G, communities, backend=label_backend, gods=gods + G, communities, backend=label_backend, model=label_model, gods=gods ) questions = suggest_questions(G, communities, labels) tokens = {"input": 0, "output": 0} diff --git a/graphify/llm.py b/graphify/llm.py index 2f56b7a7..44628c0b 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -1716,7 +1716,13 @@ def _merge_into(merged: dict, result: dict) -> None: merged["output_tokens"] += result.get("output_tokens", 0) -def _call_llm(prompt: str, *, backend: str, max_tokens: int = 200) -> str: +def _call_llm( + prompt: str, + *, + backend: str, + max_tokens: int = 200, + model: str | None = None, +) -> str: """Send a plain-text prompt to `backend` and return the model's text reply. Used by lightweight callers (e.g. `graphify.dedup` LLM tiebreaker) that @@ -1740,7 +1746,7 @@ def _call_llm(prompt: str, *, backend: str, max_tokens: int = 200) -> str: raise ValueError( f"No API key for backend '{backend}'. Set {_format_backend_env_keys(backend)}." ) - mdl = _default_model_for_backend(backend) + mdl = model or _default_model_for_backend(backend) if backend == "claude": try: @@ -1769,8 +1775,11 @@ def _call_llm(prompt: str, *, backend: str, max_tokens: int = 200) -> str: raise RuntimeError("Claude Code CLI not found on $PATH") elif shutil.which("claude") is None: raise RuntimeError("Claude Code CLI not found on $PATH") + cli_args = [claude_cmd, "-p", "--output-format", "json", "--no-session-persistence"] + if model is not None: + cli_args.extend(["--model", mdl]) proc = subprocess.run( - [claude_cmd, "-p", "--output-format", "json", "--no-session-persistence"], + cli_args, input=prompt, capture_output=True, text=True, @@ -2033,6 +2042,7 @@ def label_communities( communities, *, backend: str, + model: str | None = None, gods=None, max_communities: int | None = None, top_k: int = _LABEL_TOP_K, @@ -2084,7 +2094,10 @@ def label_communities( # _resolve_max_tokens so GRAPHIFY_MAX_OUTPUT_TOKENS applies here too (#1200). max_tokens = _resolve_max_tokens(min(64 + 24 * len(batch_cids), 8192)) try: - text = _call_llm(prompt, backend=backend, max_tokens=max_tokens) + call_kwargs = {"backend": backend, "max_tokens": max_tokens} + if model is not None: + call_kwargs["model"] = model + text = _call_llm(prompt, **call_kwargs) parsed = _parse_label_response(text, batch_cids) labels.update(parsed) written += len(parsed) @@ -2109,6 +2122,7 @@ def generate_community_labels( communities, *, backend: str | None = None, + model: str | None = None, gods=None, quiet: bool = False, ) -> tuple[dict[int, str], str]: @@ -2130,7 +2144,7 @@ def generate_community_labels( ) return _placeholder_community_labels(communities), "placeholder" try: - labels = label_communities(G, communities, backend=backend, gods=gods) + labels = label_communities(G, communities, backend=backend, model=model, gods=gods) return labels, "llm" except Exception as exc: if not quiet: diff --git a/tests/test_labeling.py b/tests/test_labeling.py index 97c0d311..a228b6d7 100644 --- a/tests/test_labeling.py +++ b/tests/test_labeling.py @@ -3,6 +3,9 @@ Backend calls are mocked - no network. Covers the happy path, partial replies, malformed replies, and the no-backend fallback. """ +import json +import sys + import networkx as nx import pytest @@ -40,6 +43,71 @@ def test_label_communities_happy_path(monkeypatch): assert captured["backend"] == "gemini" +def test_label_communities_passes_model_override(monkeypatch): + G, communities = _graph() + captured = {} + + def fake_call(prompt, *, backend, max_tokens=200, model=None): + captured["backend"] = backend + captured["model"] = model + return '{"0": "Order Management", "1": "Payment Flow"}' + + monkeypatch.setattr("graphify.llm._call_llm", fake_call) + labels = label_communities( + G, + communities, + backend="gemini", + model="gemini-3.1-flash-lite", + ) + + assert labels == {0: "Order Management", 1: "Payment Flow"} + assert captured == {"backend": "gemini", "model": "gemini-3.1-flash-lite"} + + +def test_label_cli_passes_model_override(tmp_path, monkeypatch): + import graphify.__main__ as cli + + out = tmp_path / "graphify-out" + out.mkdir() + graph = { + "directed": False, + "multigraph": False, + "nodes": [ + {"id": "n1", "label": "OrderService", "community": 0}, + ], + "links": [], + } + (out / "graph.json").write_text(json.dumps(graph), encoding="utf-8") + + captured = {} + + def fake_generate(G, communities, *, backend=None, model=None, gods=None, quiet=False): + captured["backend"] = backend + captured["model"] = model + return {0: "Orders"}, "llm" + + monkeypatch.setattr("graphify.llm.generate_community_labels", fake_generate) + monkeypatch.setattr("graphify.export.to_html", lambda *args, **kwargs: None) + monkeypatch.setattr( + sys, + "argv", + [ + "graphify", + "label", + str(tmp_path), + "--backend", + "gemini", + "--model", + "gemini-3.1-flash-lite", + "--no-viz", + ], + ) + + cli.main() + + assert captured == {"backend": "gemini", "model": "gemini-3.1-flash-lite"} + + def test_label_communities_partial_reply_fills_placeholder(monkeypatch): G, communities = _graph() monkeypatch.setattr("graphify.llm._call_llm",