69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
from src.thinking_parser import parse_thinking_trace
|
|
|
|
|
|
def test_parse_xml_thinking_tag():
|
|
raw = "<thinking>\nLet me analyze this problem step by step.\n</thinking>\nHere is the answer."
|
|
segments, response = parse_thinking_trace(raw)
|
|
assert len(segments) == 1
|
|
assert segments[0].content == "Let me analyze this problem step by step."
|
|
assert segments[0].marker == "thinking"
|
|
assert response == "Here is the answer."
|
|
|
|
|
|
def test_parse_xml_thought_tag():
|
|
raw = "<thought>This is my reasoning process</thought>\nFinal response here."
|
|
segments, response = parse_thinking_trace(raw)
|
|
assert len(segments) == 1
|
|
assert segments[0].content == "This is my reasoning process"
|
|
assert segments[0].marker == "thought"
|
|
assert response == "Final response here."
|
|
|
|
|
|
def test_parse_text_thinking_prefix():
|
|
raw = "Thinking:\nThis is a text-based thinking trace.\n\nNow for the actual response."
|
|
segments, response = parse_thinking_trace(raw)
|
|
assert len(segments) == 1
|
|
assert segments[0].content == "This is a text-based thinking trace."
|
|
assert segments[0].marker == "Thinking:"
|
|
assert response == "Now for the actual response."
|
|
|
|
|
|
def test_parse_no_thinking():
|
|
raw = "This is a normal response without any thinking markers."
|
|
segments, response = parse_thinking_trace(raw)
|
|
assert len(segments) == 0
|
|
assert response == raw
|
|
|
|
|
|
def test_parse_empty_response():
|
|
segments, response = parse_thinking_trace("")
|
|
assert len(segments) == 0
|
|
assert response == ""
|
|
|
|
|
|
def test_parse_multiple_markers():
|
|
raw = "<thinking>First thinking</thinking>\n<thought>Second thought</thought>\nResponse"
|
|
segments, response = parse_thinking_trace(raw)
|
|
assert len(segments) == 2
|
|
assert segments[0].content == "First thinking"
|
|
assert segments[1].content == "Second thought"
|
|
|
|
|
|
def test_parse_thinking_with_empty_response():
|
|
raw = "<thinking>Just thinking, no response</thinking>"
|
|
segments, response = parse_thinking_trace(raw)
|
|
assert len(segments) == 1
|
|
assert segments[0].content == "Just thinking, no response"
|
|
assert response == ""
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_parse_xml_thinking_tag()
|
|
test_parse_xml_thought_tag()
|
|
test_parse_text_thinking_prefix()
|
|
test_parse_no_thinking()
|
|
test_parse_empty_response()
|
|
test_parse_multiple_markers()
|
|
test_parse_thinking_with_empty_response()
|
|
print("All thinking trace tests passed!")
|