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

chapter9/gaia-experience/AWorld/aworld/replay_buffer/README.md

@@ -0,0 +1,111 @@
+# Replay Buffer
+
+A multi-process capable replay buffer system for storing and sampling experience data.
+
+## Features
+
+- **Multi-process Support**: Safe concurrent access using shared memory and locks
+- **Flexible Querying**: Powerful query builder for filtering stored data
+- **Task-based Organization**: Data organized by task_id and agent_id
+- **Capacity Management**: FIFO eviction when reaching max capacity
+- **Custom Sampling**: Implement custom sampling logic through Sampler interface
+- **Data Conversion**: Custom data conversion through Converter interface
+
+## Basic Usage
+
+### Writing Data
+
+```python
+from aworld.replay_buffer import ReplayBuffer, DataRow, ExpMeta, Experience
+from aworld.core.common import ActionModel, Observation
+
+# Create a data row
+data = DataRow(
+    exp_meta=ExpMeta(
+        task_id="task_1",
+        task_name="my_task",
+        agent_id="agent_1",
+        step=1,
+        execute_time=time.time()
+    ),
+    exp_data=Experience(
+        state=Observation(),
+        action=ActionModel()
+    )
+)
+
+# Store data
+replay_buffer.store(data)
+```
+
+### Reading Data
+
+```python
+from aworld.replay_buffer.query_filter import QueryBuilder
+
+# Basic example
+replay_buffer = ReplayBuffer()
+query_condition = QueryBuilder().eq("exp_meta.task_name", "test_task").build()
+data = replay_buffer.sample(sampler=RandomTaskSample(),
+                            query_condition=query_condition,
+                            converter=DefaultConverter(),
+                            batch_size=1000)
+
+# Query Task by task_id
+query = QueryBuilder().eq("exp_meta.task_id", "task_1").build()
+data = replay_buffer.sample_task(query_condition=query, batch_size=10)
+
+# Query Task by agent_id 
+query = QueryBuilder().eq("exp_meta.agent_id", "agent_1").build()
+data = replay_buffer.sample_task(query_condition=query, batch_size=5)
+```
+## Multi-processing Example
+
+```python
+import multiprocessing
+from aworld.replay_buffer.storage.multi_proc_mem import MultiProcMemoryStorage
+
+manager = multiprocessing.Manager()
+replay_buffer = ReplayBuffer(
+    storage=MultiProcMemoryStorage(
+        data_dict=manager.dict(),
+        fifo_queue=manager.list(),
+        lock=manager.Lock(),
+        max_capacity=10000
+    )
+)
+
+# Start writer processes
+processes = [
+    multiprocessing.Process(target=write_processing, args=(replay_buffer, f"task_{i}"))
+    for i in range(4)
+]
+```
+## Query Builder Examples
+
+### Simple Equality
+```python
+QueryBuilder().eq("exp_meta.task_id", "123").build()
+```
+
+### Complex Conditions
+```python
+QueryBuilder()
+    .eq("exp_meta.task_id", "123")
+    .and_()
+    .eq("exp_meta.agent_id", "456")
+    .build()
+```
+### Nested Conditions
+```python
+QueryBuilder()
+    .eq("exp_meta.task_id", "123")
+    .and_()
+    .nested(
+        QueryBuilder()
+            .eq("exp_meta.agent_id", "111")
+            .or_()
+            .eq("exp_meta.agent_id", "222")
+    )
+    .build()
+```

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

@@ -0,0 +1,13 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+
+from aworld.replay_buffer.base import ReplayBuffer, DataRow, ExpMeta, Experience
+from aworld.replay_buffer.event_replay_buffer import EventReplayBuffer
+
+__all__ = [
+    'ReplayBuffer',
+    'DataRow',
+    'ExpMeta',
+    'Experience',
+    'EventReplayBuffer',
+]

chapter9/gaia-experience/AWorld/aworld/replay_buffer/base.py

