ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/tests/replay_buffer/__init__.py

@@ -0,0 +1,2 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.

chapter9/gaia-experience/AWorld/tests/replay_buffer/buffer.py

@@ -0,0 +1,54 @@
+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}")

chapter9/gaia-experience/AWorld/tests/replay_buffer/multi_processing.py

@@ -0,0 +1,98 @@
+import time
+import traceback
+import multiprocessing
+from aworld import replay_buffer
+from aworld.core.common import ActionModel, Observation
+from aworld.replay_buffer.base import ReplayBuffer, DataRow, ExpMeta, Experience
+from aworld.replay_buffer.query_filter import QueryBuilder
+from aworld.replay_buffer.storage.multi_proc_mem import MultiProcMemoryStorage
+from aworld.logs.util import logger
+
+
+def write_processing(replay_buffer: ReplayBuffer, task_id: str):
+    for i in range(10):
+        try:
+            data = DataRow(
+                exp_meta=ExpMeta(
+                    task_id=task_id,
+                    task_name=task_id,
+                    agent_id=f"agent_{i+1}",
+                    step=i,
+                    execute_time=time.time()
+                ),
+                exp_data=Experience(state=Observation(),
+                                    actions=[ActionModel()])
+            )
+            replay_buffer.store(data)
+        except Exception as e:
+            stack_trace = traceback.format_exc()
+            logger.error(
+                f"write_processing error: {e}\nStack trace:\n{stack_trace}")
+        time.sleep(1)
+
+
+def read_processing_by_task(replay_buffer: ReplayBuffer, task_id: str):
+    while True:
+        try:
+            query_condition = QueryBuilder().eq("exp_meta.task_id", task_id).build()
+            data = replay_buffer.sample_task(
+                query_condition=query_condition, batch_size=2)
+            logger.info(f"read data of task[{task_id}]: {data}")
+        except Exception as e:
+            stack_trace = traceback.format_exc()
+            logger.error(
+                f"read_processing_by_task error: {e}\nStack trace:\n{stack_trace}")
+        time.sleep(1)
+
+
+def read_processing_by_agent(replay_buffer: ReplayBuffer, agent_id: str):
+    while True:
+        try:
+            query_condition = QueryBuilder().eq("exp_meta.agent_id", agent_id).build()
+            data = replay_buffer.sample_task(
+                query_condition=query_condition, batch_size=2)
+            logger.info(f"read data of agent[{agent_id}]: {data}")
+        except Exception as e:
+            logger.info(f"read_processing_by_agent error: {e}")
+        time.sleep(1)
+
+
+def run():
+    multiprocessing.freeze_support()
+    multiprocessing.set_start_method('spawn')
+    manager = multiprocessing.Manager()
+
+    replay_buffer = ReplayBuffer(storage=MultiProcMemoryStorage(
+        data_dict=manager.dict(),
+        fifo_queue=manager.list(),
+        lock=manager.Lock(),
+        max_capacity=10000
+    ))
+
+    processes = [
+        multiprocessing.Process(target=write_processing,
+                                args=(replay_buffer, "task_1",)),
+        multiprocessing.Process(target=write_processing,
+                                args=(replay_buffer, "task_2",)),
+        multiprocessing.Process(target=write_processing,
+                                args=(replay_buffer, "task_3",)),
+        multiprocessing.Process(target=write_processing,
+                                args=(replay_buffer, "task_4",)),
+        # multiprocessing.Process(
+        #     target=read_processing_by_task, args=(replay_buffer, "task_1",)),
+        multiprocessing.Process(
+            target=read_processing_by_agent, args=(replay_buffer, "agent_3",))
+    ]
+    for p in processes:
+        p.start()
+
+    try:
+        for p in processes:
+            p.join()
+    except KeyboardInterrupt:
+        for p in processes:
+            p.terminate()
+        for p in processes:
+            p.join()
+    finally:
+        logger.info("Processes terminated.")

chapter9/gaia-experience/AWorld/tests/replay_buffer/query_builder.py

@@ -0,0 +1,96 @@
+from aworld.replay_buffer.query_filter import QueryBuilder
+from aworld.logs.util import logger
+
+
+def example():
+    '''
+    expression: task_id = "123"
+    return :
+        {
+            'field': 'task_id',
+            'value': '123',
+            'op': 'eq'
+        }
+    '''
+    qb = QueryBuilder()
+    query = qb.eq("task_id", "123").build()
+    logger.info(query)
+
+
+def example1():
+    '''
+    expression: (task_id = "123" and agent_id = "111") or (task_id = "456" and agent_id = "222")
+    return : 
+        {
+            'or_': [{
+                'and_': [{
+                    'field': 'task_id',
+                    'value': '123',
+                    'op': 'eq'
+                }, {
+                    'field': 'agent_id',
+                    'value': '111',
+                    'op': 'eq'
+                }]
+            }, {
+                'and_': [{
+                    'field': 'task_id',
+                    'value': '456',
+                    'op': 'eq'
+                }, {
+                    'field': 'agent_id',
+                    'value': '222',
+                    'op': 'eq'
+                }]
+            }]
+        }
+    '''
+    qb = QueryBuilder()
+    query = (qb.eq("task_id", "123")
+             .and_()
+             .eq("agent_id", "111")
+             .or_()
+             .nested(QueryBuilder()
+                     .eq("task_id", "456")
+                     .and_()
+                     .eq("agent_id", "222"))
+             .build())
+    logger.info(query)
+
+
+def example2():
+    '''
+    expression: task_id = "123" and (agent_id = "111" or agent_id = "222")
+    return :
+        {
+            'and_': [{
+                'field': 'task_id',
+                'value': '123',
+                'op': 'eq'
+            }, {
+                'or_': [{
+                    'field': 'agent_id',
+                    'value': '111',
+                    'op': 'eq'
+                }, {
+                    'field': 'agent_id',
+                    'value': '222',
+                    'op': 'eq'
+                }
+            }
+        }   
+    '''
+    qb = QueryBuilder()
+    query = (qb.eq("task_id", "123")
+             .and_()
+             .nested(QueryBuilder()
+                     .eq("agent_id", "111")
+                     .or_()
+                     .eq("agent_id", "222"))
+             .build())
+    logger.info(query)
+
+if __name__ == "__main__":
+    example()
+    example1()
+    example2()

