diff --git a/test/registered/unit/parser/test_code_completion_parser.py b/test/registered/unit/parser/test_code_completion_parser.py new file mode 100644 index 000000000..50ca71519 --- /dev/null +++ b/test/registered/unit/parser/test_code_completion_parser.py @@ -0,0 +1,188 @@ +"""Unit tests for srt/parser/code_completion_parser.py""" + +import unittest +from unittest.mock import patch + +from sglang.srt.entrypoints.openai.protocol import CompletionRequest +from sglang.srt.parser.code_completion_parser import ( + CompletionTemplate, + FimPosition, + completion_template_exists, + completion_templates, + generate_completion_prompt, + generate_completion_prompt_from_request, + is_completion_template_defined, + register_completion_template, + set_completion_template, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="stage-a-cpu-only") + + +class TestFimPosition(CustomTestCase): + def test_middle_and_end_are_distinct(self): + """Test that MIDDLE and END are different enum values.""" + self.assertNotEqual(FimPosition.MIDDLE, FimPosition.END) + + +class TestCompletionTemplate(CustomTestCase): + def test_dataclass_fields(self): + """Test creating a CompletionTemplate with all fields.""" + t = CompletionTemplate( + name="test", + fim_begin_token="", + fim_middle_token="", + fim_end_token="", + fim_position=FimPosition.MIDDLE, + ) + self.assertEqual(t.name, "test") + self.assertEqual(t.fim_begin_token, "") + self.assertEqual(t.fim_position, FimPosition.MIDDLE) + + +class TestRegisterCompletionTemplate(CustomTestCase): + def test_builtin_templates_registered(self): + """Test that deepseek_coder, star_coder, qwen_coder are pre-registered.""" + self.assertTrue(completion_template_exists("deepseek_coder")) + self.assertTrue(completion_template_exists("star_coder")) + self.assertTrue(completion_template_exists("qwen_coder")) + + def test_unregistered_template_not_found(self): + """Test that a non-existent template returns False.""" + self.assertFalse(completion_template_exists("nonexistent_template")) + + def test_register_new_template(self): + """Test registering a new template.""" + t = CompletionTemplate( + name="_test_new_template", + fim_begin_token="", + fim_middle_token="", + fim_end_token="", + fim_position=FimPosition.END, + ) + register_completion_template(t) + self.assertTrue(completion_template_exists("_test_new_template")) + # Cleanup + del completion_templates["_test_new_template"] + + def test_register_duplicate_raises(self): + """Test that registering a duplicate name without override raises.""" + with self.assertRaises(AssertionError): + register_completion_template( + CompletionTemplate( + name="deepseek_coder", + fim_begin_token="x", + fim_middle_token="y", + fim_end_token="z", + fim_position=FimPosition.MIDDLE, + ) + ) + + def test_register_duplicate_with_override(self): + """Test that override=True allows re-registration.""" + original = completion_templates["deepseek_coder"] + try: + register_completion_template( + CompletionTemplate( + name="deepseek_coder", + fim_begin_token="", + fim_middle_token="", + fim_end_token="", + fim_position=FimPosition.END, + ), + override=True, + ) + self.assertEqual( + completion_templates["deepseek_coder"].fim_begin_token, "" + ) + finally: + # Restore original + completion_templates["deepseek_coder"] = original + + +class TestGenerateCompletionPrompt(CustomTestCase): + def test_deepseek_coder_middle_position(self): + """Test FIM prompt with MIDDLE position (deepseek_coder style).""" + result = generate_completion_prompt( + "prefix_code", "suffix_code", "deepseek_coder" + ) + t = completion_templates["deepseek_coder"] + expected = f"{t.fim_begin_token}prefix_code{t.fim_middle_token}suffix_code{t.fim_end_token}" + self.assertEqual(result, expected) + + def test_star_coder_end_position(self): + """Test FIM prompt with END position (star_coder style).""" + result = generate_completion_prompt("prefix_code", "suffix_code", "star_coder") + t = completion_templates["star_coder"] + expected = f"{t.fim_begin_token}prefix_code{t.fim_end_token}suffix_code{t.fim_middle_token}" + self.assertEqual(result, expected) + + def test_qwen_coder_end_position(self): + """Test FIM prompt with END position (qwen_coder style).""" + result = generate_completion_prompt("prefix", "suffix", "qwen_coder") + t = completion_templates["qwen_coder"] + expected = ( + f"{t.fim_begin_token}prefix{t.fim_end_token}suffix{t.fim_middle_token}" + ) + self.assertEqual(result, expected) + + def test_empty_prompt_and_suffix(self): + """Test FIM prompt generation with empty strings.""" + result = generate_completion_prompt("", "", "deepseek_coder") + t = completion_templates["deepseek_coder"] + expected = f"{t.fim_begin_token}{t.fim_middle_token}{t.fim_end_token}" + self.assertEqual(result, expected) + + +class TestGenerateCompletionPromptFromRequest(CustomTestCase): + def test_empty_suffix_returns_prompt_directly(self): + """Test that empty suffix bypasses FIM formatting.""" + request = CompletionRequest(prompt="just code", suffix="") + result = generate_completion_prompt_from_request(request) + self.assertEqual(result, "just code") + + def test_nonempty_suffix_uses_fim_template(self): + """Test that non-empty suffix triggers FIM formatting.""" + with patch( + "sglang.srt.parser.code_completion_parser.completion_template_name", + "deepseek_coder", + ): + request = CompletionRequest(prompt="prefix", suffix="suffix") + result = generate_completion_prompt_from_request(request) + t = completion_templates["deepseek_coder"] + expected = ( + f"{t.fim_begin_token}prefix{t.fim_middle_token}suffix{t.fim_end_token}" + ) + self.assertEqual(result, expected) + + +class TestSetCompletionTemplate(CustomTestCase): + def test_set_only_once(self): + """Test that set_completion_template only sets the name once.""" + import sglang.srt.parser.code_completion_parser as module + + with patch.object(module, "completion_template_name", None): + set_completion_template("star_coder") + self.assertEqual(module.completion_template_name, "star_coder") + # Second call should be ignored + set_completion_template("qwen_coder") + self.assertEqual(module.completion_template_name, "star_coder") + + def test_is_completion_template_defined(self): + """Test the defined check before and after setting.""" + import sglang.srt.parser.code_completion_parser as module + + old_name = module.completion_template_name + try: + module.completion_template_name = None + self.assertFalse(is_completion_template_defined()) + set_completion_template("deepseek_coder") + self.assertTrue(is_completion_template_defined()) + finally: + module.completion_template_name = old_name + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/parser/test_conversation.py b/test/registered/unit/parser/test_conversation.py new file mode 100644 index 000000000..5c3f4818e --- /dev/null +++ b/test/registered/unit/parser/test_conversation.py @@ -0,0 +1,1300 @@ +"""Unit tests for srt/parser/conversation.py""" + +import json +import os +import tempfile +import unittest + +from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionMessageContentAudioPart, + ChatCompletionMessageContentAudioURL, + ChatCompletionMessageContentImagePart, + ChatCompletionMessageContentImageURL, + ChatCompletionMessageContentTextPart, + ChatCompletionMessageContentVideoPart, + ChatCompletionMessageContentVideoURL, + ChatCompletionMessageGenericParam, + ChatCompletionMessageUserParam, + ChatCompletionRequest, +) +from sglang.srt.parser.conversation import ( + Conversation, + SeparatorStyle, + _get_full_multimodal_text_prompt, + chat_template_exists, + chat_templates, + generate_chat_conv, + generate_embedding_convs, + get_conv_template_by_model_path, + get_model_type, + register_conv_template, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="stage-a-cpu-only") + + +class TestConversationGetPrompt(CustomTestCase): + def test_add_colon_single(self): + """Test prompt generation with ADD_COLON_SINGLE style.""" + conv = Conversation( + name="test", + system_message="System msg", + roles=("User", "Assistant"), + messages=[["User", "Hello"], ["Assistant", "Hi"], ["User", None]], + sep_style=SeparatorStyle.ADD_COLON_SINGLE, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("System msg\n", prompt) + self.assertIn("User: Hello\n", prompt) + self.assertIn("Assistant: Hi\n", prompt) + self.assertTrue(prompt.endswith("User:")) + + def test_add_colon_two(self): + """Test prompt generation with ADD_COLON_TWO style (alternating separators).""" + conv = Conversation( + name="test", + system_message="Sys", + roles=("User", "Assistant"), + messages=[["User", "Q"], ["Assistant", "A"], ["User", None]], + sep_style=SeparatorStyle.ADD_COLON_TWO, + sep="", + sep2="", + ) + prompt = conv.get_prompt() + self.assertIn("User: Q", prompt) + self.assertIn("Assistant: A", prompt) + self.assertTrue(prompt.endswith("User:")) + + def test_chatml(self): + """Test prompt generation with CHATML style.""" + conv = Conversation( + name="test", + system_message="<|im_start|>system\nYou are helpful", + roles=("<|im_start|>user", "<|im_start|>assistant"), + messages=[ + ["<|im_start|>user", "Hello"], + ["<|im_start|>assistant", None], + ], + sep_style=SeparatorStyle.CHATML, + sep="<|im_end|>", + ) + prompt = conv.get_prompt() + self.assertIn("You are helpful<|im_end|>", prompt) + self.assertIn("<|im_start|>user\nHello<|im_end|>", prompt) + self.assertTrue(prompt.endswith("<|im_start|>assistant\n")) + + def test_llama3(self): + """Test prompt generation with LLAMA3 style.""" + conv = Conversation( + name="test", + system_message="<|start_header_id|>system<|end_header_id|>\n\nBe helpful<|eot_id|>", + roles=("user", "assistant"), + messages=[["user", "Hi"], ["assistant", None]], + sep_style=SeparatorStyle.LLAMA3, + ) + prompt = conv.get_prompt() + self.assertIn("Be helpful<|eot_id|>", prompt) + self.assertIn( + "<|start_header_id|>user<|end_header_id|>\n\nHi<|eot_id|>", prompt + ) + self.assertTrue( + prompt.endswith("<|start_header_id|>assistant<|end_header_id|>\n\n") + ) + + def test_no_colon_single(self): + """Test prompt generation with NO_COLON_SINGLE style.""" + conv = Conversation( + name="test", + system_message="", + roles=("[USER]", "[ASST]"), + messages=[["[USER]", "Hello"], ["[ASST]", None]], + sep_style=SeparatorStyle.NO_COLON_SINGLE, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("[USER]Hello\n", prompt) + self.assertTrue(prompt.endswith("[ASST]")) + + def test_none_message_in_prompt(self): + """Test that None message produces role-only output (no content).""" + conv = Conversation( + name="test", + system_message="", + roles=("User", "Assistant"), + messages=[["User", "Q"], ["Assistant", None]], + sep_style=SeparatorStyle.ADD_COLON_SINGLE, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertTrue(prompt.endswith("Assistant:")) + + def test_empty_system_message(self): + """Test that empty system message produces empty prefix for LLAMA3.""" + conv = Conversation( + name="test", + system_message="", + roles=("User", "Assistant"), + messages=[["User", "Hello"], ["Assistant", None]], + sep_style=SeparatorStyle.LLAMA3, + ) + prompt = conv.get_prompt() + self.assertNotIn("system", prompt.lower()) + + def test_add_colon_space_single(self): + """Test prompt generation with ADD_COLON_SPACE_SINGLE style.""" + conv = Conversation( + name="test", + system_message="Sys", + roles=("User", "Bot"), + messages=[["User", "Hi"], ["Bot", None]], + sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("User: Hi\n", prompt) + # None message should end with ": " (space after colon) + self.assertTrue(prompt.endswith("Bot: ")) + + def test_add_new_line_single(self): + """Test prompt generation with ADD_NEW_LINE_SINGLE style.""" + conv = Conversation( + name="test", + system_message="Sys", + roles=("User", "Bot"), + messages=[["User", "Hi"], ["Bot", None]], + sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("User\nHi\n", prompt) + self.assertTrue(prompt.endswith("Bot\n")) + + def test_no_colon_two(self): + """Test prompt generation with NO_COLON_TWO style (alternating separators).""" + conv = Conversation( + name="test", + system_message="", + roles=("[U]", "[A]"), + messages=[["[U]", "Q"], ["[A]", "A"], ["[U]", None]], + sep_style=SeparatorStyle.NO_COLON_TWO, + sep="", + sep2="", + ) + prompt = conv.get_prompt() + self.assertIn("[U]Q", prompt) + self.assertIn("[A]A", prompt) + self.assertTrue(prompt.endswith("[U]")) + + def test_llama2_with_system(self): + """Test LLAMA2 with system message.""" + conv = Conversation( + name="test", + system_message="<>\nBe helpful\n<>\n\n", + system_template="[INST] {system_message}", + roles=("[INST]", "[/INST]"), + messages=[["[INST]", "Hi"], ["[/INST]", None]], + sep_style=SeparatorStyle.LLAMA2, + sep=" ", + sep2=" ", + ) + prompt = conv.get_prompt() + self.assertIn("Be helpful", prompt) + self.assertIn("Hi ", prompt) + + def test_llama2_without_system(self): + """Test LLAMA2 without system message falls back to '[INST] ' prefix.""" + conv = Conversation( + name="test", + system_message="", + roles=("[INST]", "[/INST]"), + messages=[["[INST]", "Hi"], ["[/INST]", None]], + sep_style=SeparatorStyle.LLAMA2, + sep=" ", + sep2=" ", + ) + prompt = conv.get_prompt() + self.assertTrue(prompt.startswith("[INST] Hi")) + + def test_llama2_multi_turn(self): + """Test LLAMA2 with multi-turn (i>0 uses tag+sep pattern).""" + conv = Conversation( + name="test", + system_message="<>\nSys\n<>\n\n", + system_template="[INST] {system_message}", + roles=("[INST]", "[/INST]"), + messages=[ + ["[INST]", "Q1"], + ["[/INST]", "A1"], + ["[INST]", "Q2"], + ["[/INST]", None], + ], + sep_style=SeparatorStyle.LLAMA2, + sep=" ", + sep2=" ", + ) + prompt = conv.get_prompt() + # i=0: message + " " (no tag prefix) + self.assertIn("Q1 ", prompt) + # i=1: tag + " " + message + sep2 + self.assertIn("[/INST] A1 ", prompt) + + def test_llama4(self): + """Test prompt generation with LLAMA4 style.""" + conv = Conversation( + name="test", + system_message="Be helpful", + system_template="{system_message}", + roles=("user", "assistant"), + messages=[["user", "Hello"], ["assistant", None]], + sep_style=SeparatorStyle.LLAMA4, + ) + prompt = conv.get_prompt() + self.assertIn("Be helpful", prompt) + self.assertIn("<|header_start|>user<|header_end|>", prompt) + self.assertIn("Hello<|eot|>", prompt) + + def test_llama4_empty_system(self): + """Test LLAMA4 with empty system message omits system prefix.""" + conv = Conversation( + name="test", + system_message="", + roles=("user", "assistant"), + messages=[["user", "Hello"], ["assistant", None]], + sep_style=SeparatorStyle.LLAMA4, + ) + prompt = conv.get_prompt() + self.assertTrue(prompt.startswith("<|header_start|>user")) + + def test_chatglm3(self): + """Test prompt generation with CHATGLM3 style.""" + conv = Conversation( + name="test", + system_message="<|system|>\nBe helpful", + roles=("<|user|>", "<|assistant|>"), + messages=[["<|user|>", "Hi"], ["<|assistant|>", None]], + sep_style=SeparatorStyle.CHATGLM3, + ) + prompt = conv.get_prompt() + self.assertIn("Be helpful", prompt) + self.assertIn("<|user|>\nHi", prompt) + self.assertTrue(prompt.endswith("<|assistant|>")) + + def test_deepseek_chat(self): + """Test prompt generation with DEEPSEEK_CHAT style.""" + conv = Conversation( + name="test", + system_message="", + roles=("User", "Assistant"), + messages=[["User", "Q"], ["Assistant", "A"], ["User", None]], + sep_style=SeparatorStyle.DEEPSEEK_CHAT, + sep="\n\n", + sep2="", + ) + prompt = conv.get_prompt() + self.assertIn("User: Q\n\n", prompt) + self.assertIn("Assistant: A", prompt) + self.assertTrue(prompt.endswith("User:")) + + def test_robin(self): + """Test prompt generation with ROBIN style.""" + conv = Conversation( + name="test", + system_message="Sys", + roles=("###Human", "###Assistant"), + messages=[["###Human", "Hi"], ["###Assistant", None]], + sep_style=SeparatorStyle.ROBIN, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("###Human:\nHi\n", prompt) + self.assertTrue(prompt.endswith("###Assistant:\n")) + + def test_falcon_chat(self): + """Test prompt generation with FALCON_CHAT style.""" + conv = Conversation( + name="test", + system_message="System prompt.", + roles=("User", "Falcon"), + messages=[["User", "Hi"], ["Falcon", None]], + sep_style=SeparatorStyle.FALCON_CHAT, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("System prompt.\n", prompt) + self.assertIn("User: Hi\n", prompt) + self.assertTrue(prompt.endswith("Falcon:")) + + def test_metamath(self): + """Test prompt generation with METAMATH style.""" + conv = Conversation( + name="test", + system_message="", + roles=("Query", "Response"), + messages=[["Query", "2+2?"], ["Response", None]], + sep_style=SeparatorStyle.METAMATH, + sep="\n", + sep2="Let's think step by step.\n", + ) + prompt = conv.get_prompt() + self.assertIn("Query:\n2+2?\n", prompt) + self.assertIn("Response: Let's think step by step.\n", prompt) + + def test_mpt(self): + """Test prompt generation with MPT style.""" + conv = Conversation( + name="test", + system_message="<|system|>", + roles=("<|user|>", "<|assistant|>"), + messages=[["<|user|>", "Hi"], ["<|assistant|>", None]], + sep_style=SeparatorStyle.MPT, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("<|user|>Hi\n", prompt) + self.assertTrue(prompt.endswith("<|assistant|>")) + + def test_chatintern(self): + """Test prompt generation with CHATINTERN style.""" + conv = Conversation( + name="test", + system_message="", + roles=("HUMAN", "BOT"), + messages=[["HUMAN", "Hi"], ["BOT", "Hello"], ["HUMAN", None]], + sep_style=SeparatorStyle.CHATINTERN, + sep="\n", + sep2="", + ) + prompt = conv.get_prompt() + self.assertIn("HUMAN:Hi\n", prompt) + self.assertIn("BOT:Hello", prompt) + + def test_dolly(self): + """Test prompt generation with DOLLY style.""" + conv = Conversation( + name="test", + system_message="", + roles=("Instruction", "Response"), + messages=[["Instruction", "Q"], ["Response", "A"], ["Instruction", None]], + sep_style=SeparatorStyle.DOLLY, + sep="\n\n", + sep2="", + ) + prompt = conv.get_prompt() + self.assertIn("Instruction:\nQ\n\n", prompt) + self.assertIn("Response:\nA", prompt) + self.assertTrue(prompt.endswith("Instruction:\n")) + + def test_phoenix(self): + """Test prompt generation with PHOENIX style.""" + conv = Conversation( + name="test", + system_message="", + roles=("Human", "Phoenix"), + messages=[["Human", "Hi"], ["Phoenix", None]], + sep_style=SeparatorStyle.PHOENIX, + ) + prompt = conv.get_prompt() + self.assertIn("Human: Hi", prompt) + self.assertTrue(prompt.endswith("Phoenix: ")) + + def test_deepseek_vl2(self): + """Test prompt generation with DeepSeekVL2 style.""" + conv = Conversation( + name="test", + system_message="Sys", + roles=("User", "Assistant"), + messages=[["User", "Q"], ["Assistant", None]], + sep_style=SeparatorStyle.DeepSeekVL2, + sep="\n", + sep2="", + ) + prompt = conv.get_prompt() + self.assertIn("Sys\n", prompt) + self.assertIn("User: Q\n", prompt) + self.assertTrue(prompt.endswith("Assistant:")) + + def test_deepseek_vl2_empty_system(self): + """Test DeepSeekVL2 with empty system message omits system prefix.""" + conv = Conversation( + name="test", + system_message="", + roles=("User", "Assistant"), + messages=[["User", "Q"], ["Assistant", None]], + sep_style=SeparatorStyle.DeepSeekVL2, + sep="\n", + sep2="", + ) + prompt = conv.get_prompt() + self.assertTrue(prompt.startswith("User: Q")) + + def test_gemma3(self): + """Test prompt generation with GEMMA3 style (first message special).""" + conv = Conversation( + name="test", + system_message="", + roles=("", ""), + messages=[["", "Hello"], ["", "Hi"], ["", None]], + sep_style=SeparatorStyle.GEMMA3, + sep="", + ) + prompt = conv.get_prompt() + # First message: no role prefix, just message + sep + self.assertTrue(prompt.startswith("Hello")) + # Subsequent: role + message + sep + self.assertIn("Hi", prompt) + + def test_rwkv(self): + """Test prompt generation with RWKV style (newline replacement).""" + conv = Conversation( + name="test", + system_message="", + roles=("Bob", "Alice"), + messages=[["Bob", "Hello\n\nWorld"], ["Alice", None]], + sep_style=SeparatorStyle.RWKV, + ) + prompt = conv.get_prompt() + # RWKV replaces \n\n with \n in message + self.assertIn("Bob: Hello\nWorld\n\n", prompt) + + def test_qwen2_vl_embed(self): + """Test prompt generation with QWEN2_VL_EMBED style.""" + conv = Conversation( + name="test", + system_message="Sys", + roles=("user", "assistant"), + messages=[["user", "Hi"], ["assistant", None]], + sep_style=SeparatorStyle.QWEN2_VL_EMBED, + sep="\n", + stop_str="<|endoftext|>", + ) + prompt = conv.get_prompt() + self.assertIn("user\nHi\n", prompt) + self.assertTrue(prompt.endswith("<|endoftext|>")) + + def test_chatglm(self): + """Test prompt generation with CHATGLM style (round numbering).""" + conv = Conversation( + name="chatglm", + system_message="", + roles=("问", "答"), + messages=[["问", "Hello"], ["答", "Hi"], ["问", None]], + sep_style=SeparatorStyle.CHATGLM, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("[Round 0]\n", prompt) + self.assertIn("问:Hello\n", prompt) + self.assertIn("答:Hi\n", prompt) + self.assertTrue(prompt.endswith("问:")) + + def test_chatglm2_round_offset(self): + """Test CHATGLM style with chatglm2 name (round starts at 1 instead of 0).""" + conv = Conversation( + name="chatglm2", + system_message="", + roles=("问", "答"), + messages=[["问", "Hello"], ["答", None]], + sep_style=SeparatorStyle.CHATGLM, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("[Round 1]\n", prompt) + + def test_chatglm_with_system(self): + """Test CHATGLM with non-empty system message.""" + conv = Conversation( + name="chatglm", + system_message="You are helpful", + roles=("问", "答"), + messages=[["问", "Hi"], ["答", None]], + sep_style=SeparatorStyle.CHATGLM, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertTrue(prompt.startswith("You are helpful\n")) + + def test_qwen2_audio(self): + """Test QWEN2_AUDIO style with audio token counter replacement.""" + conv = Conversation( + name="test", + system_message="", + roles=("user", "assistant"), + messages=[ + ["user", "Listen: and "], + ["assistant", None], + ], + sep_style=SeparatorStyle.QWEN2_AUDIO, + sep="\n", + audio_token="", + ) + prompt = conv.get_prompt() + # Audio tokens should be replaced with counter: idx=1, idx=2 + self.assertIn("", prompt) + self.assertIn("", prompt) + self.assertNotIn("{idx}", prompt) + + def test_paddle_ocr(self): + """Test prompt generation with PADDLE_OCR style.""" + conv = Conversation( + name="test", + system_message="", + roles=("USER", "ASSISTANT"), + messages=[["USER", "Describe image"], ["ASSISTANT", None]], + sep_style=SeparatorStyle.PADDLE_OCR, + sep="", + ) + prompt = conv.get_prompt() + self.assertIn("USER: Describe image", prompt) + self.assertTrue(prompt.endswith("ASSISTANT: ")) + + def test_paddle_ocr_with_image_token(self): + """Test PADDLE_OCR strips newline after image token for USER role.""" + conv = Conversation( + name="test", + system_message="", + roles=("USER", "ASSISTANT"), + messages=[ + ["USER", "\nDescribe this"], + ["ASSISTANT", "It shows a cat"], + ], + sep_style=SeparatorStyle.PADDLE_OCR, + sep="", + image_token="", + ) + prompt = conv.get_prompt() + # image_token + "\n" should be replaced with just image_token + self.assertIn("USER: Describe this\n", prompt) + self.assertIn("ASSISTANT: It shows a cat", prompt) + + def test_mpt_with_tuple_message(self): + """Test MPT style extracts first element from tuple messages.""" + conv = Conversation( + name="test", + system_message="<|system|>", + roles=("<|user|>", "<|assistant|>"), + messages=[ + ["<|user|>", ("Hello", "extra1", "extra2")], + ["<|assistant|>", None], + ], + sep_style=SeparatorStyle.MPT, + sep="\n", + ) + prompt = conv.get_prompt() + self.assertIn("<|user|>Hello\n", prompt) + self.assertNotIn("extra1", prompt) + + def test_invalid_sep_style_raises(self): + """Test that an invalid SeparatorStyle raises ValueError.""" + conv = Conversation( + name="test", + system_message="", + roles=("A", "B"), + messages=[["A", "Hi"]], + sep_style=999, + sep="\n", + ) + with self.assertRaises(ValueError): + conv.get_prompt() + + +class TestConversationMethods(CustomTestCase): + def _make_conv(self): + return Conversation( + name="test", + roles=("User", "Assistant"), + messages=[], + sep_style=SeparatorStyle.ADD_COLON_SINGLE, + sep="\n", + ) + + def test_append_message(self): + """Test appending messages to conversation.""" + conv = self._make_conv() + conv.append_message("User", "Hello") + conv.append_message("Assistant", "Hi") + self.assertEqual(len(conv.messages), 2) + self.assertEqual(conv.messages[0], ["User", "Hello"]) + + def test_set_system_message(self): + """Test setting the system message.""" + conv = self._make_conv() + conv.set_system_message("Be helpful") + self.assertEqual(conv.system_message, "Be helpful") + + def test_update_last_message(self): + """Test updating the last message in-place.""" + conv = self._make_conv() + conv.append_message("User", "Q") + conv.append_message("Assistant", None) + conv.update_last_message("Answer") + self.assertEqual(conv.messages[-1][1], "Answer") + + def test_to_openai_api_messages_with_system(self): + """Test conversion to OpenAI format with system message.""" + conv = self._make_conv() + conv.system_message = "Be helpful" + conv.append_message("User", "Hello") + conv.append_message("Assistant", "Hi") + result = conv.to_openai_api_messages() + self.assertEqual(result[0], {"role": "system", "content": "Be helpful"}) + self.assertEqual(result[1], {"role": "user", "content": "Hello"}) + self.assertEqual(result[2], {"role": "assistant", "content": "Hi"}) + + def test_to_openai_api_messages_without_system(self): + """Test conversion to OpenAI format without system message.""" + conv = self._make_conv() + conv.append_message("User", "Hello") + result = conv.to_openai_api_messages() + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["role"], "user") + + def test_to_openai_api_messages_skips_none_assistant(self): + """Test that None assistant message is omitted from OpenAI format.""" + conv = self._make_conv() + conv.append_message("User", "Hello") + conv.append_message("Assistant", None) + result = conv.to_openai_api_messages() + self.assertEqual(len(result), 1) # only user message + + def test_to_gradio_chatbot(self): + """Test conversion to Gradio chatbot format (user/assistant pairs).""" + conv = self._make_conv() + conv.append_message("User", "Q1") + conv.append_message("Assistant", "A1") + conv.append_message("User", "Q2") + conv.append_message("Assistant", "A2") + result = conv.to_gradio_chatbot() + self.assertEqual(len(result), 2) + self.assertEqual(result[0], ["Q1", "A1"]) + self.assertEqual(result[1], ["Q2", "A2"]) + + def test_to_gradio_chatbot_pending_response(self): + """Test Gradio format with pending assistant response (None).""" + conv = self._make_conv() + conv.append_message("User", "Q1") + conv.append_message("Assistant", None) + result = conv.to_gradio_chatbot() + self.assertEqual(result, [["Q1", None]]) + + def test_append_image(self): + """Test appending image data to conversation.""" + conv = self._make_conv() + conv.image_data = [] + conv.append_image("http://example.com/img.jpg", "auto") + self.assertEqual(len(conv.image_data), 1) + self.assertEqual(conv.image_data[0].url, "http://example.com/img.jpg") + self.assertEqual(conv.image_data[0].detail, "auto") + + def test_append_video(self): + """Test appending video data to conversation.""" + conv = self._make_conv() + conv.video_data = [] + conv.append_video("http://example.com/vid.mp4") + self.assertEqual(len(conv.video_data), 1) + self.assertEqual(conv.video_data[0], "http://example.com/vid.mp4") + + def test_append_audio(self): + """Test appending audio data to conversation.""" + conv = self._make_conv() + conv.audio_data = [] + conv.append_audio("http://example.com/audio.wav") + self.assertEqual(len(conv.audio_data), 1) + self.assertEqual(conv.audio_data[0], "http://example.com/audio.wav") + + def test_copy_is_independent(self): + """Test that copy() creates an independent conversation.""" + conv = self._make_conv() + conv.append_message("User", "Hello") + copied = conv.copy() + copied.append_message("Assistant", "Hi") + self.assertEqual(len(conv.messages), 1) + self.assertEqual(len(copied.messages), 2) + + def test_dict_serialization(self): + """Test dict() returns expected keys.""" + conv = self._make_conv() + conv.append_message("User", "Hello") + d = conv.dict() + self.assertEqual(d["template_name"], "test") + self.assertIn("messages", d) + self.assertIn("roles", d) + + +class TestTemplateRegistry(CustomTestCase): + def test_builtin_templates_exist(self): + """Test that common built-in templates are registered.""" + self.assertTrue(chat_template_exists("chatml")) + self.assertTrue(chat_template_exists("llama-2")) + + def test_unregistered_template_not_found(self): + """Test that non-existent template returns False.""" + self.assertFalse(chat_template_exists("_nonexistent_template_xyz")) + + def test_register_and_lookup(self): + """Test registering and looking up a custom template.""" + t = Conversation( + name="_test_conv_template", + roles=("A", "B"), + messages=[], + sep_style=SeparatorStyle.ADD_COLON_SINGLE, + sep="\n", + ) + register_conv_template(t) + self.assertTrue(chat_template_exists("_test_conv_template")) + # Cleanup + del chat_templates["_test_conv_template"] + + def test_register_duplicate_raises(self): + """Test that registering a duplicate name without override raises.""" + with self.assertRaises(AssertionError): + register_conv_template( + Conversation( + name="chatml", + roles=("A", "B"), + messages=[], + sep_style=SeparatorStyle.CHATML, + sep="", + ) + ) + + def test_get_conv_template_by_model_path_returns_none_for_unknown(self): + """Test that unknown model path returns None.""" + result = get_conv_template_by_model_path("totally-unknown-model-xyz") + self.assertIsNone(result) + + def test_get_conv_template_by_model_path_vicuna(self): + """Test that vicuna model path is matched correctly.""" + result = get_conv_template_by_model_path("lmsys/vicuna-7b-v1.5") + self.assertEqual(result, "vicuna_v1.1") + + def test_get_conv_template_by_model_path_internvl(self): + """Test that internvl model path is matched correctly.""" + result = get_conv_template_by_model_path("OpenGVLab/InternVL2-8B") + self.assertEqual(result, "internvl-2-5") + + def test_get_conv_template_by_model_path_deepseek_vl2(self): + """Test that deepseek-vl2 model path is matched correctly.""" + result = get_conv_template_by_model_path("deepseek-ai/deepseek-vl2") + self.assertEqual(result, "deepseek-vl2") + + def test_get_conv_template_by_model_path_whisper(self): + """Test that whisper model path is matched correctly.""" + result = get_conv_template_by_model_path("openai/whisper-large-v3") + self.assertEqual(result, "whisper") + + def test_get_conv_template_by_model_path_janus(self): + """Test that janus model path is matched correctly.""" + result = get_conv_template_by_model_path("deepseek-ai/Janus-Pro-7B") + self.assertEqual(result, "janus-pro") + + def test_get_conv_template_by_model_path_phi4_mm(self): + """Test that phi-4-multimodal model path is matched correctly.""" + result = get_conv_template_by_model_path("microsoft/phi-4-multimodal") + self.assertEqual(result, "phi-4-mm") + + def test_get_conv_template_by_model_path_llava_next(self): + """Test that llava-next-video-34b model path returns chatml-llava.""" + result = get_conv_template_by_model_path("llava-hf/llava-next-video-34b") + self.assertEqual(result, "chatml-llava") + + def test_get_conv_template_by_model_path_paddle_ocr(self): + """Test that paddleocr model path is matched correctly.""" + result = get_conv_template_by_model_path("PaddleOCR/PaddleOCR-2.9") + self.assertEqual(result, "paddle-ocr") + + def test_get_conv_template_by_model_path_deepseek_ocr(self): + """Test that deepseek-ocr model path is matched correctly.""" + result = get_conv_template_by_model_path("deepseek-ai/deepseek-ocr-base") + self.assertEqual(result, "deepseek-ocr") + + def test_get_conv_template_by_model_path_points(self): + """Test that points model path is matched correctly.""" + result = get_conv_template_by_model_path("WePOINTS/points-v1.5") + self.assertEqual(result, "points-v15-chat") + + def test_get_conv_template_by_model_path_minicpm_v(self): + """Test that minicpm-v model path returns minicpmv.""" + result = get_conv_template_by_model_path("openbmb/MiniCPM-V-2_6") + self.assertEqual(result, "minicpmv") + + def test_get_conv_template_by_model_path_minicpm_o(self): + """Test that minicpm-o model path returns minicpmo.""" + result = get_conv_template_by_model_path("openbmb/MiniCPM-o-2_6") + self.assertEqual(result, "minicpmo") + + +class TestGenerateEmbeddingConvs(CustomTestCase): + def test_text_only(self): + """Test generating embedding conversations with text only.""" + convs = generate_embedding_convs( + texts=["Hello world"], + images=[None], + videos=[None], + template_name="chatml", + ) + self.assertEqual(len(convs), 1) + self.assertEqual(len(convs[0].messages), 2) + self.assertIn("Hello world", convs[0].messages[0][1]) + self.assertIsNone(convs[0].messages[1][1]) # assistant placeholder + + def test_with_image(self): + """Test generating embedding conversations with image.""" + convs = generate_embedding_convs( + texts=["Describe"], + images=["http://example.com/img.jpg"], + videos=[None], + template_name="chatml", + ) + self.assertEqual(len(convs), 1) + msg = convs[0].messages[0][1] + self.assertIn("", msg) + self.assertIn("Describe", msg) + + def test_with_video(self): + """Test generating embedding conversations with video.""" + convs = generate_embedding_convs( + texts=["Describe"], + images=[None], + videos=["http://example.com/vid.mp4"], + template_name="chatml", + ) + self.assertEqual(len(convs), 1) + msg = convs[0].messages[0][1] + self.assertIn("