|
| 1 | +"""Tests for ProviderZhipu GLM non-standard special token handling. |
| 2 | +
|
| 3 | +Covers the three layers of cleaning introduced to fix issue #5556: |
| 4 | +1. ``_clean_glm_special_tokens`` — pure regex-based cleaner |
| 5 | +2. ``_normalize_content`` — overrides the base static method |
| 6 | +3. ``_parse_openai_completion`` — second-pass cleaning on assembled text |
| 7 | +""" |
| 8 | + |
| 9 | +from unittest.mock import AsyncMock, MagicMock, patch |
| 10 | + |
| 11 | +import pytest |
| 12 | + |
| 13 | +from astrbot.core.agent.tool import ( |
| 14 | + ToolSet, # noqa: F401 – ensures the module is importable |
| 15 | +) |
| 16 | +from astrbot.core.message.message_event_result import MessageChain |
| 17 | +from astrbot.core.provider.entities import LLMResponse |
| 18 | +from astrbot.core.provider.sources.openai_source import ProviderOpenAIOfficial |
| 19 | +from astrbot.core.provider.sources.zhipu_source import ProviderZhipu |
| 20 | + |
| 21 | +# ────────────────────────────────────────────────────────────────────────────── |
| 22 | +# Helpers |
| 23 | +# ────────────────────────────────────────────────────────────────────────────── |
| 24 | + |
| 25 | + |
| 26 | +def _make_provider() -> ProviderZhipu: |
| 27 | + return ProviderZhipu( |
| 28 | + provider_config={ |
| 29 | + "id": "test-zhipu", |
| 30 | + "type": "zhipu_chat_completion", |
| 31 | + "model": "glm-4.6v-flash", |
| 32 | + "key": ["test-key"], |
| 33 | + }, |
| 34 | + provider_settings={}, |
| 35 | + ) |
| 36 | + |
| 37 | + |
| 38 | +def _make_llm_response(text: str) -> LLMResponse: |
| 39 | + """Return an LLMResponse whose completion_text equals *text*.""" |
| 40 | + r = LLMResponse("assistant") |
| 41 | + r.result_chain = MessageChain().message(text) |
| 42 | + return r |
| 43 | + |
| 44 | + |
| 45 | +# ────────────────────────────────────────────────────────────────────────────── |
| 46 | +# _clean_glm_special_tokens |
| 47 | +# ────────────────────────────────────────────────────────────────────────────── |
| 48 | + |
| 49 | + |
| 50 | +class TestCleanGLMSpecialTokens: |
| 51 | + """Unit tests for the pure-function token cleaner.""" |
| 52 | + |
| 53 | + # <None> — null-response signal ---------------------------------------- |
| 54 | + |
| 55 | + def test_null_token_alone(self): |
| 56 | + assert ProviderZhipu._clean_glm_special_tokens("<None>") == "" |
| 57 | + |
| 58 | + def test_null_token_with_leading_newline(self): |
| 59 | + # Exact pattern observed from glm-4.6v-flash: content='\n<None>' |
| 60 | + assert ProviderZhipu._clean_glm_special_tokens("\n<None>") == "" |
| 61 | + |
| 62 | + def test_null_token_surrounded_by_whitespace(self): |
| 63 | + assert ProviderZhipu._clean_glm_special_tokens(" <None> ") == "" |
| 64 | + |
| 65 | + def test_null_token_case_insensitive_lower(self): |
| 66 | + assert ProviderZhipu._clean_glm_special_tokens("<none>") == "" |
| 67 | + |
| 68 | + def test_null_token_case_insensitive_upper(self): |
| 69 | + assert ProviderZhipu._clean_glm_special_tokens("<NONE>") == "" |
| 70 | + |
| 71 | + def test_null_token_in_middle_of_text(self): |
| 72 | + result = ProviderZhipu._clean_glm_special_tokens("hello <None> world") |
| 73 | + # Token itself must be gone; surrounding spaces are collapsed to one |
| 74 | + assert "<None>" not in result |
| 75 | + assert "hello" in result and "world" in result |
| 76 | + |
| 77 | + # Role / control tokens ------------------------------------------------ |
| 78 | + |
| 79 | + def test_endoftext_token(self): |
| 80 | + assert ProviderZhipu._clean_glm_special_tokens("<|endoftext|>") == "" |
| 81 | + |
| 82 | + def test_user_role_token(self): |
| 83 | + assert ProviderZhipu._clean_glm_special_tokens("<|user|>") == "" |
| 84 | + |
| 85 | + def test_assistant_role_token(self): |
| 86 | + assert ProviderZhipu._clean_glm_special_tokens("<|assistant|>") == "" |
| 87 | + |
| 88 | + def test_system_role_token(self): |
| 89 | + assert ProviderZhipu._clean_glm_special_tokens("<|system|>") == "" |
| 90 | + |
| 91 | + def test_observation_role_token(self): |
| 92 | + assert ProviderZhipu._clean_glm_special_tokens("<|observation|>") == "" |
| 93 | + |
| 94 | + def test_role_token_prefix_removed(self): |
| 95 | + result = ProviderZhipu._clean_glm_special_tokens( |
| 96 | + "<|assistant|>Hello, how can I help?" |
| 97 | + ) |
| 98 | + assert result == "Hello, how can I help?" |
| 99 | + |
| 100 | + def test_multiple_role_tokens(self): |
| 101 | + result = ProviderZhipu._clean_glm_special_tokens("<|user|>Hi<|assistant|>Hello") |
| 102 | + assert result == "HiHello" |
| 103 | + |
| 104 | + def test_endoftext_at_end_of_reply(self): |
| 105 | + result = ProviderZhipu._clean_glm_special_tokens( |
| 106 | + "Python 最新版本是 3.13。<|endoftext|>" |
| 107 | + ) |
| 108 | + assert result == "Python 最新版本是 3.13。" |
| 109 | + |
| 110 | + # Normal text must not be affected ------------------------------------ |
| 111 | + |
| 112 | + def test_normal_text_unchanged(self): |
| 113 | + text = "我是 GLM,很高兴认识你!" |
| 114 | + assert ProviderZhipu._clean_glm_special_tokens(text) == text |
| 115 | + |
| 116 | + def test_empty_string(self): |
| 117 | + assert ProviderZhipu._clean_glm_special_tokens("") == "" |
| 118 | + |
| 119 | + def test_angle_bracket_in_normal_text_unchanged(self): |
| 120 | + """Angle brackets that are not special tokens must survive.""" |
| 121 | + text = "if a < b and b > 0: pass" |
| 122 | + assert ProviderZhipu._clean_glm_special_tokens(text) == text |
| 123 | + |
| 124 | + def test_html_like_tag_unchanged(self): |
| 125 | + """HTML-style tags (not GLM tokens) must not be stripped.""" |
| 126 | + text = "Use <strong>bold</strong> for emphasis." |
| 127 | + assert ProviderZhipu._clean_glm_special_tokens(text) == text |
| 128 | + |
| 129 | + |
| 130 | +# ────────────────────────────────────────────────────────────────────────────── |
| 131 | +# _normalize_content (static override) |
| 132 | +# ────────────────────────────────────────────────────────────────────────────── |
| 133 | + |
| 134 | + |
| 135 | +class TestNormalizeContent: |
| 136 | + """Verify that ProviderZhipu._normalize_content applies GLM cleaning on top |
| 137 | + of the base ProviderOpenAIOfficial normalisation.""" |
| 138 | + |
| 139 | + def test_null_token_string_returns_empty(self): |
| 140 | + assert ProviderZhipu._normalize_content("\n<None>") == "" |
| 141 | + |
| 142 | + def test_normal_string_unchanged(self): |
| 143 | + text = "Hello, world!" |
| 144 | + assert ProviderZhipu._normalize_content(text) == text |
| 145 | + |
| 146 | + def test_list_content_null_token(self): |
| 147 | + raw = [{"type": "text", "text": "<None>"}] |
| 148 | + assert ProviderZhipu._normalize_content(raw) == "" |
| 149 | + |
| 150 | + def test_list_content_normal_text(self): |
| 151 | + raw = [{"type": "text", "text": "Hello"}] |
| 152 | + assert ProviderZhipu._normalize_content(raw) == "Hello" |
| 153 | + |
| 154 | + def test_list_content_endoftext(self): |
| 155 | + raw = [{"type": "text", "text": "Done<|endoftext|>"}] |
| 156 | + assert ProviderZhipu._normalize_content(raw) == "Done" |
| 157 | + |
| 158 | + def test_dict_content_null_token(self): |
| 159 | + raw = {"type": "text", "text": "<None>"} |
| 160 | + assert ProviderZhipu._normalize_content(raw) == "" |
| 161 | + |
| 162 | + def test_override_is_distinct_from_base(self): |
| 163 | + """The Zhipu override should differ from the base when GLM tokens are present.""" |
| 164 | + text = "\n<None>" |
| 165 | + base_result = ProviderOpenAIOfficial._normalize_content(text) |
| 166 | + zhipu_result = ProviderZhipu._normalize_content(text) |
| 167 | + # Base keeps "<None>" after strip; Zhipu must remove it |
| 168 | + assert "<None>" not in zhipu_result |
| 169 | + assert zhipu_result == "" |
| 170 | + # Confirm the base does NOT clean it (so the override is meaningful) |
| 171 | + assert base_result == "<None>" |
| 172 | + |
| 173 | + |
| 174 | +# ────────────────────────────────────────────────────────────────────────────── |
| 175 | +# _parse_openai_completion — second-pass cleaning |
| 176 | +# ────────────────────────────────────────────────────────────────────────────── |
| 177 | + |
| 178 | + |
| 179 | +class TestParseOpenAICompletionCleaning: |
| 180 | + """Integration tests for the post-processing pass in _parse_openai_completion. |
| 181 | +
|
| 182 | + We patch ProviderOpenAIOfficial._parse_openai_completion so that we can |
| 183 | + control what the base class "returns" and verify that ProviderZhipu |
| 184 | + correctly applies the extra GLM cleaning pass on top. |
| 185 | + """ |
| 186 | + |
| 187 | + @pytest.mark.asyncio |
| 188 | + async def test_null_token_content_becomes_empty(self): |
| 189 | + """content='\\n<None>' (real API response) should produce an empty reply.""" |
| 190 | + provider = _make_provider() |
| 191 | + try: |
| 192 | + fake_completion = MagicMock() |
| 193 | + parent_response = _make_llm_response("\n<None>") |
| 194 | + |
| 195 | + with patch.object( |
| 196 | + ProviderOpenAIOfficial, |
| 197 | + "_parse_openai_completion", |
| 198 | + new=AsyncMock(return_value=parent_response), |
| 199 | + ): |
| 200 | + result = await provider._parse_openai_completion(fake_completion, None) |
| 201 | + |
| 202 | + assert result.completion_text == "" |
| 203 | + finally: |
| 204 | + await provider.terminate() |
| 205 | + |
| 206 | + @pytest.mark.asyncio |
| 207 | + async def test_endoftext_token_stripped_from_end(self): |
| 208 | + provider = _make_provider() |
| 209 | + try: |
| 210 | + parent_response = _make_llm_response("当然可以!<|endoftext|>") |
| 211 | + |
| 212 | + with patch.object( |
| 213 | + ProviderOpenAIOfficial, |
| 214 | + "_parse_openai_completion", |
| 215 | + new=AsyncMock(return_value=parent_response), |
| 216 | + ): |
| 217 | + result = await provider._parse_openai_completion(MagicMock(), None) |
| 218 | + |
| 219 | + assert result.completion_text == "当然可以!" |
| 220 | + finally: |
| 221 | + await provider.terminate() |
| 222 | + |
| 223 | + @pytest.mark.asyncio |
| 224 | + async def test_assistant_role_token_prefix_stripped(self): |
| 225 | + provider = _make_provider() |
| 226 | + try: |
| 227 | + parent_response = _make_llm_response("<|assistant|>我是一个AI助手。") |
| 228 | + |
| 229 | + with patch.object( |
| 230 | + ProviderOpenAIOfficial, |
| 231 | + "_parse_openai_completion", |
| 232 | + new=AsyncMock(return_value=parent_response), |
| 233 | + ): |
| 234 | + result = await provider._parse_openai_completion(MagicMock(), None) |
| 235 | + |
| 236 | + assert result.completion_text == "我是一个AI助手。" |
| 237 | + finally: |
| 238 | + await provider.terminate() |
| 239 | + |
| 240 | + @pytest.mark.asyncio |
| 241 | + async def test_normal_content_unchanged(self): |
| 242 | + """Normal GLM replies must not be modified.""" |
| 243 | + provider = _make_provider() |
| 244 | + try: |
| 245 | + normal = "好的,我来帮你解答这个问题。" |
| 246 | + parent_response = _make_llm_response(normal) |
| 247 | + |
| 248 | + with patch.object( |
| 249 | + ProviderOpenAIOfficial, |
| 250 | + "_parse_openai_completion", |
| 251 | + new=AsyncMock(return_value=parent_response), |
| 252 | + ): |
| 253 | + result = await provider._parse_openai_completion(MagicMock(), None) |
| 254 | + |
| 255 | + assert result.completion_text == normal |
| 256 | + finally: |
| 257 | + await provider.terminate() |
| 258 | + |
| 259 | + @pytest.mark.asyncio |
| 260 | + async def test_empty_completion_text_not_modified(self): |
| 261 | + """When the base class returns empty completion_text, don't error out.""" |
| 262 | + provider = _make_provider() |
| 263 | + try: |
| 264 | + parent_response = LLMResponse("assistant") |
| 265 | + parent_response.result_chain = None |
| 266 | + parent_response._completion_text = "" |
| 267 | + |
| 268 | + with patch.object( |
| 269 | + ProviderOpenAIOfficial, |
| 270 | + "_parse_openai_completion", |
| 271 | + new=AsyncMock(return_value=parent_response), |
| 272 | + ): |
| 273 | + result = await provider._parse_openai_completion(MagicMock(), None) |
| 274 | + |
| 275 | + assert result.completion_text == "" |
| 276 | + finally: |
| 277 | + await provider.terminate() |
| 278 | + |
| 279 | + @pytest.mark.asyncio |
| 280 | + async def test_reasoning_content_preserved(self): |
| 281 | + """Cleaning must not touch reasoning_content.""" |
| 282 | + provider = _make_provider() |
| 283 | + try: |
| 284 | + parent_response = _make_llm_response("\n<None>") |
| 285 | + parent_response.reasoning_content = "思考过程:用户打了招呼,不需要回复。" |
| 286 | + |
| 287 | + with patch.object( |
| 288 | + ProviderOpenAIOfficial, |
| 289 | + "_parse_openai_completion", |
| 290 | + new=AsyncMock(return_value=parent_response), |
| 291 | + ): |
| 292 | + result = await provider._parse_openai_completion(MagicMock(), None) |
| 293 | + |
| 294 | + assert result.completion_text == "" |
| 295 | + assert "思考过程" in result.reasoning_content |
| 296 | + finally: |
| 297 | + await provider.terminate() |
| 298 | + |
| 299 | + @pytest.mark.asyncio |
| 300 | + async def test_other_response_fields_preserved(self): |
| 301 | + """id, usage and other metadata must survive the cleaning pass.""" |
| 302 | + provider = _make_provider() |
| 303 | + try: |
| 304 | + parent_response = _make_llm_response("普通回复") |
| 305 | + parent_response.id = "cmp-test-id-123" |
| 306 | + |
| 307 | + with patch.object( |
| 308 | + ProviderOpenAIOfficial, |
| 309 | + "_parse_openai_completion", |
| 310 | + new=AsyncMock(return_value=parent_response), |
| 311 | + ): |
| 312 | + result = await provider._parse_openai_completion(MagicMock(), None) |
| 313 | + |
| 314 | + assert result.id == "cmp-test-id-123" |
| 315 | + assert result.completion_text == "普通回复" |
| 316 | + finally: |
| 317 | + await provider.terminate() |
0 commit comments