feat: extra_body for custom providers + multi-batch label_communities (#1197)

Adds extra_body parameter support for custom/OpenAI-compat providers so users can pass provider-specific params (e.g. thinking budget for Claude via Bedrock compat). Adds multi-batch label_communities for 16k-context models — batches multiple community descriptions into a single LLM call instead of one per community. Partial batch failures are handled gracefully.

Co-authored-by: EirikWolf <EirikWolf@users.noreply.github.com>

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 Full release notes with details on each version: [GitHub Releases](https://github.com/safishamsi/graphify/releases)
 
+## Unreleased
+
+- Feat: `extra_body` field in `providers.json` is now sent to the OpenAI-compatible client at both the extraction and labeling code paths. Lets self-hosted OpenAI-compatible endpoints (vLLM serving Qwen3, Llama-3.1, etc.) pass model-specific request shapes — most importantly `{"chat_template_kwargs": {"enable_thinking": false}}` for Qwen3, which otherwise emits chain-of-thought instead of the JSON the parser expects. An explicit `extra_body` from a custom provider also bypasses the Ollama `num_ctx` auto-derive, so a provider that points at Ollama can opt out of that default if it knows better.
+- Feat: `label_communities` now batches communities in chunks of 100 (configurable via `batch_size=`) instead of a single call capped at 200. This was the only obstacle to running community labeling against any backend with a 16k token context window (Qwen3.6-27B served by vLLM, Llama-3.1 8B-Instruct, etc.): 200 communities × 12 sampled node labels routinely overflowed the prompt. The default `max_communities` is now `None` (label them all); explicit integer caps still work for back-compat. Partial batch failures no longer drop the whole pass — successful batches still contribute labels, only the failed batch's communities stay as placeholders.
+
 ## 0.8.35 (2026-06-07)
 
 - Feat: CodeBuddy platform support. `graphify codebuddy install` installs the graphify skill to `~/.codebuddy/skills/graphify/SKILL.md`, writes a `CODEBUDDY.md` always-on section, and registers Bash + Read|Glob PreToolUse hooks in `.codebuddy/settings.json` that nudge the agent toward `graphify query` instead of grepping raw files when a graph exists. `graphify install --platform codebuddy` and `graphify codebuddy uninstall` also supported. Thanks to @studyzy (#1136).

graphify/llm.py

@@ -689,6 +689,7 @@ def _call_openai_compat(
     backend: str = "",
     deep_mode: bool = False,
     images: list[_ImageRef] | None = None,
+    extra_body: dict | None = None,
 ) -> dict:
     """Call any OpenAI-compatible API (Kimi, OpenAI, etc.) and return parsed JSON."""
     try:
@@ -715,8 +716,14 @@ def _call_openai_compat(
         kwargs["temperature"] = temperature
     if reasoning_effort is not None:
         kwargs["reasoning_effort"] = reasoning_effort
+    # A custom provider in providers.json can pass its own extra_body (e.g.
+    # `chat_template_kwargs.enable_thinking=false` for self-hosted Qwen3 served
+    # by vLLM). When supplied, it wins over the moonshot default — the user has
+    # explicitly chosen the request shape for their endpoint.
+    if extra_body is not None:
+        kwargs["extra_body"] = extra_body
     # Kimi-k2.6 is a reasoning model — disable thinking so content isn't empty
-    if "moonshot" in base_url:
+    elif "moonshot" in base_url:
         kwargs["extra_body"] = {"thinking": {"type": "disabled"}}
     # Ollama defaults num_ctx to 2048 and silently truncates prompts larger
     # than that — the symptom is hollow 200 OK responses after the first few
@@ -726,7 +733,9 @@ def _call_openai_compat(
     # hollow-200 symptom — just from a different direction (#798 follow-up).
     # Formula: actual input tokens + output cap + system prompt headroom.
     # Capped at 131072 (enough for the default 60k token_budget); env var wins.
-    if backend == "ollama":
+    # The ollama num_ctx auto-derive is a default. A custom provider that
+    # explicitly sets extra_body has opted out — respect their request shape.
+    if backend == "ollama" and extra_body is None:
         num_ctx_raw = os.environ.get("GRAPHIFY_OLLAMA_NUM_CTX", "").strip()
         # Auto-derive num_ctx from actual chunk size regardless — used as the
         # fallback and for the mismatch check below.
@@ -1153,6 +1162,7 @@ def extract_files_direct(
         backend=backend,
         deep_mode=deep_mode,
         images=image_refs,
+        extra_body=cfg.get("extra_body"),
     )
 
 
@@ -1638,7 +1648,11 @@ def _call_llm(prompt: str, *, backend: str, max_tokens: int = 200) -> str:
         kwargs["temperature"] = temperature
     if cfg.get("reasoning_effort"):
         kwargs["reasoning_effort"] = cfg["reasoning_effort"]
-    if "moonshot" in cfg["base_url"]:
+    # Custom providers can override via providers.json `extra_body`; falls back
+    # to the moonshot default to preserve existing behavior.
+    if cfg.get("extra_body") is not None:
+        kwargs["extra_body"] = cfg["extra_body"]
+    elif "moonshot" in cfg["base_url"]:
         kwargs["extra_body"] = {"thinking": {"type": "disabled"}}
     resp = client.chat.completions.create(**kwargs)
     if not resp.choices or resp.choices[0].message is None:
@@ -1767,9 +1781,10 @@ def detect_backend() -> str | None:
 # batched call and return a complete ``{cid: name}`` map (#1097).
 
 _LABEL_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*|\s*```\s*$", re.IGNORECASE)
-_LABEL_MAX_COMMUNITIES = 200   # cap LLM-named communities; tail stays placeholder
+_LABEL_MAX_COMMUNITIES = 200   # legacy soft-cap; kept for callers that pin it.
 _LABEL_TOP_K = 12              # node labels sampled per community for the prompt
 _LABEL_MAXLEN = 60             # truncate individual labels to keep the prompt small
+_LABEL_BATCH_SIZE = 100        # communities per LLM call; sized for ~16k context windows
 
 
 def _placeholder_community_labels(communities) -> dict[int, str]:
@@ -1830,31 +1845,70 @@ def label_communities(
     *,
     backend: str,
     gods=None,
-    max_communities: int = _LABEL_MAX_COMMUNITIES,
+    max_communities: int | None = None,
     top_k: int = _LABEL_TOP_K,
+    batch_size: int = _LABEL_BATCH_SIZE,
 ) -> dict[int, str]:
     """Return a complete ``{cid: name}`` map using ``backend`` for naming.
 
-    Placeholders (``Community N``) are used for any community the backend did not
-    name. Raises on backend/parse failure - callers that want graceful
-    degradation should use :func:`generate_community_labels`.
+    Communities are labeled in batches of ``batch_size`` so the prompt fits in a
+    16k-token context window (which is enough for one batch of ~100 communities
+    × ``top_k`` node labels). With the previous hard cap of 200 communities in a
+    single call, self-hosted 16k models (Qwen3, Llama 3.1 8B-Instruct, etc.)
+    routinely overflowed context and dropped the entire labeling pass to
+    placeholders.
+
+    ``max_communities=None`` (the default) labels every community. Pass an
+    integer to cap the total (the legacy 200 default preserved this behavior;
+    explicit callers can still pin it). Placeholders (``Community N``) are used
+    for any community the backend did not name. Per-batch failures are logged
+    to stderr and skipped — the surviving batches still contribute labels.
+
+    Raises on the first batch's backend/parse failure if it leaves *no* labels
+    written. Callers that want graceful degradation should use
+    :func:`generate_community_labels`.
     """
     labels = _placeholder_community_labels(communities)
-    lines, labeled_cids = _community_label_lines(G, communities, gods, max_communities, top_k)
+    cap = len(communities) if max_communities is None else max_communities
+    lines, labeled_cids = _community_label_lines(G, communities, gods, cap, top_k)
     if not lines:
         return labels
 
-    prompt = (
-        "You are naming clusters in a knowledge graph. For each community below, "
-        "return a concise 2-5 word plain-language name describing what it is about "
-        "(e.g. \"Order Management\", \"Payment Flow\", \"Auth Middleware\"). "
-        "Respond ONLY with a JSON object mapping the community id (as a string) to "
-        "its name - no prose, no markdown fences.\n\n" + "\n".join(lines)
-    )
+    n_batches = (len(labeled_cids) + batch_size - 1) // batch_size
+    written = 0
+    first_error: Exception | None = None
+    for batch_idx in range(n_batches):
+        start = batch_idx * batch_size
+        end = min(start + batch_size, len(labeled_cids))
+        batch_lines = lines[start:end]
+        batch_cids = labeled_cids[start:end]
 
-    max_tokens = min(40 + 16 * len(labeled_cids), 4096)
-    text = _call_llm(prompt, backend=backend, max_tokens=max_tokens)
-    labels.update(_parse_label_response(text, labeled_cids))
+        prompt = (
+            "You are naming clusters in a knowledge graph. For each community below, "
+            "return a concise 2-5 word plain-language name describing what it is about "
+            "(e.g. \"Order Management\", \"Payment Flow\", \"Auth Middleware\"). "
+            "Respond ONLY with a JSON object mapping the community id (as a string) to "
+            "its name - no prose, no markdown fences.\n\n" + "\n".join(batch_lines)
+        )
+        max_tokens = min(40 + 16 * len(batch_cids), 4096)
+        try:
+            text = _call_llm(prompt, backend=backend, max_tokens=max_tokens)
+            parsed = _parse_label_response(text, batch_cids)
+            labels.update(parsed)
+            written += len(parsed)
+        except Exception as exc:
+            if first_error is None:
+                first_error = exc
+            print(
+                f"[graphify label] batch {batch_idx + 1}/{n_batches} "
+                f"({len(batch_cids)} communities) failed: {exc}",
+                file=sys.stderr,
+            )
+            continue
+
+    if written == 0 and first_error is not None:
+        # Every batch failed; propagate so generate_community_labels degrades cleanly.
+        raise first_error
     return labels
 
 

tests/test_labeling.py

@@ -117,3 +117,96 @@ def test_empty_communities_returns_placeholders(monkeypatch):
     labels = label_communities(G, {0: []}, backend="gemini")
     assert labels == {0: "Community 0"}
     assert called is False
+
+
+# ---------------------------------------------------------------------------
+# Multi-batch labeling: a single prompt with >100 communities overflows the
+# 16k context window of self-hosted reasoning models (Qwen3, Llama-3.1 8B).
+# label_communities now splits into batches so coverage stays complete.
+# ---------------------------------------------------------------------------
+
+
+def _wide_graph(n_communities: int):
+    G = nx.Graph()
+    communities: dict[int, list[str]] = {}
+    for cid in range(n_communities):
+        a, b = f"c{cid}_a", f"c{cid}_b"
+        G.add_node(a, label=f"node_{cid}_a")
+        G.add_node(b, label=f"node_{cid}_b")
+        communities[cid] = [a, b]
+    return G, communities
+
+
+def test_label_communities_batches_when_over_batch_size(monkeypatch):
+    G, communities = _wide_graph(250)
+    calls = []
+
+    def fake_call(prompt, *, backend, max_tokens=200):
+        # The fake reads which cids the prompt asks about and answers all of them.
+        cids = [int(line.split(":", 1)[0].removeprefix("Community ").strip())
+                for line in prompt.splitlines() if line.startswith("Community ")]
+        calls.append(len(cids))
+        return "{" + ", ".join(f'"{c}": "Cluster {c}"' for c in cids) + "}"
+
+    monkeypatch.setattr("graphify.llm._call_llm", fake_call)
+    labels = label_communities(G, communities, backend="gemini", batch_size=100)
+
+    # 250 communities / 100 per batch -> 3 batches (100, 100, 50)
+    assert calls == [100, 100, 50]
+    # And every community got a real name, none left as a placeholder.
+    assert all(name.startswith("Cluster ") for name in labels.values()), \
+        f"some communities still have placeholders: {[k for k, v in labels.items() if not v.startswith('Cluster ')][:5]}"
+    assert len(labels) == 250
+
+
+def test_label_communities_partial_batch_failure_keeps_successful_batches(monkeypatch):
+    G, communities = _wide_graph(150)
+    n_calls = [0]
+
+    def fake_call(prompt, *, backend, max_tokens=200):
+        n_calls[0] += 1
+        cids = [int(line.split(":", 1)[0].removeprefix("Community ").strip())
+                for line in prompt.splitlines() if line.startswith("Community ")]
+        if n_calls[0] == 2:
+            raise RuntimeError("simulated transient backend failure")
+        return "{" + ", ".join(f'"{c}": "Named {c}"' for c in cids) + "}"
+
+    monkeypatch.setattr("graphify.llm._call_llm", fake_call)
+    labels = label_communities(G, communities, backend="gemini", batch_size=50)
+
+    # 3 batches; second one fails. First and third produce real labels;
+    # the failed batch's cids stay as placeholders.
+    real = [cid for cid, name in labels.items() if name.startswith("Named ")]
+    placeholder = [cid for cid, name in labels.items() if name.startswith("Community ")]
+    assert len(real) == 100, f"expected 100 real labels from 2 successful batches, got {len(real)}"
+    assert len(placeholder) == 50, f"expected 50 placeholders from the failed batch, got {len(placeholder)}"
+
+
+def test_label_communities_all_batches_fail_raises(monkeypatch):
+    G, communities = _wide_graph(150)
+
+    def always_fail(prompt, *, backend, max_tokens=200):
+        raise RuntimeError("backend down")
+
+    monkeypatch.setattr("graphify.llm._call_llm", always_fail)
+    # Every batch fails -> propagate so generate_community_labels can degrade.
+    with pytest.raises(RuntimeError, match="backend down"):
+        label_communities(G, communities, backend="gemini", batch_size=50)
+
+
+def test_label_communities_max_communities_caps_total(monkeypatch):
+    # Backwards compat: explicit max_communities still caps the total labeled,
+    # so callers that pinned the legacy 200-default keep their behavior.
+    G, communities = _wide_graph(150)
+    captured_cids = []
+
+    def fake_call(prompt, *, backend, max_tokens=200):
+        cids = [int(line.split(":", 1)[0].removeprefix("Community ").strip())
+                for line in prompt.splitlines() if line.startswith("Community ")]
+        captured_cids.extend(cids)
+        return "{" + ", ".join(f'"{c}": "X{c}"' for c in cids) + "}"
+
+    monkeypatch.setattr("graphify.llm._call_llm", fake_call)
+    label_communities(G, communities, backend="gemini", max_communities=40, batch_size=100)
+    # Only 40 communities should have been sent to the backend.
+    assert len(captured_cids) == 40

tests/test_llm_backends.py

@@ -429,6 +429,58 @@ def test_non_ollama_backend_gets_no_num_ctx_extra_body(monkeypatch):
     assert eb is None or "options" not in eb, "non-ollama backends must not get num_ctx injection"
 
 
+# ---------------------------------------------------------------------------
+# Custom-provider extra_body: lets providers.json route around the moonshot-only
+# default. Self-hosted Qwen3 served by vLLM needs
+# `chat_template_kwargs.enable_thinking=false` or the model emits chain-of-thought
+# instead of the JSON the extraction parser expects.
+# ---------------------------------------------------------------------------
+
+
+def test_call_openai_compat_uses_explicit_extra_body(monkeypatch):
+    captured = _install_capturing_openai(monkeypatch)
+
+    llm._call_openai_compat(
+        "https://kitor.example/vllm/v1", "tk", "Qwen3.6-27B",
+        "u", temperature=0, max_completion_tokens=8192, backend="kitor-vllm",
+        extra_body={"chat_template_kwargs": {"enable_thinking": False}},
+    )
+
+    assert captured["extra_body"] == {"chat_template_kwargs": {"enable_thinking": False}}
+
+
+def test_call_openai_compat_extra_body_wins_over_moonshot_default(monkeypatch):
+    # A user could legitimately set up a moonshot-compatible custom provider
+    # and want a different extra_body — explicit kwarg must override the default.
+    captured = _install_capturing_openai(monkeypatch)
+
+    llm._call_openai_compat(
+        "https://api.moonshot.ai/v1", "tk", "kimi-k2-thinking",
+        "u", temperature=0, max_completion_tokens=8192, backend="kimi",
+        extra_body={"thinking": {"type": "enabled"}},
+    )
+
+    assert captured["extra_body"] == {"thinking": {"type": "enabled"}}
+
+
+def test_call_openai_compat_explicit_extra_body_skips_ollama_auto_derive(monkeypatch):
+    # An explicit extra_body means "I own this request shape" — Ollama's
+    # num_ctx auto-derive (a default) must step aside or we'd clobber it.
+    captured = _install_capturing_openai(monkeypatch)
+    monkeypatch.delenv("GRAPHIFY_OLLAMA_NUM_CTX", raising=False)
+    monkeypatch.delenv("GRAPHIFY_OLLAMA_KEEP_ALIVE", raising=False)
+
+    llm._call_openai_compat(
+        "http://localhost:11434/v1", "ollama", "qwen2.5-coder:7b",
+        "u", temperature=0, max_completion_tokens=8192, backend="ollama",
+        extra_body={"options": {"num_ctx": 4096}},
+    )
+
+    assert captured["extra_body"] == {"options": {"num_ctx": 4096}}, (
+        "explicit extra_body must replace the ollama auto-derived num_ctx"
+    )
+
+
 def test_extract_corpus_parallel_ollama_runs_serially(tmp_path, monkeypatch):
     # With 3 chunks and backend=ollama, ThreadPoolExecutor must NOT be used
     # (workers=1 takes the sequential path). We verify by ensuring all chunks