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

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

@@ -0,0 +1,98 @@
+# Checkpoint Module
+
+## Overview
+The Checkpoint module provides a robust and extensible framework for managing state snapshots (checkpoints) in Python applications. It is designed for scenarios where you need to persist, restore, and version the state of a process, session, or task.
+
+```mermaid
+sequenceDiagram
+    participant Application
+    participant CheckpointRepository
+    participant BackendStorage
+
+    Note over Application,BackendStorage: Create and store a checkpoint
+    %% Create and store a checkpoint
+    Application->>CheckpointRepository: create checkpoint
+    CheckpointRepository->>BackendStorage: put(checkpoint)
+    BackendStorage-->>CheckpointRepository: success
+    CheckpointRepository-->>Application: ack
+    
+    Note over Application,BackendStorage: Retrieve the latest checkpoint by session
+
+    %% Retrieve the latest checkpoint by session
+    Application->>CheckpointRepository: get checkpoint by session_id
+    CheckpointRepository->>BackendStorage: get_by_session(session_id)
+    BackendStorage-->>CheckpointRepository: Checkpoint
+    CheckpointRepository-->>Application: Checkpoint
+
+```
+
+## Key Features
+
+- **Structured Data Model**: Uses Pydantic's `BaseModel` for strong typing and validation of checkpoint data and metadata.
+- **Versioning Support**: Built-in version management utilities for checkpoint evolution and comparison.
+- **Extensible Repository Pattern**: Abstract base class (`BaseCheckpointRepository`) defines a standard interface for checkpoint storage, supporting both synchronous and asynchronous operations.
+- **In-Memory Implementation**: Includes a simple, ready-to-use in-memory repository for development and testing.
+- **Utility Functions**: Helper methods for creating, copying, and managing checkpoints.
+
+## Data Structures
+
+```mermaid
+classDiagram
+    class Application {
+        +CheckpointRepository repo
+        +create_checkpoint()
+        +get_checkpoint_by_session()
+    }
+    class CheckpointRepository {
+        +put(checkpoint)
+        +get_by_session(session_id)
+        +delete_by_session(session_id)
+        -BackendStorage backend
+    }
+    class BackendStorage {
+        +put(checkpoint)
+        +get_by_session(session_id)
+        +delete_by_session(session_id)
+    }
+    Application --> CheckpointRepository : uses
+    CheckpointRepository --> BackendStorage : delegates
+    class Checkpoint {
+        +id: str
+        +ts: str
+        +metadata: CheckpointMetadata
+        +values: dict
+        +version: int
+        +parent_id: str
+        +namespace: str
+    }
+    class CheckpointMetadata {
+        +session_id: str
+        +task_id: str
+    }
+    Checkpoint o-- CheckpointMetadata
+    CheckpointRepository o-- Checkpoint
+    BackendStorage o-- Checkpoint
+```
+
+
+## Usage Example
+
+```python
+from aworld.checkpoint import (
+    Checkpoint, CheckpointMetadata, empty_checkpoint, create_checkpoint, InMemoryCheckpointRepository
+)
+
+# Create a new checkpoint
+metadata = CheckpointMetadata(session_id="session-123", task_id="task-456")
+values = {"step": 1, "score": 100}
+checkpoint = create_checkpoint(values=values, metadata=metadata)
+
+# Store and retrieve using the in-memory repository
+repo = InMemoryCheckpointRepository()
+repo.put(checkpoint)
+restored = repo.get(checkpoint.id)
+```
+
+## Extensibility
+- Implement custom repositories by inheriting from `BaseCheckpointRepository` (e.g., for database, file, or cloud storage).
+- Extend versioning logic via the `VersionUtils` class.

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

