Skip to content

Commit 2409ade

Browse files
authored
enh: The addition of the think closure tag supports multi-turn thin… (#267)
* ENH: The addition of the `think` closure tag supports multi-turn thinking scenarios. * fix lint * add test * fix lint * fix lint * fix lint * fix lint * add unittest * fix lint * fix lint
1 parent c7fca3d commit 2409ade

File tree

2 files changed

+123
-7
lines changed

2 files changed

+123
-7
lines changed

python/dify_plugin/interfaces/model/large_language_model.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -536,17 +536,22 @@ def _wrap_thinking_by_reasoning_content(self, delta: dict, is_reasoning: bool) -
536536

537537
content = delta.get("content") or ""
538538
reasoning_content = delta.get("reasoning_content")
539-
539+
output = content
540540
if reasoning_content:
541541
if not is_reasoning:
542-
content = "<think>\n" + reasoning_content
542+
output = "<think>\n" + reasoning_content
543543
is_reasoning = True
544544
else:
545-
content = reasoning_content
546-
elif is_reasoning and content:
547-
content = "\n</think>" + content
548-
is_reasoning = False
549-
return content, is_reasoning
545+
output = reasoning_content
546+
else:
547+
if is_reasoning:
548+
is_reasoning = False
549+
if not reasoning_content:
550+
output = "\n</think>"
551+
if content:
552+
output += content
553+
554+
return output, is_reasoning
550555

551556
############################################################
552557
# For executor use only #
Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
import unittest
2+
3+
from dify_plugin.entities.model import AIModelEntity, ModelPropertyKey, ModelType
4+
from dify_plugin.entities.model.llm import LLMMode, LLMResult
5+
from dify_plugin.interfaces.model.large_language_model import LargeLanguageModel
6+
7+
8+
class MockLLM(LargeLanguageModel):
9+
"""
10+
Concrete Mock class for testing non-abstract methods of LargeLanguageModel.
11+
"""
12+
13+
def _invoke(
14+
self,
15+
model: str,
16+
credentials: dict,
17+
prompt_messages: list,
18+
model_parameters: dict,
19+
tools: list,
20+
stop: list,
21+
stream: bool,
22+
user: str,
23+
) -> LLMResult:
24+
pass
25+
26+
def get_num_tokens(self, model: str, credentials: dict, prompt_messages: list, tools: list) -> int:
27+
return 0
28+
29+
def validate_credentials(self, model: str, credentials: dict) -> None:
30+
pass
31+
32+
@property
33+
def _invoke_error_mapping(self) -> dict:
34+
return {}
35+
36+
37+
class TestWrapThinking(unittest.TestCase):
38+
def setUp(self):
39+
# Create a dummy model schema to satisfy AIModel.__init__
40+
dummy_schema = AIModelEntity(
41+
model="mock_model",
42+
label={"en_US": "Mock Model"},
43+
model_type=ModelType.LLM,
44+
features=[],
45+
model_properties={ModelPropertyKey.MODE: LLMMode.CHAT.value, ModelPropertyKey.CONTEXT_SIZE: 4096},
46+
parameter_rules=[],
47+
pricing=None,
48+
deprecated=False,
49+
)
50+
self.llm = MockLLM(model_schemas=[dummy_schema])
51+
52+
def test_wrap_thinking_logic_closure(self):
53+
"""
54+
Test that when reasoning_content ends, even if content is empty (e.g. followed immediately by tool_calls),
55+
the <think> tag should be closed correctly.
56+
"""
57+
58+
# Simulate simulated streaming data:
59+
# 1. Has reasoning_content
60+
# 2. reasoning_content ends, followed immediately by tool_calls (content is None)
61+
62+
chunks = [
63+
# Chunk 1: Thinking started
64+
{"reasoning_content": "Thinking started.", "content": ""},
65+
# Chunk 2: Still thinking
66+
{"reasoning_content": " Still thinking.", "content": ""},
67+
# Chunk 3: Thinking ended, transitioned to Tool Call (reasoning_content=None, content=None/Empty)
68+
# This is a critical point, old logic would fail here because content is empty
69+
{"reasoning_content": None, "content": "", "tool_calls": [{"id": "call_1", "function": {}}]},
70+
# Chunk 4: Subsequent tool parameter stream
71+
{"reasoning_content": None, "content": "", "tool_calls": [{"function": {"arguments": "{"}}]},
72+
]
73+
74+
# Use the "new logic" from PR for testing.
75+
# We can directly call self.llm._wrap_thinking_by_reasoning_content.
76+
77+
# Assume we are testing the logic function itself:
78+
is_reasoning = False
79+
full_output = ""
80+
81+
for chunk in chunks:
82+
# Directly call the implementation in SDK to verify real code logic
83+
output, is_reasoning = self.llm._wrap_thinking_by_reasoning_content(chunk, is_reasoning)
84+
full_output += output
85+
86+
# Verify results
87+
print(f"DEBUG Output: {full_output!r}")
88+
89+
assert "<think>" in full_output
90+
assert "Thinking started. Still thinking." in full_output
91+
assert "</think>" in full_output, "Should verify <think> tag is closed properly"
92+
93+
# Verify the position of the closing tag: should be after the thinking content
94+
expected_part = "Thinking started. Still thinking.\n</think>"
95+
assert expected_part in full_output
96+
97+
def test_standard_reasoning_flow(self):
98+
"""Test standard reasoning -> text flow"""
99+
chunks = [
100+
{"reasoning_content": "Thinking.", "content": ""},
101+
{"reasoning_content": None, "content": "Hello world."},
102+
]
103+
104+
is_reasoning = False
105+
full_output = ""
106+
for chunk in chunks:
107+
# Directly call the implementation in SDK
108+
output, is_reasoning = self.llm._wrap_thinking_by_reasoning_content(chunk, is_reasoning)
109+
full_output += output
110+
111+
assert full_output == "<think>\nThinking.\n</think>Hello world."

0 commit comments

Comments
 (0)