add test_id for test_json_mode

Signed-off-by: Linkun Chen <github@lkchen.net>

Author: lk-chen

release/llm_tests/serve/probes/test_json_mode.py

@@ -58,28 +58,34 @@ class NestedResponse(BaseModel):
     sorted_list: List[Person] = Field(description="List of the sorted objects")
 
 
-def get_params_and_expected_type(response_type: str):
+def get_params_and_expected_type(response_type: str, test_id: str):
     params = {}
     if response_type == "basic":
         params.update(
             **messages(
-                system("You are a helpful assistant designed to output JSON."),
+                system(
+                    f"{test_id} You are a helpful assistant designed to output JSON."
+                ),
                 user("Who won the world series in 2020?"),
             )
         )
         expected_type = BasicResponse
     elif response_type == "array":
         params.update(
             **messages(
-                system("You are a helpful assistant designed to output JSON."),
+                system(
+                    f"{test_id} You are a helpful assistant designed to output JSON."
+                ),
                 user("Sort the numbers 3, 1, 2, 4, 5"),
             )
         )
         expected_type = ArrayResponse
     elif response_type == "nested":
         params.update(
             **messages(
-                system("You are a helpful assistant designed to output JSON."),
+                system(
+                    f"{test_id} You are a helpful assistant designed to output JSON."
+                ),
                 user(
                     "Sort these people by age: John, 20 years old, Mary, 30 years old, Bob, 10 years old."
                 ),
@@ -116,11 +122,11 @@ def get_response_formats():
 
 
 async def query_json_model(
-    model: str, response_type: str, stream: bool, openai_async_client
+    model: str, response_type: str, stream: bool, openai_async_client, test_id: str
 ):
     querier = TextGenerationProbeQuerier(openai_async_client, {"temperature": 0.0})
 
-    params, expected_type = get_params_and_expected_type(response_type)
+    params, expected_type = get_params_and_expected_type(response_type, test_id)
     response = await querier.query(model, stream=stream, **params)
     response_str = response.full()
 
@@ -148,7 +154,7 @@ async def test_json_mode(
 
     responses = await asyncio.gather(
         *[
-            query_json_model(model, response_type, stream, openai_async_client)
+            query_json_model(model, response_type, stream, openai_async_client, test_id)
             for _ in range(n_concurrent_requests)
         ]
     )
@@ -172,7 +178,7 @@ async def test_response_format_options(
 
     params = {
         **messages(
-            system("You are a helpful assistant designed to output JSON."),
+            system(f"{test_id} You are a helpful assistant designed to output JSON."),
             user("Who won the world series in 2020?"),
         ),
         "response_format": response_format,