WibblyWall/Paper_Gear_Cardboard
08
1"""2Test suite for DeepSeek-V4 Encoding.3 4Run: python test_encoding_dsv4.py5"""6 7import json8import os9 10from encoding_dsv4 import encode_messages, parse_message_from_completion_text11 12TESTS_DIR = os.path.join(os.path.dirname(__file__), "tests")13 14 15def test_case_1():16 """Thinking mode with tool calls (multi-turn, tool results merged into user)."""17 with open(os.path.join(TESTS_DIR, "test_input_1.json")) as f:18 td = json.load(f)19 messages = td["messages"]20 messages[0]["tools"] = td["tools"]21 gold = open(os.path.join(TESTS_DIR, "test_output_1.txt")).read()22 prompt = encode_messages(messages, thinking_mode="thinking")23 assert prompt == gold24 25 # Parse: assistant turn with tool call26 marker = "<|Assistant|><think>"27 first_start = prompt.find(marker) + len(marker)28 first_end = prompt.find("<|User|>", first_start)29 parsed_tc = parse_message_from_completion_text(prompt[first_start:first_end], thinking_mode="thinking")30 assert parsed_tc["reasoning_content"] == "The user wants to know the weather in Beijing. I should use the get_weather tool."31 assert parsed_tc["content"] == ""32 assert len(parsed_tc["tool_calls"]) == 133 assert parsed_tc["tool_calls"][0]["function"]["name"] == "get_weather"34 assert json.loads(parsed_tc["tool_calls"][0]["function"]["arguments"]) == {"location": "Beijing", "unit": "celsius"}35 36 # Parse: final assistant turn with content37 last_start = prompt.rfind(marker) + len(marker)38 parsed_final = parse_message_from_completion_text(prompt[last_start:], thinking_mode="thinking")39 assert parsed_final["reasoning_content"] == "Got the weather data. Let me format a nice response."40 assert "22°C" in parsed_final["content"]41 assert parsed_final["tool_calls"] == []42 43 print(" [PASS] case 1: thinking with tools (encode + parse)")44 45 46def test_case_2():47 """Thinking mode without tools (drop_thinking removes earlier reasoning)."""48 messages = json.load(open(os.path.join(TESTS_DIR, "test_input_2.json")))49 gold = open(os.path.join(TESTS_DIR, "test_output_2.txt")).read()50 prompt = encode_messages(messages, thinking_mode="thinking")51 assert prompt == gold52 53 # Parse: last assistant turn54 marker = "<|Assistant|><think>"55 last_start = prompt.rfind(marker) + len(marker)56 parsed = parse_message_from_completion_text(prompt[last_start:], thinking_mode="thinking")57 assert parsed["reasoning_content"] == "The user asks about the capital of France. It is Paris."58 assert parsed["content"] == "The capital of France is Paris."59 assert parsed["tool_calls"] == []60 61 # Verify drop_thinking: first assistant's reasoning should be absent62 assert "The user said hello" not in prompt63 64 print(" [PASS] case 2: thinking without tools (encode + parse)")65 66 67def test_case_3():68 """Interleaved thinking + search (developer with tools, latest_reminder)."""69 messages = json.load(open(os.path.join(TESTS_DIR, "test_input_3.json")))70 gold = open(os.path.join(TESTS_DIR, "test_output_3.txt")).read()71 assert encode_messages(messages, thinking_mode="thinking") == gold72 print(" [PASS] case 3: interleaved thinking + search")73 74 75def test_case_4():76 """Quick instruction task with latest_reminder (chat mode, action task)."""77 messages = json.load(open(os.path.join(TESTS_DIR, "test_input_4.json")))78 gold = open(os.path.join(TESTS_DIR, "test_output_4.txt")).read()79 assert encode_messages(messages, thinking_mode="chat") == gold80 print(" [PASS] case 4: quick instruction task")81 82 83if __name__ == "__main__":84 print("Running DeepSeek-V4 Encoding Tests...\n")85 test_case_1()86 test_case_2()87 test_case_3()88 test_case_4()89 print("\nAll 4 tests passed!")90 