beyhan/graphify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add model override for community labeling

Browse Source
This commit is contained in:
hanmo1
2026-06-12 17:57:19 +08:00
parent 1bb30fcc56
commit b304331d24
3 changed files with 100 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
graphify/__main__.py
+13 -1
View File
@@ -2141,8 +2141,10 @@ def main() -> None:
print(" --graph <path> path to graph.json (default <path>/graphify-out/graph.json)")
print(" --no-label keep 'Community N' placeholders (skip LLM community naming)")
print(" --backend=<name> backend to use for community naming (default: auto-detect)")
print(" --model=<name> model to use for community naming")
print(" label <path> (re)name communities with the configured LLM backend, regenerate report")
print(" --backend=<name> backend to use (default: auto-detect from API keys)")
print(" --model=<name> model to use for community naming")
print(" query \"<question>\" 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}
graphify/llm.py
+19 -5
View File
@@ -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:
tests/test_labeling.py
+68
View File
@@ -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",