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(