diff --git a/lightrag/llm/zhipu.py b/lightrag/llm/zhipu.py
index bceb15cc8b..9ce15dbd37 100644
--- a/lightrag/llm/zhipu.py
+++ b/lightrag/llm/zhipu.py
@@ -26,6 +26,7 @@
)
from lightrag.utils import (
+ TruncatedResponse,
wrap_embedding_func_with_attrs,
logger,
)
@@ -177,15 +178,19 @@ async def zhipu_complete_if_cache(
if not response.choices or response.choices[0].message is None:
return ""
- message = response.choices[0].message
+ choice = response.choices[0]
+ message = choice.message
content = message.content or ""
reasoning_content = getattr(message, "reasoning_content", "") or ""
if enable_cot and reasoning_content.strip():
if content:
- return f"{reasoning_content}{content}"
- return f"{reasoning_content}"
+ content = f"{reasoning_content}{content}"
+ else:
+ content = f"{reasoning_content}"
+ if getattr(choice, "finish_reason", None) == "length":
+ return TruncatedResponse(content)
return content
diff --git a/tests/llm/zhipu_impl/test_zhipu_llm.py b/tests/llm/zhipu_impl/test_zhipu_llm.py
index 6474a20c25..6c4e604630 100644
--- a/tests/llm/zhipu_impl/test_zhipu_llm.py
+++ b/tests/llm/zhipu_impl/test_zhipu_llm.py
@@ -7,17 +7,24 @@
import numpy as np
import pytest
+from lightrag.utils import is_truncated_response
+
def _fake_embedding_vector(dim=1024):
return [0.1] * dim
-def _fake_chat_response(content="", reasoning_content="", usage=None):
+def _fake_chat_response(
+ content="", reasoning_content="", usage=None, finish_reason="stop"
+):
message = SimpleNamespace(
content=content,
reasoning_content=reasoning_content,
)
- return SimpleNamespace(choices=[SimpleNamespace(message=message)], usage=usage)
+ return SimpleNamespace(
+ choices=[SimpleNamespace(message=message, finish_reason=finish_reason)],
+ usage=usage,
+ )
def _load_zhipu_module(monkeypatch, client_factory):
@@ -157,6 +164,48 @@ def add_usage(self, token_counts):
]
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_zhipu_length_finish_reason_marks_result_truncated(monkeypatch):
+ class FakeClient:
+ def __init__(self, api_key=None):
+ self.api_key = api_key
+ self.chat = SimpleNamespace(completions=SimpleNamespace(create=self.create))
+
+ def create(self, **kwargs):
+ return _fake_chat_response(content="partial answer", finish_reason="length")
+
+ zhipu_module = _load_zhipu_module(monkeypatch, FakeClient)
+
+ result = await zhipu_module.zhipu_complete_if_cache(
+ prompt="hello", api_key="test-key"
+ )
+
+ assert is_truncated_response(result)
+ assert result == "partial answer"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_zhipu_stop_finish_reason_keeps_plain_response(monkeypatch):
+ class FakeClient:
+ def __init__(self, api_key=None):
+ self.api_key = api_key
+ self.chat = SimpleNamespace(completions=SimpleNamespace(create=self.create))
+
+ def create(self, **kwargs):
+ return _fake_chat_response(content="complete answer", finish_reason="stop")
+
+ zhipu_module = _load_zhipu_module(monkeypatch, FakeClient)
+
+ result = await zhipu_module.zhipu_complete_if_cache(
+ prompt="hello", api_key="test-key"
+ )
+
+ assert not is_truncated_response(result)
+ assert result == "complete answer"
+
+
@pytest.mark.offline
@pytest.mark.asyncio
async def test_zhipu_complete_token_tracker_never_reaches_the_raw_client_call(