chapter9/gaia-experience/AWorld/tests/replay_buffer/query_filter.py

@@ -0,0 +1,37 @@
+import time
+from aworld.replay_buffer.base import (
+    DataRow,
+    DefaultConverter,
+    ReplayBuffer,
+    ExpMeta,
+    Experience,
+)
+from aworld.core.common import ActionModel, Observation
+from aworld.replay_buffer.query_filter import QueryBuilder, QueryFilter
+from aworld.logs.util import logger
+
+
+def filter():
+    row = DataRow(
+        exp_meta=ExpMeta(
+            task_id="task_1",
+            task_name="default_task_name",
+            agent_id="agent_1",
+            step=1,
+            execute_time=time.time(),
+        ),
+        exp_data=Experience(state=Observation(), action=[ActionModel()])
+    )
+
+    query = QueryBuilder().eq("exp_meta.task_id", "task_1").build()
+    filter1 = QueryFilter(query)
+    assert filter1.check_condition(row)
+
+    query = QueryBuilder().eq("exp_meta.task_id", "task_2").build()
+    filter2 = QueryFilter(query)
+    assert not filter2.check_condition(row)
+
+    query = QueryBuilder().eq("exp_meta.task_id", "task_1").and_().eq(
+        "exp_meta.agent_id", "agent_2").build()
+    filter3 = QueryFilter(query)
+    assert not filter3.check_condition(row)

chapter9/gaia-experience/AWorld/tests/replay_buffer/storage_odps.py

@@ -0,0 +1,65 @@
+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
+from aworld.replay_buffer.storage.odps import OdpsStorage
+
+
+buffer = ReplayBuffer(storage=OdpsStorage(
+    table_name="adm_aworld_replay_buffer",
+    project="alifin_jtest_dev",
+    endpoint="",
+    access_id="",
+    access_key=""
+))
+
+
+def write_data():
+    rows = []
+    for id in range(5):
+        task_id = f"task_{id+1}"
+        for i in range(5):
+            agent_id = f"agent_{i+1}"
+            for j in range(5):
+                step = j + 1
+                execute_time = time.time() + j
+                row = DataRow(
+                    exp_meta=ExpMeta(
+                        task_id=task_id,
+                        task_name="default_task_name",
+                        agent_id=agent_id,
+                        step=step,
+                        execute_time=execute_time,
+                        pre_agent="pre_agent_id"
+                    ),
+                    exp_data=Experience(state=Observation(),
+                                        actions=[ActionModel()])
+                )
+                rows.append(row)
+    buffer.store_batch(rows)
+
+
+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=1)
+    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}")

chapter9/gaia-experience/AWorld/tests/replay_buffer/storage_redis.py

@@ -0,0 +1,68 @@
+import time
+from aworld.replay_buffer.base import DataRow, ExpMeta, Experience
+from aworld.replay_buffer.storage.redis import RedisStorage
+from aworld.replay_buffer.query_filter import QueryBuilder
+from aworld.core.common import Observation, ActionModel
+from aworld.logs.util import logger
+
+
+def generate_data_row() -> list[DataRow]:
+    rows: list[DataRow] = []
+    for id in range(5):
+        task_id = f"task_{id+1}"
+        for i in range(5):
+            agent_id = f"agent_{i+1}"
+            for j in range(5):
+                step = j + 1
+                execute_time = time.time() + j
+                row = DataRow(
+                    exp_meta=ExpMeta(
+                        task_id=task_id,
+                        task_name="default_task_name",
+                        agent_id=agent_id,
+                        step=step,
+                        execute_time=execute_time,
+                        pre_agent="pre_agent_id"
+                    ),
+                    exp_data=Experience(state=Observation(),
+                                        actions=[ActionModel()])
+                )
+                rows.append(row)
+    return rows
+
+
+def wriete_data(storage):
+    storage.clear()
+    rows = generate_data_row()
+    storage.add_batch(rows)
+    logger.info(f"Add {len(rows)} rows to storage.")
+
+
+def read_data(storage):
+    query_condition = (QueryBuilder()
+                       .eq("exp_meta.task_id", "task_1")
+                       .and_()
+                       .eq("exp_meta.agent_id", "agent_1")
+                       .or_()
+                       .nested(QueryBuilder()
+                               .eq("exp_meta.task_id", "task_4")
+                               .and_()
+                               .eq("exp_meta.agent_id", "agent_3")
+                               .and_()
+                               .gt("exp_meta.step", 4)).build())
+
+    rows = storage.get_all(query_condition)
+    for row in rows:
+        logger.info(row)
+
+    rows = storage.get_paginated(
+        page=2, page_size=2, query_condition=query_condition)
+    for row in rows:
+        logger.info(f"get_paginated: {row}")
+
+
+# if __name__ == "__main__":
+#     storage = RedisStorage(host="localhost", port=6379,
+#                            recreate_idx_if_exists=False)
+#     wriete_data(storage)
+#     read_data(storage)