@@ -0,0 +1,245 @@
+from typing import Any, Dict, Optional, List
+import copy
+import uuid
+from datetime import datetime, timezone
+from abc import ABC, abstractmethod
+import asyncio
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class CheckpointMetadata(BaseModel):
+    """
+    Metadata for a checkpoint, including session and task identifiers.
+
+    Attributes:
+        session_id (str): The session identifier (required).
+        task_id (Optional[str]): The task identifier (optional).
+        artifact_id (Optional[str]): The artifact identifier (optional).
+    """
+    session_id: str = Field(..., description="The session identifier.")
+    task_id: Optional[str] = Field(None, description="The task identifier.")
+    artifact_id: Optional[str] = Field(None, description="The artifact identifier.")
+
+    model_config = ConfigDict(extra="allow")
+
+class Checkpoint(BaseModel):
+    """
+    Core structure for a state checkpoint.
+    
+    Attributes:
+        id (str): Unique identifier for the checkpoint.
+        ts (str): Timestamp of the checkpoint.
+        metadata (CheckpointMetadata): Metadata associated with the checkpoint.
+        values (dict[str, Any]): State values stored in the checkpoint.
+        version (str): Version of the checkpoint format.
+        parent_id (Optional[str]): Parent checkpoint identifier, if any.
+        namespace (str): Namespace for the checkpoint, default is 'aworld'.
+    """
+    id: str = Field(..., description="Unique identifier for the checkpoint.")
+    ts: str = Field(..., description="Timestamp of the checkpoint.")
+    metadata: CheckpointMetadata = Field(..., description="Metadata associated with the checkpoint.")
+    values: Dict[str, Any] = Field(..., description="State values stored in the checkpoint.")
+    version: int = Field(..., description="Version of the checkpoint format.")
+    parent_id: Optional[str] = Field(default=None, description="Parent checkpoint identifier, if any.")
+    namespace: str = Field(default="aworld", description="Namespace for the checkpoint, default is 'aworld'.")
+
+def empty_checkpoint() -> Checkpoint:
+    """
+    Create an empty checkpoint with default values.
+    
+    Returns:
+        Checkpoint: An empty checkpoint structure.
+    """
+    return Checkpoint(
+        id=str(uuid.uuid4()),
+        ts=datetime.now(timezone.utc).isoformat(),
+        metadata=CheckpointMetadata(session_id="", task_id=None),
+        values={},
+        version=1,
+        parent_id=None,
+        namespace="aworld",
+    )
+
+def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
+    """
+    Create a deep copy of a checkpoint.
+    
+    Args:
+        checkpoint (Checkpoint): The checkpoint to copy.
+    Returns:
+        Checkpoint: A deep copy of the provided checkpoint.
+    """
+    return copy.deepcopy(checkpoint)
+
+def create_checkpoint(
+    values: Dict[str, Any],
+    metadata: CheckpointMetadata,
+    parent_id: Optional[str] = None,
+    version: int = 1,
+    namespace: str = 'aworld',
+) -> Checkpoint:
+    """
+    Create a new checkpoint from provided state values and metadata.
+    
+    Args:
+        values (dict[str, Any]): State values to store in the checkpoint.
+        metadata (CheckpointMetadata): Metadata for the checkpoint.
+        parent_id (Optional[str]): Parent checkpoint identifier, if any.
+        version (str): Version of the checkpoint format.
+        namespace (str): Namespace for the checkpoint.
+    Returns:
+        Checkpoint: The newly created checkpoint.
+    """
+    return Checkpoint(
+        id=str(uuid.uuid4()),
+        ts=datetime.now(timezone.utc).isoformat(),
+        metadata=metadata,
+        values=values,
+        version=version,
+        parent_id=parent_id,
+        namespace=namespace,
+    )
+
+class BaseCheckpointRepository(ABC):
+    """
+    Abstract base class for a checkpoint repository.
+    Provides synchronous and asynchronous methods for checkpoint management.
+    """
+
+    @abstractmethod
+    def get(self, checkpoint_id: str) -> Optional[Checkpoint]:
+        """
+        Retrieve a checkpoint by its unique identifier.
+        
+        Args:
+            checkpoint_id (str): The unique identifier of the checkpoint.
+        Returns:
+            Optional[Checkpoint]: The checkpoint if found, otherwise None.
+        """
+        pass
+
+    @abstractmethod
+    def list(self, params: Dict[str, Any]) -> List[Checkpoint]:
+        """
+        List checkpoints matching the given parameters.
+        
+        Args:
+            params (dict): Parameters to filter checkpoints.
+        Returns:
+            List[Checkpoint]: List of matching checkpoints.
+        """
+        pass
+
+    @abstractmethod
+    def put(self, checkpoint: Checkpoint) -> None:
+        """
+        Store a checkpoint.
+        
+        Args:
+            checkpoint (Checkpoint): The checkpoint to store.
+        """
+        pass
+
+    @abstractmethod
+    def get_by_session(self, session_id: str) -> Optional[Checkpoint]:
+        """
+        Get the latest checkpoint for a session.
+        
+        Args:
+            session_id (str): The session identifier.
+        Returns:
+            Optional[Checkpoint]: The latest checkpoint if found, otherwise None.
+        """
+        pass
+
+    @abstractmethod
+    def delete_by_session(self, session_id: str) -> None:
+        """
+        Delete all checkpoints related to a session.
+        
+        Args:
+            session_id (str): The session identifier.
+        """
+        pass
+
+    # Async methods
+    async def aget(self, checkpoint_id: str) -> Optional[Checkpoint]:
+        """
+        Asynchronously retrieve a checkpoint by its unique identifier.
+        
+        Args:
+            checkpoint_id (str): The unique identifier of the checkpoint.
+        Returns:
+            Optional[Checkpoint]: The checkpoint if found, otherwise None.
+        """
+        return await asyncio.to_thread(self.get, checkpoint_id)
+
+    async def alist(self, params: Dict[str, Any]) -> List[Checkpoint]:
+        """
+        Asynchronously list checkpoints matching the given parameters.
+        
+        Args:
+            params (dict): Parameters to filter checkpoints.
+        Returns:
+            List[Checkpoint]: List of matching checkpoints.
+        """
+        return await asyncio.to_thread(self.list, params)
+
+    async def aput(self, checkpoint: Checkpoint) -> None:
+        """
+        Asynchronously store a checkpoint.
+        
+        Args:
+            checkpoint (Checkpoint): The checkpoint to store.
+        """
+        await asyncio.to_thread(self.put, checkpoint)
+
+    async def aget_by_session(self, session_id: str) -> Optional[Checkpoint]:
+        """
+        Asynchronously get the latest checkpoint for a session.
+        
+        Args:
+            session_id (str): The session identifier.
+        Returns:
+            Optional[Checkpoint]: The latest checkpoint if found, otherwise None.
+        """
+        return await asyncio.to_thread(self.get_by_session, session_id)
+
+    async def adelete_by_session(self, session_id: str) -> None:
+        """
+        Asynchronously delete all checkpoints related to a session.
+        
+        Args:
+            session_id (str): The session identifier.
+        """
+        await asyncio.to_thread(self.delete_by_session, session_id)
+
+class VersionUtils:
+
+    @staticmethod
+    def get_next_version(version: int) -> int:
+        """
+        Get the next version of the checkpoint.
+        """
+        return version + 1
+    
+    @staticmethod
+    def get_previous_version(version: int)  -> int:
+        """
+        Get the previous version of the checkpoint.
+        """
+        return version - 1
+        
+    @staticmethod
+    def is_version_greater(checkpoint: Checkpoint, version: int) -> bool:
+        """
+        Check if the checkpoint version is greater than the given version.
+        """
+        return checkpoint.version > version
+    
+    @staticmethod
+    def is_version_less(checkpoint: Checkpoint, version: int) -> bool:
+        """
+        Check if the checkpoint version is less than the given version.
+        """
+        return checkpoint.version < version

