diff --git a/graphify/llm.py b/graphify/llm.py index 31759eef2..40449f7d7 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -1483,8 +1483,11 @@ def _call_openai_compat( raise ValueError("LLM returned empty or filtered response") raw_content = resp.choices[0].message.content result = _parse_llm_json(raw_content or "{}") - result["input_tokens"] = resp.usage.prompt_tokens if resp.usage else 0 - result["output_tokens"] = resp.usage.completion_tokens if resp.usage else 0 + # Some OpenAI-compatible gateways report null token counts (#4325) — coerce + # to 0 so the ollama low-token warning and the per-chunk merge below don't + # crash comparing/adding None. + result["input_tokens"] = (resp.usage.prompt_tokens or 0) if resp.usage else 0 + result["output_tokens"] = (resp.usage.completion_tokens or 0) if resp.usage else 0 result["model"] = model # `finish_reason == "length"` means the model hit max_completion_tokens # mid-generation. The JSON we got back is truncated; callers should @@ -1531,8 +1534,8 @@ def _call_claude(api_key: str, model: str, user_message: str, max_tokens: int = ) raw_content = _anthropic_response_text(resp.content) result = _parse_llm_json(raw_content or "{}") - result["input_tokens"] = resp.usage.input_tokens if resp.usage else 0 - result["output_tokens"] = resp.usage.output_tokens if resp.usage else 0 + result["input_tokens"] = (resp.usage.input_tokens or 0) if resp.usage else 0 + result["output_tokens"] = (resp.usage.output_tokens or 0) if resp.usage else 0 result["model"] = model # Normalise Anthropic's `stop_reason` to the OpenAI-compat `finish_reason` # vocabulary so the adaptive-retry layer doesn't have to know which @@ -1885,8 +1888,8 @@ def _call_azure( raise ValueError("Azure OpenAI returned empty or filtered response") raw_content = resp.choices[0].message.content result = _parse_llm_json(raw_content or "{}") - result["input_tokens"] = resp.usage.prompt_tokens if resp.usage else 0 - result["output_tokens"] = resp.usage.completion_tokens if resp.usage else 0 + result["input_tokens"] = (resp.usage.prompt_tokens or 0) if resp.usage else 0 + result["output_tokens"] = (resp.usage.completion_tokens or 0) if resp.usage else 0 result["model"] = model result["finish_reason"] = resp.choices[0].finish_reason _mark_hollow(result, raw_content, "azure") @@ -1936,8 +1939,8 @@ def _call_bedrock(model: str, user_message: str, max_tokens: int = 8192, *, deep text = _bedrock_response_text(resp, default="{}") result = _parse_llm_json(text) usage = resp.get("usage", {}) - result["input_tokens"] = usage.get("inputTokens", 0) - result["output_tokens"] = usage.get("outputTokens", 0) + result["input_tokens"] = usage.get("inputTokens", 0) or 0 + result["output_tokens"] = usage.get("outputTokens", 0) or 0 result["model"] = model result["finish_reason"] = "length" if resp.get("stopReason") == "max_tokens" else "stop" _mark_hollow(result, text, "bedrock") diff --git a/tests/test_llm_null_token_counts.py b/tests/test_llm_null_token_counts.py new file mode 100644 index 000000000..841a5f400 --- /dev/null +++ b/tests/test_llm_null_token_counts.py @@ -0,0 +1,113 @@ +"""Regression tests for #4325: null token counts from LLM backends must not crash extraction. + +Some OpenAI-compatible gateways (and custom providers) return a `usage` object +whose token-count fields are `null`. Reading them verbatim left `None` in +`result["input_tokens"]` / `result["output_tokens"]`, which crashed downstream +(`'<' not supported between 'NoneType' and 'int'` in the ollama low-token +warning, and `None` propagating through the per-chunk merge). +""" + +import sys +import types +from types import SimpleNamespace + +from graphify import llm + +_NODE_JSON = '{"nodes": [{"id": "a", "label": "A"}], "edges": [], "hyperedges": []}' + + +def _fake_openai_null_usage(monkeypatch): + class _Completions: + def create(self, **kw): + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content=_NODE_JSON), + finish_reason="stop", + ) + ], + usage=SimpleNamespace(prompt_tokens=None, completion_tokens=None), + ) + + mod = types.ModuleType("openai") + client = SimpleNamespace(chat=SimpleNamespace(completions=_Completions())) + mod.OpenAI = lambda **kw: client + mod.AzureOpenAI = lambda **kw: client + monkeypatch.setitem(sys.modules, "openai", mod) + + +def _fake_anthropic_null_usage(monkeypatch): + class _Messages: + def create(self, **kw): + return SimpleNamespace( + content=[SimpleNamespace(text=_NODE_JSON, type="text")], + usage=SimpleNamespace(input_tokens=None, output_tokens=None), + stop_reason="end_turn", + ) + + mod = types.ModuleType("anthropic") + mod.Anthropic = lambda **kw: SimpleNamespace(messages=_Messages()) + monkeypatch.setitem(sys.modules, "anthropic", mod) + + +def _fake_boto3_null_usage(monkeypatch): + converse_resp = { + "output": {"message": {"content": [{"text": _NODE_JSON}]}}, + "usage": {"inputTokens": None, "outputTokens": None}, + "stopReason": "end_turn", + } + + class _Client: + def converse(self, **kw): + return converse_resp + + class _Session: + def __init__(self, *a, **k): + pass + + def client(self, *a, **k): + return _Client() + + mod = types.ModuleType("boto3") + mod.Session = _Session + botocore_mod = types.ModuleType("botocore") + config_mod = types.ModuleType("botocore.config") + config_mod.Config = lambda **kw: SimpleNamespace(**kw) + exceptions_mod = types.ModuleType("botocore.exceptions") + exceptions_mod.ClientError = type("ClientError", (Exception,), {}) + botocore_mod.config = config_mod + botocore_mod.exceptions = exceptions_mod + monkeypatch.setitem(sys.modules, "boto3", mod) + monkeypatch.setitem(sys.modules, "botocore", botocore_mod) + monkeypatch.setitem(sys.modules, "botocore.config", config_mod) + monkeypatch.setitem(sys.modules, "botocore.exceptions", exceptions_mod) + + +def test_openai_compat_null_token_counts_coerced_to_zero(monkeypatch): + _fake_openai_null_usage(monkeypatch) + result = llm._call_openai_compat( + "https://example.invalid/v1", "k", "m", "text", backend="openai" + ) + assert result["input_tokens"] == 0 + assert result["output_tokens"] == 0 + + +def test_azure_null_token_counts_coerced_to_zero(monkeypatch): + _fake_openai_null_usage(monkeypatch) + result = llm._call_azure("k", "https://example.invalid", "m", "text") + assert result["input_tokens"] == 0 + assert result["output_tokens"] == 0 + + +def test_claude_null_token_counts_coerced_to_zero(monkeypatch): + _fake_anthropic_null_usage(monkeypatch) + result = llm._call_claude("k", "claude-sonnet-4-6", "text") + assert result["input_tokens"] == 0 + assert result["output_tokens"] == 0 + + +def test_bedrock_null_token_counts_coerced_to_zero(monkeypatch): + _fake_boto3_null_usage(monkeypatch) + result = llm._call_bedrock("model", "text") + assert result["input_tokens"] == 0 + assert result["output_tokens"] == 0