fix Ollama num_ctx: derive from actual chunk size instead of hardcoding 131072 (#798)

graphify/llm.py

@@ -281,17 +281,25 @@ def _call_openai_compat(
         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
-    # chunks (#798). We send num_ctx large enough to fit our default 60k-token
-    # chunk budget plus system prompt and JSON output headroom. Ollama caps
-    # gracefully at the model's built-in limit if it's lower, so a large
-    # default is safe. keep_alive pins the model in VRAM across chunks so it
-    # isn't unloaded/reloaded mid-run under concurrency pressure.
+    # chunks (#798). We derive num_ctx from the actual prompt size so we don't
+    # over-allocate KV-cache VRAM. Over-allocation (e.g. 128k slots for an 8k
+    # prompt on a 31B model) exhausts VRAM by chunk 4 and produces the same
+    # 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":
         num_ctx_raw = os.environ.get("GRAPHIFY_OLLAMA_NUM_CTX", "").strip()
-        try:
-            num_ctx = int(num_ctx_raw) if num_ctx_raw else 131072
-        except ValueError:
-            num_ctx = 131072
+        if num_ctx_raw:
+            try:
+                num_ctx = int(num_ctx_raw)
+            except ValueError:
+                num_ctx = 131072
+        else:
+            # Estimate input tokens: user_message chars / 4 (standard BPE
+            # heuristic) + 400 for the system prompt, then add output headroom.
+            estimated_input = len(user_message) // _CHARS_PER_TOKEN + 400
+            num_ctx = min(estimated_input + max_completion_tokens + 2000, 131072)
+            num_ctx = max(num_ctx, 8192)  # floor: never under-allocate badly
         keep_alive = os.environ.get("GRAPHIFY_OLLAMA_KEEP_ALIVE", "30m")
         kwargs["extra_body"] = {"options": {"num_ctx": num_ctx}, "keep_alive": keep_alive}
     resp = client.chat.completions.create(**kwargs)
@@ -321,9 +329,12 @@ def _call_openai_compat(
     output_tokens = result["output_tokens"]
     if output_tokens < 50 and backend == "ollama":
         print(
-            "[graphify] warning: ollama returned very few tokens — the model may be "
-            "too small or not following the JSON instruction format. "
-            "Try a larger model with --model (e.g. --model qwen2.5-coder:14b).",
+            "[graphify] warning: ollama returned very few tokens — likely causes: "
+            "(1) VRAM pressure: check `nvidia-smi` and reduce chunk size with "
+            "--token-budget (e.g. --token-budget 4096) or set "
+            "GRAPHIFY_OLLAMA_NUM_CTX to a smaller value; "
+            "(2) model too small for JSON instruction following — "
+            "try a larger model with --model (e.g. --model qwen2.5-coder:14b).",
             file=sys.stderr,
         )
     return result

tests/test_llm_backends.py

@@ -369,10 +369,38 @@ def test_ollama_extra_body_sets_num_ctx_and_keep_alive(monkeypatch):
 
     assert "extra_body" in captured, "extra_body must be sent to Ollama"
     eb = captured["extra_body"]
-    assert eb.get("options", {}).get("num_ctx") == 131072, "default num_ctx must be 131072"
+    # num_ctx is now dynamic: derived from message size, not hardcoded 131072
+    assert "num_ctx" in eb.get("options", {}), "num_ctx must be present"
+    assert eb["options"]["num_ctx"] >= 8192, "num_ctx must be at least the floor value"
     assert eb.get("keep_alive") == "30m", "default keep_alive must be 30m"
 
 
+def test_ollama_num_ctx_scales_with_small_token_budget(monkeypatch):
+    # Regression for #798 follow-up: with --token-budget 8192, the old hardcoded
+    # 131072 forced Ollama to allocate 128k KV-cache slots on a 31B model, causing
+    # VRAM exhaustion by chunk 4. num_ctx must now reflect actual chunk size.
+    captured = _install_capturing_openai(monkeypatch)
+    monkeypatch.delenv("GRAPHIFY_OLLAMA_NUM_CTX", raising=False)
+    monkeypatch.delenv("GRAPHIFY_OLLAMA_KEEP_ALIVE", raising=False)
+
+    # Simulate an 8k-token chunk: ~32k chars of content
+    small_chunk_msg = "x" * 32_000
+
+    llm._call_openai_compat(
+        "http://localhost:11434/v1", "ollama", "qwen2.5-coder:7b",
+        small_chunk_msg, temperature=0, max_completion_tokens=16384, backend="ollama",
+    )
+
+    num_ctx = captured["extra_body"]["options"]["num_ctx"]
+    # Should be far less than 131072 for an 8k input — VRAM-friendly
+    assert num_ctx < 131072, (
+        f"num_ctx={num_ctx} is too large for a small chunk; "
+        "this wastes VRAM and causes OOM on large models (#798)"
+    )
+    # But still large enough to fit input + output
+    assert num_ctx >= 8192, "num_ctx must cover at least the output cap"
+
+
 def test_ollama_num_ctx_env_override(monkeypatch):
     captured = _install_capturing_openai(monkeypatch)
     monkeypatch.setenv("GRAPHIFY_OLLAMA_NUM_CTX", "65536")