From 5a860b7b54751e82b4557dcd5db56eda34608332 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Tue, 12 Nov 2024 11:29:35 +0900
Subject: [PATCH] Update mix python script for new log format

---
 scripts/latency.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/scripts/latency.py b/scripts/latency.py
index 6fd970c..6dd84b4 100755
--- a/scripts/latency.py
+++ b/scripts/latency.py
@@ -6,10 +6,6 @@ from typing import Dict, Optional
 import statistics
 import argparse
 
-from json_stream.base import TransientStreamingJSONObject
-
-JsonStream = Iterable[TransientStreamingJSONObject]
-
 
 class Message:
     def __init__(self, message_id: str, step_a: Optional[int]):
@@ -71,8 +67,13 @@ def parse_record_stream(record_stream: Iterable[str]) -> MessageStorage:
     storage: MessageStorage = {}
 
     for record in record_stream:
-        json_record = json.loads(record)
-        payload_id = json_record["payload_id"]
+        try:
+            json_record = json.loads(record)
+        except json.decoder.JSONDecodeError:
+            continue
+
+        if (payload_id := json_record.get("payload_id")) is None:
+            continue
         step_id = json_record["step_id"]
 
         if (stored_message := storage.get(payload_id)) is None: