Upload buffer.py

examples/buffer.py

@@ -0,0 +1,59 @@
+import time
+from aworld.core.common import ActionModel, Observation
+from aworld.replay_buffer.base import (
+    DataRow,
+    DefaultConverter,
+    ReplayBuffer,
+    ExpMeta,
+    Experience,
+    RandomTaskSample
+)
+from aworld.replay_buffer.query_filter import QueryBuilder
+from aworld.logs.util import logger
+
+
+buffer = ReplayBuffer()
+
+
+def write_data():
+    for task_id in range(5):
+        for i in range(10):
+            task_id = f"task_{task_id}"
+            agent_id = f"agent_{i+1}"
+            step = i + 1
+            execute_time = time.time() + i
+            row = DataRow(
+                exp_meta=ExpMeta(
+                    task_id=task_id,
+                    task_name="default_task_name",
+                    agent_id=agent_id,
+                    step=step,
+                    execute_time=execute_time,
+                ),
+                exp_data=Experience(state=Observation(),
+                                    actions=[ActionModel()])
+            )
+            buffer.store(row)
+
+
+def read_data():
+    query = QueryBuilder().eq("exp_meta.task_id", "task_1").build()
+    datas = buffer.sample_task(query_condition=query,
+                               sampler=RandomTaskSample(),
+                               converter=DefaultConverter(),
+                               batch_size=2)
+    for data in datas:
+        logger.info(f"task_1 data: {data}")
+
+    query = QueryBuilder().eq("exp_meta.agent_id", "agent_5").build()
+    datas = buffer.sample_task(query_condition=query,
+                               sampler=RandomTaskSample(),
+                               converter=DefaultConverter(),
+                               batch_size=2)
+    for data in datas:
+        logger.info(f"agent_5 data: {data}")
+
+
+if __name__ == "__main__":
+    write_data()
+    read_data()