@@ -0,0 +1,410 @@
+import random
+import uuid
+from dataclasses import dataclass, field
+from typing import Dict, List, TypeVar
+from abc import ABC, abstractmethod
+from math import ceil
+
+from aworld.core.common import ActionModel, Observation
+from aworld.replay_buffer.query_filter import QueryCondition, QueryFilter
+from aworld.logs.util import logger
+from aworld.utils.serialized_util import to_serializable
+
+T = TypeVar('T')
+
+
+@dataclass
+class Experience:
+    '''
+    Experience of agent.
+    '''
+    state: Observation
+    actions: List[ActionModel]
+    reward_t: float = None
+    adv_t: float = None
+    v_t: float = None
+    messages: List[Dict] = None
+
+    def to_dict(self):
+        return {
+            "state": to_serializable(self.state),
+            "actions": to_serializable(self.actions),
+            "reward_t": self.reward_t,
+            "adv_t": self.adv_t,
+            "v_t": self.v_t,
+            "messages": self.messages
+        }
+
+
+@dataclass
+class ExpMeta:
+    '''
+    Experience meta data.
+    '''
+    task_id: str
+    task_name: str
+    agent_id: str
+    step: int
+    execute_time: float
+    pre_agent: str
+
+    def to_dict(self):
+        return {
+            "task_id": self.task_id,
+            "task_name": self.task_name,
+            "agent_id": self.agent_id,
+            "step": self.step,
+            "execute_time": self.execute_time,
+            "pre_agent": self.pre_agent
+        }
+@dataclass
+class DataRow:
+    '''
+    Data row for storing data.
+    '''
+    exp_meta: ExpMeta
+    exp_data: Experience
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+    def to_dict(self):
+        return {
+            "exp_meta": self.exp_meta.to_dict(),
+            "exp_data": self.exp_data.to_dict(),
+            "id": self.id
+        }
+
+
+class Storage(ABC):
+    '''
+    Storage for storing and sampling data.
+    '''
+
+    @abstractmethod
+    def add(self, data: DataRow):
+        '''
+        Add data to the storage.
+        Args:
+            data (DataRow): Data to add.
+        '''
+
+    @abstractmethod
+    def add_batch(self, data_batch: List[DataRow]):
+        '''
+        Add batch of data to the storage.
+        Args:
+            data_batch (List[DataRow]): List of data to add.
+        '''
+
+    @abstractmethod
+    def size(self, query_condition: QueryCondition = None) -> int:
+        '''
+        Get the size of the storage.
+        Returns:
+            int: Size of the storage.
+        '''
+
+    @abstractmethod
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        '''
+        Get paginated data from the storage.
+        Args:
+            page (int): Page number.
+            page_size (int): Number of data per page.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+
+    @abstractmethod
+    def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        '''
+        Get all data from the storage.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+
+    @abstractmethod
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        '''
+        Get data by task_id from the storage.
+        Args:
+            task_id (str): Task id.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+
+    @abstractmethod
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        '''
+        Get batch of data by task_ids from the storage.
+        Args:
+            task_ids (List[str]): List of task ids.
+        Returns:
+            Dict[str, List[DataRow]]: Dictionary of data.
+            The key is the task_id and the value is the list of data.
+            The list of data is sorted by step.
+        '''
+
+
+class Sampler(ABC):
+    '''
+    Sample data from the storage.
+    '''
+
+    def sample(self,
+               storage: Storage,
+               batch_size: int,
+               query_condition: QueryCondition = None) -> List[DataRow]:
+        '''
+        Sample data from the storage.
+        Args:
+            storage (Storage): Storage to sample from.
+            batch_size (int): Number of data to sample.
+            query_condition (QueryCondition, optional): Query condition. Defaults to None.
+        Returns:
+            List[DataRow]
+        '''
+
+
+class TaskSampler(Sampler):
+    '''
+    Sample task data from storage, returns Dict[str, List[DataRow]] where:
+    - key is task_id
+    - value is list of task all data rows
+    '''
+
+    def sorted_by_step(self, task_experience: List[DataRow]) -> List[DataRow]:
+        '''
+        Sort the task experience by step and execute_time.
+        Args:
+            task_experience (List[DataRow]): List of task experience.
+        Returns:
+            List[DataRow]: List of task experience sorted by step and execute_time.
+        '''
+        return sorted(task_experience, key=lambda x: (x.exp_meta.step, x.exp_meta.execute_time))
+
+    def sample(self,
+               storage: Storage,
+               batch_size: int,
+               query_condition: QueryCondition = None) -> List[DataRow]:
+        task_ids = self.sample_task_ids(storage, batch_size, query_condition)
+        return storage.get_bacth_by_task_ids(task_ids)
+
+    def sample_tasks(self,
+                     storage: Storage,
+                     batch_size: int,
+                     query_condition: QueryCondition = None) -> Dict[str, List[DataRow]]:
+        '''
+        Sample data from the storage.
+        Args:
+            storage (Storage): Storage to sample from.
+            batch_size (int): Number of data to sample.
+            query_condition (QueryCondition, optional): Query condition. Defaults to None.
+        Returns:
+            Dict[str, List[DataRow]]: Dictionary of sampled data.
+            The key is the task_id and the value is the list of data.
+            The list of data is sorted by step.
+        '''
+        task_ids = self.sample_task_ids(storage, batch_size, query_condition)
+        raws = storage.get_bacth_by_task_ids(task_ids)
+        return {task_id: self.sorted_by_step(raws) for task_id, raws in raws.items()}
+
+    @abstractmethod
+    def sample_task_ids(self,
+                        storage: Storage,
+                        batch_size: int,
+                        query_condition: QueryCondition = None) -> List[str]:
+        '''
+        Sample task_ids from the storage.
+        Args:
+            storage (Storage): Storage to sample from.
+            batch_size (int): Number of task_ids to sample.
+            query_condition (QueryCondition, optional): Query condition. Defaults to None.
+        Returns:
+            List[str]: List of task_ids.
+        '''
+
+
+class Converter(ABC):
+    '''
+    Convert data to dataset row.
+    '''
+
+    @abstractmethod
+    def to_dataset_row(self, task_experience: List[DataRow]) -> T:
+        '''
+        Convert task experience to dataset row.
+        Args:
+            task_experience (List[DataRow]): List of task experience.
+        Returns:
+            T: type of dataset row.
+        '''
+
+
+class InMemoryStorage(Storage):
+    '''
+    In-memory storage for storing and sampling data.
+    '''
+
+    def __init__(self, max_capacity: int = 10000):
+        self._data: Dict[str, List[DataRow]] = {}
+        self._max_capacity = max_capacity
+        self._fifo_queue = []  # (task_id)
+
+    def add(self, data: DataRow):
+        if not data:
+            raise ValueError("Data is required")
+        if not data.exp_meta:
+            raise ValueError("exp_meta is required")
+
+        while self.size() >= self._max_capacity and self._fifo_queue:
+            oldest_task_id = self._fifo_queue.pop(0)
+            if oldest_task_id in self._data:
+                del self._data[oldest_task_id]
+
+        if data.exp_meta.task_id not in self._data:
+            self._data[data.exp_meta.task_id] = []
+        self._data[data.exp_meta.task_id].append(data)
+        self._fifo_queue.append(data.exp_meta.task_id)
+
+        if data.exp_meta.task_id not in self._data:
+            self._data[data.exp_meta.task_id] = []
+        self._data[data.exp_meta.task_id].append(data)
+
+    def add_batch(self, data_batch: List[DataRow]):
+        for data in data_batch:
+            self.add(data)
+
+    def size(self, query_condition: QueryCondition = None) -> int:
+        return len(self.get_all(query_condition))
+
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        if page < 1:
+            raise ValueError("Page must be greater than 0")
+        if page_size < 1:
+            raise ValueError("Page size must be greater than 0")
+        all_data = self.get_all(query_condition)
+        start_index = (page - 1) * page_size
+        end_index = start_index + page_size
+        return all_data[start_index:end_index]
+
+    def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        all_data = []
+        query_filter = None
+        if query_condition:
+            query_filter = QueryFilter(query_condition)
+        for data in self._data.values():
+            if query_filter:
+                all_data.extend(query_filter.filter(data))
+            else:
+                all_data.extend(data)
+        return all_data
+
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        return self._data.get(task_id, [])
+
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        return {task_id: self._data.get(task_id, []) for task_id in task_ids}
+
+    def clear(self):
+        self._data = {}
+        self._fifo_queue = []
+
+
+class RandomTaskSample(TaskSampler):
+    '''
+    Randomly sample data from the storage.
+    '''
+
+    def sample_task_ids(self,
+                        storage: Storage,
+                        batch_size: int,
+                        query_condition: QueryCondition = None) -> List[str]:
+        total_size = storage.size(query_condition)
+        if total_size <= batch_size:
+            return storage.get_all(query_condition)
+
+        sampled_task_ids = set()
+        page_size = min(100, batch_size * 2)
+        total_pages = ceil(total_size/page_size)
+        visited_pages = set()
+        while len(sampled_task_ids) < batch_size and len(visited_pages) < total_pages:
+            page = random.choice(
+                [p for p in range(1, total_pages+1) if p not in visited_pages])
+            visited_pages.add(page)
+
+            current_page = storage.get_paginated(
+                page, page_size, query_condition)
+            if not current_page:
+                continue
+            current_page_task_ids = set(
+                [data.exp_meta.task_id for data in current_page if data.exp_meta.task_id not in sampled_task_ids])
+            sample_count = min(len(current_page_task_ids),
+                               batch_size - len(sampled_task_ids))
+            sampled_task_ids.update(random.sample(
+                list(current_page_task_ids), sample_count))
+
+        return list(sampled_task_ids)
+
+
+class DefaultConverter(Converter):
+    '''
+    Default converter do nothing.
+    '''
+
+    def to_dataset_row(self, task_experience: List[DataRow]) -> List[DataRow]:
+        return task_experience
+
+
+class ReplayBuffer:
+    '''
+    Replay buffer for storing and sampling data.
+    '''
+
+    def __init__(
+        self,
+        storage: Storage = InMemoryStorage()
+    ):
+        self._storage = storage
+
+    def store(self, data: DataRow):
+        '''
+        Store data in the replay buffer.
+        '''
+        if not data:
+            raise ValueError("Data is required")
+        self._storage.add(data)
+
+    def store_batch(self, data_batch: List[DataRow]):
+        '''
+        Store batch of data in the replay buffer.
+        '''
+        if not data_batch:
+            logger.warning("Data batch is required")
+            return
+        self._storage.add_batch(data_batch)
+
+    def sample_task(self,
+                    sampler: TaskSampler = RandomTaskSample(),
+                    query_condition: QueryCondition = None,
+                    converter: Converter = DefaultConverter(),
+                    batch_size: int = 1000) -> List[T]:
+        '''
+        Sample Task from the replay buffer and convert to dataset row.
+        DefaultConverter return List[DataRow]
+        '''
+        sampled_task = sampler.sample_tasks(
+            self._storage, batch_size, query_condition)
+        return [converter.to_dataset_row(task_experiences) for task_experiences in sampled_task.values()]
+
+    def sample(self,
+               sampler: Sampler = RandomTaskSample(),
+               query_condition: QueryCondition = None,
+               converter: Converter = DefaultConverter(),
+               batch_size: int = 1000) -> List[T]:
+        '''
+        Sample data from the replay buffer and convert to dataset row.
+        DefaultConverter return List[DataRow]
+        '''
+        sampled_data = sampler.sample(
+            self._storage, batch_size, query_condition)
+        return converter.to_dataset_row(sampled_data)

chapter9/gaia-experience/AWorld/aworld/replay_buffer/event_replay_buffer.py

@@ -0,0 +1,218 @@
+import json
+import os.path
+import traceback
+from datetime import datetime
+from typing import Dict, List, Any
+
+from aworld import import_package
+from aworld.core.agent.base import is_agent_by_name
+from aworld.core.event.base import Message, Constants
+from aworld.logs.util import logger
+from aworld.replay_buffer.base import ReplayBuffer, DataRow, ExpMeta, Experience, InMemoryStorage, Storage
+from aworld.runners.state_manager import RuntimeStateManager, EventRuntimeStateManager
+from aworld.utils.serialized_util import to_serializable
+from aworld.utils.common import get_local_ip
+
+
+class EventReplayBuffer(ReplayBuffer):
+    '''
+    Event replay buffer for storing and sampling data.
+    Adds the ability to build DataRow from messages and export data to files.
+    '''
+    def __init__(
+        self,
+        storage: Storage = InMemoryStorage()
+    ):
+        super().__init__(storage)
+        self.task_agent_map = {}
+
+    async def get_trajectory(self, messages: List[Message], task_id: str, state_mng: RuntimeStateManager = None) -> List[Dict[str, Any]] | None:
+        if not messages:
+            return None
+        valid_agent_messages = await self._filter_replay_messages(messages, task_id)
+        if not valid_agent_messages:
+            return None
+        data_rows = []
+        try:
+            for msg in valid_agent_messages:
+                data_row = self.build_data_row_from_message(msg, state_mng)
+                if data_row:
+                    data_rows.append(data_row)
+            if not data_rows:
+                logger.warn(f"No valid agent messages found for task: {task_id}")
+                return None
+
+            self.store_batch(data_rows)
+            trajectory = [to_serializable(data_row) for data_row in data_rows]
+
+            self.export(data_rows, task_id)
+
+            return trajectory
+        except Exception as e:
+            logger.error(f"Failed to save trajectories: {str(e)}.{traceback.format_exc()}")
+            return None
+
+    async def _filter_replay_messages(self, messages: List[Message], task_id: str) -> List[Message]:
+        results = []
+        logger.info(f"Retrieving agent messages for task: {task_id}")
+        for message in messages:
+            if message.task_id != task_id or message.category != Constants.AGENT:
+                continue
+            sender = message.sender
+            receiver = message.receiver
+            if not sender or not receiver or not is_agent_by_name(receiver):
+                continue
+            agent_as_tool = message.headers.get("agent_as_tool", False)
+            if agent_as_tool:
+                continue
+            results.append(message)
+        return results
+
+    def build_data_row_from_message(self, message: Message, state_manager: RuntimeStateManager = None) -> DataRow:
+        '''
+        Build DataRow from a message.
+        
+        Args:
+            message (Dict): Message data containing necessary metadata and experience data
+            
+        Returns:
+            DataRow: The constructed data row
+            
+        Raises:
+            ValueError: When the message is missing required fields
+        '''
+        if not message:
+            raise ValueError("Message cannot be empty")
+
+        agent_id = message.receiver
+        task_id = message.context.task_id
+        task_name = message.context.get_task().name
+        pre_agent = message.sender
+        task_agent_id = f"{task_id}_{agent_id}"
+        if task_agent_id not in self.task_agent_map:
+            self.task_agent_map[task_agent_id] = 0
+        self.task_agent_map[task_agent_id] += 1
+        id = f"{task_agent_id}_{self.task_agent_map[task_agent_id]}"
+
+        # Build ExpMeta
+        exp_meta = ExpMeta(
+            task_id=task_id,
+            task_name=task_name,
+            agent_id=agent_id,
+            step=self.task_agent_map[task_agent_id],
+            execute_time=message.timestamp,
+            pre_agent=pre_agent
+        )
+
+        if not state_manager:
+            state_manager = EventRuntimeStateManager.instance()
+        observation = message.payload
+        node = state_manager._find_node(message.id)
+        if node is None or not node.results:
+            logger.error(f"Node result not found for message id: {message.id}, node: {node}")
+            return None
+        agent_results = []
+        for handle_result in node.results:
+            result = handle_result.result
+            if isinstance(result, Message) and isinstance(result.payload, list):
+                agent_results.extend(result.payload)
+        messages = self._get_llm_messages_from_memory(message)
+
+        # Build Experience
+        exp_data = Experience(
+            state=observation,
+            actions=agent_results,
+            messages=messages
+        )
+        
+        # Build and return DataRow
+        return DataRow(exp_meta=exp_meta, exp_data=exp_data, id=id)
+
+    def _get_llm_messages_from_memory(self, message: Message):
+        context = message.context
+        return context.context_info.get("llm_input", [])
+
+    def export(self, data_rows: List[DataRow], task_id: str) -> None:
+        '''
+        Export data rows to a specified file.
+        
+        Args:
+            data_rows (List[DataRow]): List of data rows to export
+            filepath (str): Path of the export file
+            
+        Raises:
+            ValueError: When the data rows list is empty or the file path is invalid
+        '''
+        enable_file_export = os.getenv("EXPORT_REPLAY_FILES", "false").lower() == "true"
+        enable_oss_export = os.getenv("EXPORT_REPLAY_TO_OSS", "false").lower() == "true"
+        if not enable_file_export and not enable_oss_export:
+            return
+
+        if not data_rows:
+            logger.warn("Data rows list cannot be empty")
+            return
+
+        try:
+            # Convert data rows to dictionary list
+            data_dicts = [to_serializable(data_row) for data_row in data_rows]
+
+            timestamp = datetime.now().strftime("%Y%m%d")
+            export_dir = os.getenv('REPLAY_EXPORT_DIRECTORY', None)
+            replay_dir = os.path.join(export_dir or "./trace_data", timestamp, get_local_ip(), "replays")
+            os.makedirs(replay_dir, exist_ok=True)
+            filepath = os.path.join(replay_dir, f"task_replay_{task_id}.json")
+
+            if enable_file_export:
+                logger.info(f"Exporting {len(data_rows)} data rows to {filepath}")
+                # Write to JSON file
+                with open(filepath, 'w', encoding='utf-8') as f:
+                    json.dump(data_dicts, f, ensure_ascii=False, indent=2)
+                logger.info(f"Successfully exported {len(data_rows)} data rows to {os.path.abspath(filepath)}")
+
+            if enable_oss_export:
+                logger.info(f"Exporting {len(data_rows)} data rows to oss")
+                self.export_to_oss(data_dicts, filepath)
+        except Exception as e:
+            logger.error(f"Failed to export replay datas: {e}")
+            raise
+
+    def export_to_oss(self, datas, filepath):
+        import_package("oss2")
+        import oss2
+
+        # Get OSS credentials from environment variables
+        access_key_id = os.getenv('OSS_ACCESS_KEY_ID')
+        access_key_secret = os.getenv('OSS_ACCESS_KEY_SECRET')
+        endpoint = os.getenv('OSS_ENDPOINT')
+        bucket_name = os.getenv('OSS_BUCKET_NAME')
+        bucket = None
+
+        if not all([access_key_id, access_key_secret, endpoint, bucket_name]):
+            logger.warn("Missing required OSS environment variables")
+            return
+        else:
+            try:
+                # Initialize OSS client
+                auth = oss2.Auth(access_key_id, access_key_secret)
+                bucket = oss2.Bucket(auth, endpoint, bucket_name)
+            except Exception as e:
+                logger.warn(
+                    f"Failed to initialize OSS client, endpoint: {endpoint}, bucket: {bucket_name}. Error: {str(e)}")
+                return
+
+        # Upload to OSS
+        try:
+            # Get the relative path
+            abs_path = os.path.abspath(filepath)
+            path_parts = abs_path.split(os.sep)
+            if len(path_parts) >= 4:
+                # Get the last 4 parts of the path
+                relative_path = os.sep.join(path_parts[-4:])
+                oss_key = relative_path
+            else:
+                oss_key = f"replay_buffer/{os.path.basename(filepath)}"
+            logger.info(f"Uploading replay datas to OSS: {oss_key}")
+            bucket.put_object_from_file(oss_key, filepath)
+            logger.info(f"Successfully uploaded {filepath} to OSS: {oss_key}")
+        except Exception as e:
+            logger.warn(f"Failed to upload {filepath} to OSS: {str(e)}")

chapter9/gaia-experience/AWorld/aworld/replay_buffer/processor.py

@@ -0,0 +1,190 @@
+# coding: utf-8
+"""
+processor.py
+Used to clean raw trace data into standard storage structure for reinforcement learning training.
+"""
+import json
+import os
+import datetime
+from typing import Any
+import threading
+
+from aworld.utils import import_package
+from aworld.replay_buffer.base import DataRow, Experience, ExpMeta
+from aworld.logs.util import logger
+from aworld.utils.common import get_local_ip
+
+
+class ReplayBufferExporter:
+    def __init__(self):
+        """Initialize ReplayBufferExporter instance"""
+        self._file_locks = {}
+        self._lock_dict_lock = threading.Lock()
+        self._task_output_paths = {}
+
+    def _get_file_lock(self, file_path):
+        """Get the lock for the specified file"""
+        with self._lock_dict_lock:
+            if file_path not in self._file_locks:
+                self._file_locks[file_path] = threading.Lock()
+            return self._file_locks[file_path]
+
+    def replay_buffer_exporter(self, spans: list[dict[str, Any]], output_dir: str):
+        """
+        Process spans, only process spans with 'step_execution_' prefix, and group by task_id to output to different files
+
+        Args:
+            spans: span data list
+            output_dir: output directory path
+        """
+        # Ensure output directory exists
+        import_package("oss2")
+        import oss2
+
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Get OSS credentials from environment variables
+        enable_oss_export = os.getenv("EXPORT_REPLAY_TRACE_TO_OSS", "false").lower() == "true"
+        access_key_id = os.getenv('OSS_ACCESS_KEY_ID')
+        access_key_secret = os.getenv('OSS_ACCESS_KEY_SECRET')
+        endpoint = os.getenv('OSS_ENDPOINT')
+        bucket_name = os.getenv('OSS_BUCKET_NAME')
+        bucket = None
+
+        if not all([access_key_id, access_key_secret, endpoint, bucket_name]):
+            enable_oss_export = False
+            logger.warn("Missing required OSS environment variables")
+        else:
+            try:
+                # Initialize OSS client
+                auth = oss2.Auth(access_key_id, access_key_secret)
+                bucket = oss2.Bucket(auth, endpoint, bucket_name)
+            except Exception as e:
+                enable_oss_export = False
+                logger.warn(f"Failed to initialize OSS client, endpoint: {endpoint}, bucket: {bucket_name}. Error: {str(e)}")
+
+        # Group by task_id
+        task_groups = {}
+
+        for span_data in spans:
+            # Only process spans with 'step_execution_' prefix
+            if not span_data['name'].startswith('step_execution_'):
+                continue
+
+            attr = span_data.get('attributes', {})
+            exp_id = attr.get('exp_id')
+            task_id = attr.get('task_id', '')
+
+            if not exp_id or not task_id:
+                continue
+
+            if task_id not in task_groups:
+                task_groups[task_id] = {}
+
+            if exp_id not in task_groups[task_id]:
+                task_groups[task_id][exp_id] = {
+                    'exp_meta': None,
+                    'exp_data': None
+                }
+
+            # Process step_execution span
+            task_name = attr.get('task_name', '')
+            agent_id = attr.get('agent_id', '')
+            step = attr.get('step', 0)
+            execute_time = float(span_data.get('start_time', 0).split('.')[0].replace(' ', '').replace('-', '').replace(':', ''))
+
+            observation = {}
+            action = []
+            messages = []
+            pre_agent = None
+            if 'observation' in attr:
+                try:
+                    observation = json.loads(attr['observation'])
+                except:
+                    observation = attr['observation']
+
+            if 'actions' in attr:
+                try:
+                    action = json.loads(attr['actions'])
+                except:
+                    action = attr['actions']
+
+            if 'messages' in attr:
+                try:
+                    messages = json.loads(attr['messages'])
+                except:
+                    messages = attr['messages']
+
+            pre_agent = attr.get('pre_agent', '')
+            reward = attr.get('reward', 0.0)
+            adv = attr.get('adv_t', 0.0)
+            v = attr.get('v_t', 0.0)
+
+            exp_meta = ExpMeta(task_id, task_name, agent_id, step, execute_time, pre_agent)
+            exp_data = Experience(observation, action, reward, adv, v, messages)
+
+            task_groups[task_id][exp_id]['exp_meta'] = exp_meta
+            task_groups[task_id][exp_id]['exp_data'] = exp_data
+
+        # Process data for each task_id
+        for task_id, exp_groups in task_groups.items():
+            # Merge data and generate final Experience object
+            data_rows = []
+
+            # Read existing data (if any)
+            output_path = self._task_output_paths.get(task_id)
+            if not output_path:
+                timestamp = datetime.datetime.now().strftime("%Y%m%d")
+                replay_dir = os.path.join(output_dir or "./trace_data", timestamp, get_local_ip(), "replays")
+                replay_dataset_path = os.getenv("REPLAY_TRACE_DATASET_PATH", replay_dir)
+                export_dir = os.path.abspath(replay_dataset_path)
+                os.makedirs(export_dir, exist_ok=True)
+                output_path = os.path.join(export_dir, f"task_replay_{task_id}.json")
+                self._task_output_paths[task_id] = output_path
+
+            # Use thread lock to protect read and write operations
+            file_lock = self._get_file_lock(output_path)
+            with file_lock:
+                if os.path.exists(output_path):
+                    try:
+                        with open(output_path, 'r', encoding='utf-8') as f:
+                            existing_data = json.load(f)
+                            data_rows.extend([DataRow(
+                                ExpMeta(**row['exp_meta']),
+                                Experience(**row['exp_data']),
+                                row['id']
+                            ) for row in existing_data])
+                    except Exception as e:
+                        print(f"Failed to read existing file {output_path}: {str(e)}")
+
+                # Add new data
+                for exp_id, group in exp_groups.items():
+                    if group['exp_meta'] and group['exp_data']:
+                        row = DataRow(group['exp_meta'], group['exp_data'], exp_id)
+                        data_rows.append(row)
+
+                # Sort by execute_time
+                data_rows.sort(key=lambda x: x.exp_meta.execute_time)
+
+                # Export to json
+                with open(output_path, 'w', encoding='utf-8') as f:
+                    json.dump([row.to_dict() for row in data_rows], f, ensure_ascii=False, indent=2)
+                logger.info(f"Processing completed, exported {len(data_rows)} experiences to {output_path}")
+
+                if enable_oss_export:
+                    # Upload to OSS
+                    try:
+                        # Get the relative path
+                        abs_path = os.path.abspath(output_path)
+                        path_parts = abs_path.split(os.sep)
+                        if len(path_parts) >= 4:
+                            # Get the last 4 parts of the path
+                            relative_path = os.sep.join(path_parts[-4:])
+                            oss_key = relative_path
+                        else:
+                            oss_key = f"replay_buffer/{os.path.basename(output_path)}"
+                        bucket.put_object_from_file(oss_key, output_path)
+                        logger.info(f"Successfully uploaded {output_path} to OSS: {oss_key}")
+                    except Exception as e:
+                        logger.warn(f"Failed to upload {output_path} to OSS: {str(e)}")
+

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

@@ -0,0 +1,228 @@
+from typing import Any, List, TypeVar, Union, Literal, TypedDict, Dict
+
+DataRow = TypeVar('DataRow')
+
+
+class BaseCondition(TypedDict):
+    field: str
+    value: Any
+    op: Literal[
+        'eq', 'ne', 'gt', 'gte', 'lt', 'lte',
+        'in', 'not_in', 'like', 'not_like',
+        'is_null', 'is_not_null'
+    ]
+
+
+class LogicalCondition(TypedDict):
+    and_: List['QueryCondition']
+    or_: List['QueryCondition']
+
+
+QueryCondition = Union[BaseCondition, LogicalCondition]
+
+
+class QueryBuilder:
+    '''
+    Query builder for replay buffer. result example:
+    {
+        "and": [
+            {"field": "field1", "value": "value1", "op": "eq"}, 
+            {"or": [{"field": "field2", "value": "value2", "op": "eq"}, {"field": "field3", "value": "value3", "op": "eq"}]}
+        ]
+    }
+    '''
+
+    def __init__(self) -> None:
+        self.conditions: List[Dict[str, any]] = []
+        self.logical_ops: List[str] = []
+
+    def eq(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "eq"})
+        return self
+
+    def ne(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "ne"})
+        return self
+
+    def gt(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "gt"})
+        return self
+
+    def gte(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "gte"})
+        return self
+
+    def lt(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "lt"})
+        return self
+
+    def lte(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "lte"})
+        return self
+
+    def in_(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "in"})
+        return self
+
+    def not_in(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append(
+            {"field": field, "value": value, "op": "not_in"})
+        return self
+
+    def like(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "value": value, "op": "like"})
+        return self
+
+    def not_like(self, field: str, value: any) -> 'QueryBuilder':
+        self.conditions.append(
+            {"field": field, "value": value, "op": "not_like"})
+        return self
+
+    def is_null(self, field: str) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "op": "is_null"})
+        return self
+
+    def is_not_null(self, field: str) -> 'QueryBuilder':
+        self.conditions.append({"field": field, "op": "is_not_null"})
+        return self
+
+    def and_(self) -> 'QueryBuilder':
+        self.logical_ops.append("and_")
+        return self
+
+    def or_(self) -> 'QueryBuilder':
+        self.logical_ops.append("or_")
+        return self
+
+    def nested(self, builder: 'QueryBuilder') -> 'QueryBuilder':
+        self.conditions.append({"nested": builder.build()})
+        return self
+
+    def build(self) -> QueryCondition:
+        conditions = self.conditions  # all conditions（including nested）
+        operators = self.logical_ops
+
+        # Validate condition and operator counts (n conditions need n-1 operators)
+        if len(operators) != len(conditions) - 1:
+            raise ValueError("Mismatch between condition and operator counts")
+
+        # Use stack to handle operator precedence (simplified version supporting and/or)
+        stack: List[Union[Dict[str, any], str]] = []
+
+        for i, item in enumerate(conditions):
+            if i == 0:
+                # First element goes directly to stack (condition or nested)
+                stack.append(item)
+                continue
+
+            # Pop stack top as left operand
+            left = stack.pop()
+            op = operators[i-1]       # Current operator (and/or)
+            right = item              # Right operand (current condition)
+
+            # Build logical expression: {op: [left, right]}
+            expr = {op: [left, right]}
+            # Push result back to stack for further operations
+            stack.append(expr)
+
+        # Process nested conditions (recursive unfolding)
+        def process_nested(cond: any) -> any:
+            if isinstance(cond, dict):
+                if "nested" in cond:
+                    # Recursively process sub-conditions
+                    return process_nested(cond["nested"])
+                # Recursively process child elements
+                return {k: process_nested(v) for k, v in cond.items()}
+            elif isinstance(cond, list):
+                return [process_nested(item) for item in cond]
+            return cond
+
+        # Final result: only one element left in stack, return after processing nested
+        result = stack[0] if stack else None
+        return process_nested(result) if result else None
+
+
+class QueryFilter:
+    '''
+    Query filter for replay buffer.
+    '''
+
+    def __init__(self, query_condition: QueryCondition) -> None:
+        self.query_condition = query_condition
+
+    def _get_field_value(self, row: DataRow, field: str) -> Any:
+        '''
+        Get field value from row.
+        '''
+        obj = row
+        for part in field.split('.'):
+            obj = getattr(obj, part, None)
+            if obj is None:
+                break
+        return obj
+
+    def _do_check(self, row: DataRow, condition: QueryCondition) -> bool:
+        """
+        check if row match condition
+        """
+        if condition is None:
+            return True
+        if "field" in condition and "op" in condition:
+            field_val = self._get_field_value(row, condition["field"])
+            op = condition["op"]
+            target_val = condition["value"]
+
+            if op == "eq":
+                return field_val == target_val
+            if op == "ne":
+                return field_val != target_val
+            if op == "gt":
+                return field_val > target_val
+            if op == "gte":
+                return field_val >= target_val
+            if op == "lt":
+                return field_val < target_val
+            if op == "lte":
+                return field_val <= target_val
+            if op == "in":
+                return field_val in target_val
+            if op == "not_in":
+                return field_val not in target_val
+            if op == "like":
+                return target_val in field_val
+            if op == "not_like":
+                return target_val not in field_val
+            if op == "is_null":
+                return field_val is None
+            if op == "is_not_null":
+                return field_val is not None
+
+            return False
+
+        elif "and_" in condition or "or_" in condition:
+            if "and_" in condition:
+                return all(self._do_check(row, c) for c in condition["and_"])
+            if "or_" in condition:
+                return any(self._do_check(row, c) for c in condition["or_"])
+            return False
+
+        return False
+
+    def check_condition(self, row: DataRow) -> bool:
+        """
+        check if row match condition
+        """
+        return self._do_check(row, self.query_condition)
+
+    def filter(self, rows: List[DataRow]) -> List[DataRow]:
+        """filter rows by condition
+        Args:
+            rows (List[DataRow]): List of rows to filter.
+            query_condition (QueryCondition): Query condition.
+        Returns:
+            List[DataRow]: List of rows that match the condition.
+        """
+        condition = self.query_condition
+        if not condition:
+            return rows
+        return [row for row in rows if self.check_condition(row)]

chapter9/gaia-experience/AWorld/aworld/replay_buffer/storage/__init__.py

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

chapter9/gaia-experience/AWorld/aworld/replay_buffer/storage/multi_proc_mem.py

@@ -0,0 +1,155 @@
+import multiprocessing
+import traceback
+import pickle
+from typing import Dict, List
+from aworld.replay_buffer.base import Storage, DataRow
+from aworld.replay_buffer.query_filter import QueryCondition, QueryFilter
+from aworld.logs.util import logger
+
+
+class MultiProcMemoryStorage(Storage):
+
+    """
+    Memory storage for multi-process.
+    """
+
+    def __init__(self,
+                 data_dict: Dict[str, str],
+                 fifo_queue: List[str],
+                 lock: multiprocessing.Lock,
+                 max_capacity: int = 10000):
+        self._data: Dict[str, str] = data_dict
+        self._fifo_queue = fifo_queue
+        self._max_capacity = max_capacity
+        self._lock = lock
+
+    def _save_to_shared_memory(self, data, task_id):
+        serialized_data = pickle.dumps(data)
+        try:
+            if task_id not in self._data or not self._data[task_id]:
+                shm = multiprocessing.shared_memory.SharedMemory(
+                    create=True, size=len(serialized_data))
+                shm.buf[:len(serialized_data)] = serialized_data
+                self._data[task_id] = shm.name
+                shm.close()
+                return
+            shm = multiprocessing.shared_memory.SharedMemory(
+                name=self._data[task_id], create=False)
+            if len(serialized_data) > shm.size:
+                shm.close()
+                shm.unlink()
+                shm = multiprocessing.shared_memory.SharedMemory(
+                    create=True, size=len(serialized_data))
+                shm.buf[:len(serialized_data)] = serialized_data
+                self._data[task_id] = shm.name
+            else:
+                shm.buf[:len(serialized_data)] = serialized_data
+        except FileNotFoundError:
+            shm = multiprocessing.shared_memory.SharedMemory(
+                create=True, size=len(serialized_data))
+            shm.buf[:len(serialized_data)] = serialized_data
+            self._data[task_id] = shm.name
+        shm.close()
+
+    def _load_from_shared_memory(self, task_id):
+        try:
+            if task_id not in self._data or not self._data[task_id]:
+                return []
+            try:
+                multiprocessing.shared_memory.SharedMemory(
+                    name=self._data[task_id], create=False)
+            except FileNotFoundError:
+                return []
+            shm = multiprocessing.shared_memory.SharedMemory(
+                name=self._data[task_id])
+            data = pickle.loads(shm.buf.tobytes())
+            shm.close()
+            return data
+        except Exception as e:
+            stack_trace = traceback.format_exc()
+            logger.error(
+                f"_load_from_shared_memory error: {e}\nStack trace:\n{stack_trace}")
+            return []
+
+    def _delete_from_shared_memory(self, task_id):
+        try:
+            if task_id not in self._data or not self._data[task_id]:
+                return
+            shm = multiprocessing.shared_memory.SharedMemory(
+                name=self._data[task_id])
+            shm.close()
+            shm.unlink()
+            del self._data[task_id]
+        except FileNotFoundError:
+            pass
+
+    def add(self, data: DataRow):
+        if not data:
+            raise ValueError("Data is required")
+        if not data.exp_meta:
+            raise ValueError("exp_meta is required")
+
+        with self._lock:
+            current_size = sum(len(self._load_from_shared_memory(task_id))
+                               for task_id in self._data.keys())
+            while current_size >= self._max_capacity and self._fifo_queue:
+                oldest_task_id = self._fifo_queue.pop(0)
+                if oldest_task_id in self._data.keys():
+                    current_size -= len(self._load_from_shared_memory(oldest_task_id))
+                    self._delete_from_shared_memory(oldest_task_id)
+
+            task_id = data.exp_meta.task_id
+            existing_data = self._load_from_shared_memory(task_id)
+            existing_data.append(data)
+            self._save_to_shared_memory(existing_data, task_id)
+            self._fifo_queue.append(task_id)
+
+    def add_batch(self, data_batch: List[DataRow]):
+        with self._lock:
+            for data in data_batch:
+                self.add(data)
+
+    def size(self, query_condition: QueryCondition = None) -> int:
+        with self._lock:
+            return len(self._get_all_without_lock(query_condition))
+
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        with self._lock:
+            if page < 1:
+                raise ValueError("Page must be greater than 0")
+            if page_size < 1:
+                raise ValueError("Page size must be greater than 0")
+            all_data = self._get_all_without_lock(query_condition)
+            start_index = (page - 1) * page_size
+            end_index = start_index + page_size
+            return all_data[start_index:end_index]
+
+    def _get_all_without_lock(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        all_data = []
+        query_filter = None
+        if query_condition:
+            query_filter = QueryFilter(query_condition)
+        for task_id in self._data.keys():
+            local_data = self._load_from_shared_memory(task_id)
+            if query_filter:
+                all_data.extend(query_filter.filter(local_data))
+            else:
+                all_data.extend(local_data)
+        return all_data
+
+    def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        with self._lock:
+            return self._get_all_without_lock(query_condition)
+
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        with self._lock:
+            if task_id in self._data.keys():
+                return self._load_from_shared_memory(task_id)
+
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        with self._lock:
+            result = {}
+            for task_id in task_ids:
+                if task_id in self._data.keys():
+                    result[task_id] = self._load_from_shared_memory(task_id)
+            return result

chapter9/gaia-experience/AWorld/aworld/replay_buffer/storage/odps.py

@@ -0,0 +1,223 @@
+import json
+from pydantic import parse_obj_as
+from typing import Any, List, Dict
+from aworld.replay_buffer.base import Storage, DataRow, ExpMeta, Experience
+from aworld.replay_buffer.query_filter import QueryCondition, QueryBuilder
+from aworld.core.common import Observation, ActionModel
+from aworld.logs.util import logger
+from aworld.utils.import_package import import_package
+import_package("odps")  # noqa
+from odps import ODPS  # noqa
+from odps.models.record import Record  # noqa
+
+
+class OdpsSQLBuilder:
+    ''' Example:
+            query_condition = QueryBuilder().eq("field1", "value1").and_().eq("field2", "value2")
+            sql_builder = OdpsSQLBuilder(query_condition)
+            sql = sql_builder.build_sql()
+            print(sql)  # 输出: "field1 = 'value1' AND field2 = 'value2'"
+    '''
+
+    def __init__(self, query_condition: QueryCondition):
+        self.query_condition = query_condition
+
+    def _build_condition(self, condition: QueryCondition) -> str:
+        if condition is None:
+            return ""
+
+        if "field" in condition and "op" in condition:
+            field = condition["field"].split('.')[-1]
+            op = condition["op"]
+            value = condition.get("value")
+
+            if op == "eq":
+                return f"{field} = {self._format_value(value)}"
+            elif op == "ne":
+                return f"{field} != {self._format_value(value)}"
+            elif op == "gt":
+                return f"{field} > {self._format_value(value)}"
+            elif op == "gte":
+                return f"{field} >= {self._format_value(value)}"
+            elif op == "lt":
+                return f"{field} < {self._format_value(value)}"
+            elif op == "lte":
+                return f"{field} <= {self._format_value(value)}"
+            elif op == "in":
+                return f"{field} IN ({self._format_value(value)})"
+            elif op == "not_in":
+                return f"{field} NOT IN ({self._format_value(value)})"
+            elif op == "like":
+                return f"{field} LIKE '{value}'"
+            elif op == "not_like":
+                return f"{field} NOT LIKE '{value}'"
+            elif op == "is_null":
+                return f"{field} IS NULL"
+            elif op == "is_not_null":
+                return f"{field} IS NOT NULL"
+
+        elif "and_" in condition:
+            return f"({' AND '.join(self._build_condition(c) for c in condition['and_'])})"
+        elif "or_" in condition:
+            return f"({' OR '.join(self._build_condition(c) for c in condition['or_'])})"
+
+        return ""
+
+    def _format_value(self, value: Any) -> str:
+        if isinstance(value, str):
+            return f"'{value}'"
+        elif isinstance(value, (list, tuple)):
+            return ", ".join(self._format_value(v) for v in value)
+        return str(value)
+
+    def build_sql(self) -> str:
+        if not self.query_condition:
+            return ""
+        return self._build_condition(self.query_condition)
+
+
+class OdpsStorage(Storage):
+    '''
+        Aliyun ODPS storage.
+        Table schema:
+            id: int
+            task_id: string
+            task_name: string
+            agent_id: string
+            step: int
+            execute_time: string
+            state: string
+            actions: string
+            reward_t: string
+            adv_t: string
+            v_t: string
+    '''
+
+    def __init__(self, table_name: str, project: str, endpoint: str, access_id: str, access_key: str, **kwargs):
+        self.table_name = table_name
+        self.project = project
+        self.endpoint = endpoint
+        self.access_id = access_id
+        self.access_key = access_key
+        self.kwargs = kwargs
+        self._init_odps()
+
+    def _init_odps(self):
+        self.odps = ODPS(self.access_id, self.access_key,
+                         self.project, self.endpoint)
+
+    def _get_table(self):
+        return self.odps.get_table(self.table_name)
+
+    def _convert_row_to_record(self, row: DataRow) -> Record:
+        table = self._get_table()
+        record = table.new_record()
+        record["id"] = row.id
+        record["task_id"] = row.exp_meta.task_id
+        record["task_name"] = row.exp_meta.task_name
+        record["agent_id"] = row.exp_meta.agent_id
+        record["step"] = row.exp_meta.step
+        record["execute_time"] = row.exp_meta.execute_time
+        if row.exp_data.state:
+            record["state"] = row.exp_data.state.model_dump_json()
+        if row.exp_data.actions:
+            record["actions"] = "[" + ", ".join(action.model_dump_json()
+                                                for action in row.exp_data.actions) + "]"
+        if row.exp_data.reward_t:
+            record["reward_t"] = row.exp_data.reward_t
+        if row.exp_data.adv_t:
+            record["adv_t"] = row.exp_data.adv_t
+        if row.exp_data.v_t:
+            record["v_t"] = row.exp_data.v_t
+        return record
+
+    def _convert_record_to_row(self, record: Record) -> DataRow:
+        return DataRow(
+            id=record.id,
+            exp_meta=ExpMeta(
+                task_id=record['task_id'],
+                task_name=record['task_name'],
+                agent_id=record['agent_id'],
+                step=record['step'],
+                execute_time=record['execute_time'],
+                pre_agent=record['pre_agent'] if 'pre_agent' in record else None
+            ),
+            exp_data=Experience(
+                state=parse_obj_as(Observation, json.loads(record['state'])),
+                actions=[parse_obj_as(ActionModel, item)
+                         for item in json.loads(record['actions'])],
+                reward_t=record['reward_t'] if 'reward_t' in record else None,
+                adv_t=record['adv_t'] if 'adv_t' in record else None,
+                v_t=record['v_t'] if 'v_t' in record else None,
+            )
+        )
+
+    def _build_paginated_sql(self, page: int = None, page_size: int = None):
+        if page and page_size:
+            offset = (page - 1) * page_size
+            limit = page_size
+            return f" LIMIT {offset}, {limit}"
+        return ""
+
+    def _build_sql(self, query_condition: QueryCondition, page: int = None, page_size: int = None):
+        if not query_condition:
+            return f"SELECT * FROM {self.table_name}" + self._build_paginated_sql(page, page_size)
+        where_builder = OdpsSQLBuilder(query_condition)
+        sql = f"SELECT * FROM {self.table_name} WHERE {where_builder.build_sql()}" + self._build_paginated_sql(page,
+                                                                                                               page_size)
+        return sql
+
+    def _build_count_sql(self, query_condition: QueryCondition):
+        if not query_condition:
+            return f"SELECT count(1) as count FROM {self.table_name}"
+        where_builder = OdpsSQLBuilder(query_condition)
+        sql = f"SELECT count(1) as count FROM {self.table_name} WHERE {where_builder.build_sql()}"
+        return sql
+
+    def add(self, row: DataRow):
+        record = self._convert_row_to_record(row)
+        self.odps.write_table(self.table_name, [record])
+
+    def add_batch(self, rows: list[DataRow]):
+        records = [self._convert_row_to_record(row) for row in rows]
+        self.odps.write_table(self.table_name, records)
+
+    def size(self, query_condition: QueryCondition = None) -> int:
+        sql = self._build_count_sql(query_condition)
+        with self.odps.execute_sql(sql).open_reader() as reader:
+            return reader[0]["count"]
+
+    def get_all(self, query_condition: QueryCondition = None) -> list[DataRow]:
+        sql = self._build_sql(query_condition)
+        logger.info(f"get_all sql: {sql}")
+        with self.odps.execute_sql(sql).open_reader(tunnel=True) as reader:
+            rows = []
+            for record in reader:
+                rows.append(self._convert_record_to_row(record))
+            return rows
+
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        sql = self._build_sql(query_condition, page, page_size)
+        logger.info(f"get_paginated sql: {sql}")
+        with self.odps.execute_sql(sql).open_reader(tunnel=True) as reader:
+            rows = []
+            for record in reader:
+                rows.append(self._convert_record_to_row(record))
+            return rows
+
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        query_condition = QueryBuilder().eq("task_id", task_id).build()
+        return self.get_all(query_condition)
+
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        query_condition = QueryBuilder().in_("task_id", task_ids).build()
+        sql = self._build_sql(query_condition)
+        logger.info(f"get_bacth_by_task_ids sql: {sql}")
+        result = {}
+        with self.odps.execute_sql(sql).open_reader(tunnel=True) as reader:
+            for record in reader:
+                row = self._convert_record_to_row(record)
+                if row.exp_meta.task_id not in result:
+                    result[row.exp_meta.task_id] = []
+                result[row.exp_meta.task_id].append(row)
+        return result

chapter9/gaia-experience/AWorld/aworld/replay_buffer/storage/redis.py

@@ -0,0 +1,262 @@
+import json
+from typing import Dict, List
+from aworld.replay_buffer.base import Storage, DataRow, ExpMeta, Experience
+from aworld.logs.util import logger
+from aworld.utils.import_package import import_package
+from aworld.replay_buffer.query_filter import QueryCondition, QueryBuilder
+from aworld.core.common import Observation, ActionModel
+import_package("redis")  # noqa
+from redis import Redis  # noqa
+from redis.commands.json.path import Path  # noqa
+import redis.commands.search.aggregation as aggregations  # noqa
+import redis.commands.search.reducers as reducers  # noqa
+from redis.commands.search.field import TextField, NumericField, TagField  # noqa
+from redis.commands.search.index_definition import IndexDefinition, IndexType  # noqa
+from redis.commands.search.query import Query  # noqa
+import redis.exceptions  # noqa
+
+
+class RedisSearchQueryBuilder:
+    """
+    Build redis search query from query condition
+    """
+
+    def __init__(self, query_condition: QueryCondition):
+        self.query_condition = query_condition
+
+    def _build_condition(self, condition: QueryCondition) -> str:
+        if condition is None:
+            return ""
+
+        if "field" in condition and "op" in condition:
+            field = condition["field"].split('.')[-1]
+            op = condition["op"]
+            value = condition.get("value")
+
+            if op == "eq":
+                return f"@{field}:{{{value}}}"
+            elif op == "ne":
+                return f"-@{field}:{{{value}}}"
+            elif op == "gt":
+                return f"@{field}:[{value} +inf]"
+            elif op == "gte":
+                return f"@{field}:[{value} +inf]"
+            elif op == "lt":
+                return f"@{field}:[-inf {value}]"
+            elif op == "lte":
+                return f"@{field}:[-inf {value}]"
+            elif op == "in":
+                return f"@{field}:{{{'|'.join(str(v) for v in value)}}}"
+            elif op == "not_in":
+                return f"-@{field}:{{{'|'.join(str(v) for v in value)}}}"
+            elif op == "like":
+                return f"@{field}:*{value}*"
+            elif op == "not_like":
+                return f"-@{field}:*{value}*"
+            elif op == "is_null":
+                return f"-@{field}:*"
+            elif op == "is_not_null":
+                return f"@{field}:*"
+
+        elif "and_" in condition:
+            conditions = [self._build_condition(c) for c in condition["and_"]]
+            return " ".join(conditions)
+        elif "or_" in condition:
+            conditions = [self._build_condition(c) for c in condition["or_"]]
+            return f"({'|'.join(conditions)})"
+
+        return ""
+
+    def build(self) -> Query:
+        query_str = self._build_condition(self.query_condition)
+        logger.info(f"redis search query: {query_str}")
+        return Query(query_str)
+
+
+class RedisStorage(Storage):
+    def __init__(self,
+                 host: str = 'localhost',
+                 port: int = 6379,
+                 db: int = 0,
+                 password: str = None,
+                 key_prefix: str = 'AWORLD:RB:',
+                 index_name: str = 'idx:AWORLD:RB',
+                 recreate_idx_if_exists=False):
+        self._redis = Redis(host=host, port=port, db=db, password=password)
+        self._key_prefix = key_prefix
+        self._index_name = index_name
+        self._recreate_idx_if_exists = recreate_idx_if_exists
+        self._create_index()
+
+    def _create_index(self):
+        try:
+            existing_indices = self._redis.execute_command('FT._LIST')
+            if self._index_name.encode('utf-8') in existing_indices:
+                logger.info(f"Index {self._index_name} already exists")
+                if self._recreate_idx_if_exists:
+                    self._redis.ft(self._index_name).dropindex()
+                    logger.info(f"Index {self._index_name} dropped")
+                else:
+                    return
+            self._redis.ft(self._index_name).create_index(
+                (
+                    TagField("id"),
+                    TagField("task_id"),
+                    TextField("task_name"),
+                    TagField("agent_id"),
+                    NumericField("step"),
+                    NumericField("execute_time"),
+                    TagField("pre_agent")
+                ),
+                definition=IndexDefinition(
+                    prefix=[self._key_prefix], index_type=IndexType.HASH)
+            )
+        except redis.exceptions.ResponseError as e:
+            logger.error(f"Create index {self._index_name} failed. {e}")
+
+    def _get_object_key(self, key: str) -> str:
+        return f"{self._key_prefix}{key}"
+
+    def _serialize_to_str(self, value) -> str:
+        if str is None:
+            return ""
+        if isinstance(value, (int, float)):
+            return str(value)
+        return str(value) if value is not None else ""
+
+    def _serialize(self, data: DataRow) -> Dict[str, str]:
+        dict_data = {
+            'id': data.id,
+            'task_id': data.exp_meta.task_id,
+            'task_name': data.exp_meta.task_name,
+            'agent_id': data.exp_meta.agent_id,
+            'step': data.exp_meta.step,
+            'execute_time': data.exp_meta.execute_time,
+            'pre_agent': data.exp_meta.pre_agent,
+            'state': data.exp_data.state.model_dump_json(),
+            'actions': "[" + ", ".join(action.model_dump_json()
+                                       for action in data.exp_data.actions) + "]",
+            'reward_t': data.exp_data.reward_t,
+            'adv_t': data.exp_data.adv_t,
+            'v_t': data.exp_data.v_t
+        }
+        return {k: self._serialize_to_str(v) for k, v in dict_data.items()}
+
+    def _deserialize(self, data: Dict) -> DataRow:
+        if not data:
+            return None
+        return DataRow(
+            id=data.get('id'),
+            exp_meta=ExpMeta(
+                task_id=data.get('task_id'),
+                task_name=data.get('task_name'),
+                agent_id=data.get('agent_id'),
+                step=int(data.get('step', 0)),
+                execute_time=float(data.get('execute_time', 0)),
+                pre_agent=data.get('pre_agent')
+            ),
+            exp_data=Experience(
+                state=Observation.model_validate_json(data.get('state', '{}')),
+                actions=[ActionModel.model_validate_json(json.dumps(action))
+                         for action in json.loads(data.get('actions', '[]'))],
+                reward_t=float(data.get('reward_t', 0)) if data.get(
+                    'reward_t') is not '' else None,
+                adv_t=float(data.get('adv_t', 0)) if data.get(
+                    'adv_t') is not '' else None,
+                v_t=float(data.get('v_t', 0)) if data.get(
+                    'v_t') is not '' else None
+            )
+        )
+
+    def add(self, data: DataRow):
+        key = self._get_object_key(data.id)
+        self._redis.hset(key, mapping=self._serialize(data))
+
+    def add_batch(self, data_batch: List[DataRow]):
+        pipeline = self._redis.pipeline()
+        for data in data_batch:
+            if not data or not data.exp_meta:
+                continue
+            key = self._get_object_key(data.id)
+            pipeline.hset(key, mapping=self._serialize(data))
+        pipeline.execute()
+
+    def search(self, key: str, value: str) -> DataRow:
+        result = self._redis.ft(self._index_name).search(
+            Query(f"@{key}:{{{value}}}"))
+        logger.info(f"Search result: {result}")
+
+    def size(self, query_condition: QueryCondition = None) -> int:
+        '''
+        Get the size of the storage.
+        Returns:
+            int: Size of the storage.
+        '''
+        if not query_condition:
+            return self._redis.ft(self._index_name).info()['num_docs']
+        query_builder = RedisSearchQueryBuilder(query_condition)
+        query = query_builder.build()
+        return self._redis.ft(self._index_name).search(query).total
+
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        '''
+        Get paginated data from the storage.
+        Args:
+            page (int): Page number.
+            page_size (int): Number of data per page.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+        if not query_condition:
+            result = self._redis.ft(self._index_name).search(
+                Query("*").paging(page, page_size))
+        else:
+            query_builder = RedisSearchQueryBuilder(query_condition)
+            query = query_builder.build().paging(page, page_size)
+            result = self._redis.ft(self._index_name).search(query)
+        return [self._deserialize(doc.__dict__) for doc in result.docs]
+
+    def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        '''
+        Get all data from the storage.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+        if not query_condition:
+            result = self._redis.ft(self._index_name).search(Query("*"))
+        else:
+            query_builder = RedisSearchQueryBuilder(query_condition)
+            query = query_builder.build()
+            result = self._redis.ft(self._index_name).search(query)
+        return [self._deserialize(doc.__dict__) for doc in result.docs]
+
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        '''
+        Get data by task_id from the storage.
+        Args:
+            task_id (str): Task id.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+        query_condition = QueryBuilder().eq("task_id", task_id).build()
+        return self.get_all(query_condition)
+
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        '''
+        Get data by task_ids from the storage.
+        Args:
+            task_ids (List[str]): List of task ids.
+        Returns:
+            Dict[str, List[DataRow]]: Dict of task id and list of data.
+        '''
+        query_condition = QueryBuilder().in_("task_id", task_ids).build()
+        result = self.get_all(query_condition)
+        return {task_id: [data for data in result if data.exp_meta.task_id == task_id] for task_id in task_ids}
+
+    def clear(self):
+        '''
+        Clear the storage.
+        '''
+        keys = self._redis.keys(f"{self._key_prefix}*")
+        if keys:
+            self._redis.delete(*keys)