Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions graphify/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,14 @@ def _call_openai_compat(
}
if temperature is not None:
kwargs["temperature"] = temperature
# kimi-k2.6 is a reasoning model: it writes its chain-of-thought to
# `reasoning_content` and the answer to `content`. When reasoning consumes
# the entire max_completion_tokens budget on large chunks, `content` stays
# empty and we parse zero nodes while still paying for the input + reasoning
# tokens. For structured extraction we don't need chain-of-thought, so
# disable thinking on Moonshot-hosted models.
if "moonshot" in base_url:
kwargs["extra_body"] = {"thinking": {"type": "disabled"}}
resp = client.chat.completions.create(**kwargs)
result = _parse_llm_json(resp.choices[0].message.content or "{}")
result["input_tokens"] = resp.usage.prompt_tokens if resp.usage else 0
Expand Down
57 changes: 57 additions & 0 deletions tests/test_llm.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
"""Tests for the direct LLM backend (Claude / Kimi)."""
from unittest.mock import MagicMock, patch


def _mock_response(content: str = '{"nodes":[],"edges":[]}', completion_tokens: int = 10):
"""Build a fake OpenAI ChatCompletion response."""
resp = MagicMock()
resp.choices = [MagicMock(message=MagicMock(content=content))]
resp.usage = MagicMock(prompt_tokens=100, completion_tokens=completion_tokens)
return resp


def test_kimi_call_disables_thinking():
"""kimi-k2.6 is a reasoning model: when thinking is enabled and the chunk
is large, all of max_completion_tokens is consumed by reasoning_content,
leaving content empty and graphify parsing zero nodes. Disable thinking on
Moonshot endpoints so content is always populated for structured extraction.
"""
from graphify.llm import _call_openai_compat

fake_client = MagicMock()
fake_client.chat.completions.create.return_value = _mock_response()

with patch("openai.OpenAI", return_value=fake_client):
_call_openai_compat(
base_url="https://api.moonshot.ai/v1",
api_key="sk-test",
model="kimi-k2.6",
user_message="=== example.py ===\ndef hello(): pass\n",
temperature=None,
)

kwargs = fake_client.chat.completions.create.call_args.kwargs
assert "extra_body" in kwargs, "moonshot calls must pass extra_body to disable thinking"
assert kwargs["extra_body"] == {"thinking": {"type": "disabled"}}


def test_non_moonshot_call_does_not_disable_thinking():
"""The thinking-disabled flag is Moonshot-specific extra_body. Other
OpenAI-compatible providers (real OpenAI, Together, Groq, etc.) don't
accept it and would 400. Only set it for Moonshot URLs."""
from graphify.llm import _call_openai_compat

fake_client = MagicMock()
fake_client.chat.completions.create.return_value = _mock_response()

with patch("openai.OpenAI", return_value=fake_client):
_call_openai_compat(
base_url="https://api.openai.com/v1",
api_key="sk-test",
model="gpt-4o",
user_message="hi",
temperature=0,
)

kwargs = fake_client.chat.completions.create.call_args.kwargs
assert "extra_body" not in kwargs, "non-Moonshot calls must not pass thinking flag"