[fix] type annotation in test scripts

Author: Irfnfnkemed

python/mlc_llm/protocol/openai_api_protocol.py

@@ -106,7 +106,8 @@ def check_request_response_format(self) -> "RequestResponseFormat":
             for tag in self.tags:
                 if set(tag.keys()) != {"begin", "schema", "end"}:
                     raise ValueError(
-                        f"Each tag must contain exactly 'begin', 'schema' and 'end' keys. Got keys: {list(tag.keys())}."
+                        "Each tag must contain exactly 'begin', 'schema' and 'end' keys."
+                        f"Got keys: {list(tag.keys())}."
                     )
         elif self.tags is not None or self.triggers is not None:
             raise Warning(

tests/python/serve/server/test_server_structural_tag.py

@@ -11,7 +11,7 @@
 import json
 import os
 import re
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import pytest
 import requests
@@ -136,16 +136,21 @@ def check_openai_stream_response(
 
 
 def check_format(name_beg: str, name_end: str, beg_tag: str, schema: str):
-    schema = json.loads(schema)
+    try:
+        paras: Dict[str, Any] = json.loads(schema)
+    except json.JSONDecodeError as e:
+        print(f"Invalid JSON format: {e}")
+        assert False
+    assert "hash_code" in paras
     assert "hash_code" in schema
-    hash_code = schema["hash_code"]
+    hash_code = paras["hash_code"]
     assert hash_code in CHECK_INFO
     info = CHECK_INFO[hash_code]
     assert name_beg == info["name"]
     assert name_end == info["name"]
     assert beg_tag == info["beg_tag"]
     for key in info["required"]:
-        assert key in schema
+        assert key in paras
 
 
 # NOTE: the end-tag format and the hash_code number is been hidden in the SYSTEM_PROMPT.
@@ -219,7 +224,6 @@ def check_format(name_beg: str, name_end: str, beg_tag: str, schema: str):
 You are a helpful assistant.""",
 }
 
-
 STRUCTURAL_TAGS = {
     "triggers": ["<CALL--->", "<call--->"],
     "tags": [
@@ -337,7 +341,6 @@ def check_format(name_beg: str, name_end: str, beg_tag: str, schema: str):
     },
 }
 
-
 CHAT_COMPLETION_MESSAGES = [
     # messages #0
     [