chapter9/gaia-experience/AWorld/aworld/checkpoint/inmemory.py

@@ -0,0 +1,116 @@
+from typing import Any, Dict, List, Optional
+from . import Checkpoint, BaseCheckpointRepository, VersionUtils
+
+class InMemoryCheckpointRepository(BaseCheckpointRepository):
+    """
+    In-memory implementation of BaseCheckpointRepository.
+    Stores checkpoints in a simple in-memory dictionary.
+    Thread safety is not guaranteed.
+    """
+    def __init__(self) -> None:
+        """
+        Initialize the in-memory checkpoint repository.
+        """
+        self._checkpoints: Dict[str, Checkpoint] = {}
+        self._session_index: Dict[str, List[str]] = {}
+
+    def get(self, checkpoint_id: str) -> Optional[Checkpoint]:
+        """
+        Retrieve a checkpoint by its unique identifier.
+        Args:
+            checkpoint_id (str): The unique identifier of the checkpoint.
+        Returns:
+            Optional[Checkpoint]: The checkpoint if found, otherwise None.
+        """
+        return self._checkpoints.get(checkpoint_id)
+
+    def list(self, params: Dict[str, Any]) -> List[Checkpoint]:
+        """
+        List checkpoints matching the given parameters.
+        Args:
+            params (dict): Parameters to filter checkpoints.
+        Returns:
+            List[Checkpoint]: List of matching checkpoints.
+        """
+        result = []
+        for cp in self._checkpoints.values():
+            match = True
+            for k, v in params.items():
+                if k == 'session_id':
+                    if cp.metadata.session_id != v:
+                        match = False
+                        break
+                elif k == 'task_id':
+                    if cp.metadata.task_id != v:
+                        match = False
+                        break
+                elif cp.get(k) != v:
+                    match = False
+                    break
+            if match:
+                result.append(cp)
+        return result
+
+    def put(self, checkpoint: Checkpoint) -> None:
+        """
+        Store a checkpoint.
+        Args:
+            checkpoint (Checkpoint): The checkpoint to store.
+        """
+        # Find last version checkpoint by session_id
+        last_checkpoint = self.get_by_session(checkpoint.metadata.session_id)
+        
+        if last_checkpoint:
+            # Compare versions to ensure optimistic locking
+            if VersionUtils.is_version_less(checkpoint, last_checkpoint.version):
+                raise ValueError(f"New checkpoint version {checkpoint.version} must be greater than last version {last_checkpoint.version}")
+            
+        # Store the new checkpoint
+        self._checkpoints[checkpoint.id] = checkpoint
+            
+        # Update session index
+        session_id = checkpoint.metadata.session_id
+        if session_id:
+            if session_id not in self._session_index:
+                self._session_index[session_id] = []
+            self._session_index[session_id].append(checkpoint.id)
+
+    def get_by_session(self, session_id: str) -> Optional[Checkpoint]:
+        """
+        Get the latest checkpoint for a session.
+        Args:
+            session_id (str): The session identifier.
+        Returns:
+            Optional[Checkpoint]: The latest checkpoint if found, otherwise None.
+        """
+        ids = self._session_index.get(session_id, [])
+        if not ids:
+            return None
+        # Assume the last one is the latest
+        last_id = ids[-1]
+        return self._checkpoints.get(last_id)
+
+    def delete_by_session(self, session_id: str) -> None:
+        """
+        Delete all checkpoints related to a session.
+        Args:
+            session_id (str): The session identifier.
+        """
+        ids = self._session_index.pop(session_id, [])
+        for cid in ids:
+            self._checkpoints.pop(cid, None)
+
+    async def alist(self, params: Dict[str, Any]) -> List[Checkpoint]:
+        return self.list(params)
+
+    async def aget(self, checkpoint_id: str) -> Optional[Checkpoint]:
+        return self.get(checkpoint_id)
+
+    async def aput(self, checkpoint: Checkpoint) -> None:
+        self.put(checkpoint)
+
+    async def aget_by_session(self, session_id: str) -> Optional[Checkpoint]:
+        return self.get_by_session(session_id)
+
+    async def adelete_by_session(self, session_id: str) -> None:
+        self.delete_by_session(session_id)