From 4dd70afba7d029cb1d9ae4a87c6a670d290a9e1c Mon Sep 17 00:00:00 2001
From: Jason <jtoy@grids.local>
Date: Wed, 28 Feb 2024 17:00:09 -0800
Subject: [PATCH] tools for loading into redis

---
 utils/load_jsonl_to_redis.py   | 39 ++++++++++++++++++++++++++++++++++
 web/src/components/Sidebar.tsx |  1 +
 2 files changed, 40 insertions(+)
 create mode 100644 utils/load_jsonl_to_redis.py

diff --git a/utils/load_jsonl_to_redis.py b/utils/load_jsonl_to_redis.py
new file mode 100644
index 0000000..36d3387
--- /dev/null
+++ b/utils/load_jsonl_to_redis.py
@@ -0,0 +1,39 @@
+import json
+import redis
+import sys
+import os
+
+def inject_jsonl_to_redis_queue(file_path, queue_name, max_steps=None):
+    # Connect to Redis
+    redis_url = os.getenv('REDIS_URL', 'redis://localhost:6379/0')
+    redis_client = redis.StrictRedis.from_url(redis_url, decode_responses=True)
+
+    # Read JSONL file
+    with open(file_path, 'r') as jsonl_file:
+        for step, line in enumerate(jsonl_file, start=1):
+            if max_steps is not None and step > max_steps:
+                break
+
+            try:
+                # Parse JSON from each line
+                data = json.loads(line.strip())
+
+                # Push data into the Redis queue
+                redis_client.rpush(queue_name, json.dumps(data))
+            except json.JSONDecodeError as e:
+                print(f"Error decoding JSON: {e}")
+                continue
+
+    print(f"Data from {file_path} injected into Redis queue {queue_name} ({step} lines)")
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3 or len(sys.argv) > 4:
+        print("Usage: python script.py <jsonl_file_path> <redis_queue_name> [max_steps]")
+        sys.exit(1)
+
+    jsonl_file_path = sys.argv[1]
+    redis_queue_name = sys.argv[2]
+    max_steps = int(sys.argv[3]) if len(sys.argv) == 4 else None if len(sys.argv) == 3 else None
+
+    inject_jsonl_to_redis_queue(jsonl_file_path, redis_queue_name, max_steps)
+
diff --git a/web/src/components/Sidebar.tsx b/web/src/components/Sidebar.tsx
index d50a50b..8fce8a3 100644
--- a/web/src/components/Sidebar.tsx
+++ b/web/src/components/Sidebar.tsx
@@ -113,6 +113,7 @@ const Sidebar: React.FC<SidebarProps> = (
                 <div>
                     Step: {stepId}
                 </div>
+                <span style={{ display: 'none' }}>{substepId}</span>
                 <div className={styles.buttons}>
                     <button onClick={handleRewind}>Rewind</button>
                     <button onClick={handlePause} disabled={!isPlaying}>Pause</button>