Upload 5 files

aworld/core/context/README.md

@@ -0,0 +1,281 @@
+# AWorld Context Management
+
+Core context management system in the AWorld architecture, providing hierarchical state management through Context and AgentContext components for complete Agent state information storage and coordination.
+
+## Architecture Overview
+
+![Context Management](../../../readme_assets/context_management.png)
+
+The Context Management system implements intelligent context processing with multiple optimization strategies based on context length analysis and configuration parameters.
+
+## Features
+
+- Hierarchical context management with global and agent-specific scopes
+- Complete Agent state management with support for state restoration and recovery
+- Immutable configuration management and mutable runtime state tracking
+- Intelligent LLM Prompt management and context optimization
+- Complete LLM call intervention and control mechanisms
+- Hook system support for extensible processing workflows
+- Content compression and context window management
+
+## Core Components
+
+- `Context`: Global singleton context manager spanning the entire AWorld Runner lifecycle
+- `AgentContext`: Agent-specific context container for individual Agent execution periods
+- `PromptProcessor`: Intelligent context processor supporting content compression and truncate
+- `Hook System`: Extensible hook system supporting full-process LLM call intervention
+
+
+## Multi-Level Context
+
+![Context Lifecycle](../../../readme_assets/context_lifecycle.png)
+
+The AWorld framework implements a hierarchical context management system with distinct lifecycles:
+
+### Context Lifecycle (Session)
+Context is a singleton object that spans the entire AWorld Runner execution lifecycle, enabling global state sharing and coordination between multiple Agents in the AWorld framework.
+
+- **Scope**: Spans the entire AWorld Runner execution period
+- **Responsibility**: Global state management, task coordination, and resource allocation
+    - **Dictionary Interface**: Supports simple key-value state storage using `context['key'] = value` syntax
+- **Function**: Manages multiple AgentContext instances and enables cross-agent data exchange
+    - **Multi-Agent Coordination**: Manages multiple AgentContext instances and enables seamless data exchange between different Agents
+    - **Session Management**: Provides (Session)Context API for cross-agent state management
+    - **Task Coordination**: Handles task management, agent coordination, and resource allocation
+
+### AgentContext Lifecycle (Agent-Specific)
+- **Scope**: Spans individual Agent execution period
+- **Responsibility**: Agent-specific state management and runtime tracking
+- **Function**: 
+  - **Configuration Management**: Maintains immutable Agent configuration (agent_id, agent_name, agent_desc, system_prompt, agent_prompt, tool_names, context_rule)
+  - **Runtime State Tracking**: Manages mutable runtime state (tools, step, messages, context_usage)
+  - **Dynamic Prompt Management**: Supports runtime modification of system_prompt and agent_prompt based on execution context
+  - **Tool Lifecycle Management**: Handles tool object initialization, execution, and state management
+  - **Conversation History**: Maintains complete message history throughout Agent execution
+  - **Step-by-Step Execution**: Tracks current execution step and manages step transitions
+  - **Context Optimization**: Monitors context usage statistics and applies context processing rules
+  - **State Persistence**: Preserves Agent state across multiple LLM calls and tool invocations within a single execution period
+
+### Example: State Management and Recovery
+
+> **📋 Test Implementation**: See complete test implementation at [`tests/test_context_management.py::TestContextManagement::test_state_management_and_recovery()`](../../../tests/test_context_management.py)
+
+```python
+class StateModifyAgent(Agent):
+    async def async_policy(self, observation, info=None, **kwargs):
+        result = await super().async_policy(observation, info, **kwargs)
+        self.context.state['policy_executed'] = True
+        return result
+
+class StateTrackingAgent(Agent):
+    async def async_policy(self, observation, info=None, **kwargs):
+        result = await super().async_policy(observation, info, **kwargs)
+        assert self.context.state['policy_executed'] == True
+        return result
+
+# Create custom agent instance
+custom_agent = StateModifyAgent(
+    conf=AgentConfig(
+        llm_model_name=mock_model_name,
+        llm_base_url=mock_base_url,
+        llm_api_key=mock_api_key
+    ),
+    name="state_modify_agent",
+    system_prompt="You are a Python expert who provides detailed and practical answers.",
+    agent_prompt="You are a Python expert who provides detailed and practical answers.",
+)
+
+# Create a second agent for multi-agent testing
+second_agent = StateTrackingAgent(
+    conf=AgentConfig(
+        llm_model_name=mock_model_name,
+        llm_base_url=mock_base_url,
+        llm_api_key=mock_api_key
+    ),
+    name="state_tracking_agent",
+    system_prompt="You are a helpful assistant.",
+    agent_prompt="You are a helpful assistant.",
+)
+
+# Run multi-agent scenario
+response = run_multi_agent(
+    input="What is an agent. describe within 20 words",
+    agent1=custom_agent,
+    agent2=second_agent
+)
+
+# Verify state changes after execution
+assert custom_agent.context.state.get('policy_executed', True)
+assert second_agent.agent_context.state.get('policy_executed', True)
+```
+
+### Example: Hook System, LLM Call Intervention and Agent state sharing
+
+> **📋 Test Implementation**: See complete Hook system test implementations at:
+> - [`tests/test_context_management.py::TestHookSystem::test_hook_registration()`](../../../tests/test_context_management.py) - Hook registration test
+> - [`tests/test_context_management.py::TestHookSystem::test_hook_execution()`](../../../tests/test_context_management.py) - Hook execution test
+
+```python
+from aworld.runners.hook.hooks import PreLLMCallHook, PostLLMCallHook
+from aworld.runners.hook.hook_factory import HookFactory
+from aworld.utils.common import convert_to_snake
+from aworld.core.event.base import Message
+from aworld.core.context.base import Context
+
+# Test Hook System functionality
+@HookFactory.register(name="TestPreLLMHook", desc="Test pre-LLM hook")
+class TestPreLLMHook(PreLLMCallHook):
+    """Test hook for pre-LLM processing"""
+    
+    def name(self):
+        return convert_to_snake("TestPreLLMHook")
+    
+    async def exec(self, message: Message, context: Context = None) -> Message:
+        """Test hook execution"""
+        agent_context = context.get_agent_context(message.sender)
+        if agent_context is not None:
+            agent_context.step = 1 
+        
+        assert agent_context.step == 1 or agent_context.step == 2
+        return message
+
+
+@HookFactory.register(name="TestPostLLMHook", desc="Test post-LLM hook")
+class TestPostLLMHook(PostLLMCallHook):
+    """Test hook for post-LLM processing"""
+    
+    def name(self):
+        return convert_to_snake("TestPostLLMHook")
+    
+    async def exec(self, message: Message, context: Context = None) -> Message:
+        """Test hook execution with llm_output processing"""
+        agent_context = context.get_agent_context(message.sender)
+        if agent_context is not None and agent_context.llm_output is not None:
+            # Test dynamic prompt adjustment based on LLM output
+            if hasattr(agent_context.llm_output, 'content'):
+                content = agent_context.llm_output.content.lower()
+                if content is not None:
+                    agent_context.agent_prompt = "Success mode activated"
+
+        assert agent_context.agent_prompt == "Success mode activated"
+        return message
+
+# Test hook registration and retrieval
+assert "TestPreLLMHook" in HookFactory._cls
+assert "TestPostLLMHook" in HookFactory._cls
+
+# Test hook creation using __call__ method
+pre_hook = HookFactory("TestPreLLMHook")
+post_hook = HookFactory("TestPostLLMHook")
+
+assert isinstance(pre_hook, TestPreLLMHook)
+assert isinstance(post_hook, TestPostLLMHook)
+
+# Test hook execution
+response = agent.run("What is an agent. describe within 20 words")
+assert response.answer is not None
+```
+
+## Context Rule Configuration
+
+`ContextRuleConfig` provides comprehensive context management through two main configuration components:
+
+### OptimizationConfig
+
+Controls context optimization behavior:
+
+- `enabled`: Whether to enable context optimization (default: `False`)
+- `max_token_budget_ratio`: Maximum token budget ratio for context window usage (default: `0.5`, range: 0.0-1.0)
+
+### LlmCompressionConfig (Beta Feature)
+
+**⚠️ Beta Feature**: This configuration is currently in beta and may undergo changes in future versions.
+
+Controls intelligent context compression:
+
+- `enabled`: Whether to enable LLM-based compression (default: `False`)
+- `trigger_compress_token_length`: Token threshold to trigger basic compression (default: `10000`)
+- `trigger_mapreduce_compress_token_length`: Token threshold to trigger map-reduce compression (default: `100000`)
+- `compress_model`: ModelConfig for compression LLM calls (optional)
+
+
+### Example: Using Default Context Configuration (Recommended)
+
+> **📋 Test Implementation**: See default configuration test at [`tests/test_context_management.py::TestContextManagement::test_default_context_configuration()`](../../../tests/test_context_management.py)
+
+```python
+from aworld.agents.llm_agent import Agent
+from aworld.config.conf import AgentConfig
+
+# No need to explicitly configure context_rule, system automatically uses default configuration
+# Default configuration is equivalent to:
+# context_rule=ContextRuleConfig(
+#     optimization_config=OptimizationConfig(
+#         enabled=True,
+#         max_token_budget_ratio=1.0  # Use 100% of context window
+#     ),
+#     llm_compression_config=LlmCompressionConfig(
+#         enabled=False  # Compression disabled by default
+#     )
+# )
+response = agent.run("What is an agent. describe within 20 words")
+
+assert response.answer is not None
+assert agent.agent_context.model_config.llm_model_name == "llama-2-7b-chat-hf-function-calling-v2"
+
+# Test default context rule behavior
+assert agent.agent_context.context_rule is not None
+assert agent.agent_context.context_rule.optimization_config is not None
+```
+
+### Example: Custom Context Configuration
+
+> **📋 Test Implementation**: See custom configuration test at [`tests/test_context_management.py::TestContextManagement::test_custom_context_configuration()`](../../../tests/test_context_management.py)
+
+```python
+from aworld.config.conf import AgentConfig, ContextRuleConfig, OptimizationConfig, LlmCompressionConfig, ModelConfig
+
+# Create custom context rules
+context_rule = ContextRuleConfig(
+    optimization_config=OptimizationConfig(
+        enabled=True,
+        max_token_budget_ratio=0.8  # Use 80% of context window
+    ),
+    llm_compression_config=LlmCompressionConfig(
+        enabled=True,  # Enable beta compression feature
+        trigger_compress_token_length=100,
+        trigger_mapreduce_compress_token_length=1000,
+        compress_model=ModelConfig(
+            llm_model_name="llama-2-7b-chat-hf-function-calling-v2",
+            llm_base_url="http://localhost:1234/v1",
+            llm_api_key="lm-studio",
+        )
+    )
+)
+
+# Save original rule for restoration
+origin_rule = agent.agent_context.context_rule
+agent.update_context_rule(context_rule)
+
+# Test the agent with custom configuration
+response = agent.run("What is an agent. describe within 20 words")
+assert response.answer is not None
+
+# Test configuration values
+assert agent.agent_context.context_rule.optimization_config.enabled == True
+assert agent.agent_context.context_rule.llm_compression_config.enabled == True
+
+# Restore original rule
+agent.update_context_rule(origin_rule)
+```
+
+## Notes
+
+1. **Hierarchical Lifecycle**: Context spans the entire AWorld Runner execution while AgentContext spans individual Agent executions, as illustrated in the context lifecycle diagram.
+2. **Beta Features**: The `llm_compression_config` is currently in beta. Use with caution in production environments.
+3. **Performance Trade-offs**: Enabling compression can save token usage but increases processing time. Adjust configuration based on actual needs.
+4. **Model Compatibility**: Different models have different context length limitations. The system automatically adapts to model capabilities.
+5. **Default Configuration**: The system provides reasonable default configuration. Manual configuration is unnecessary for most scenarios.
+6. **State Management**: Context and AgentContext support state sharing between multiple Agents and ensures state consistency. State persistence functionality is currently under development.
+
+Through proper configuration of Context and AgentContext with context processors, you can significantly improve Agent performance in long conversations and complex tasks while optimizing token usage efficiency.

aworld/core/context/__init__.py

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

aworld/core/context/base.py

@@ -0,0 +1,347 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Dict, List, Any, Optional, TYPE_CHECKING
+
+from aworld.config import ConfigDict
+from aworld.config.conf import ContextRuleConfig, ModelConfig
+from aworld.core.context.context_state import ContextState
+from aworld.core.context.session import Session
+from aworld.core.singleton import InheritanceSingleton
+from aworld.models.model_response import ModelResponse
+from aworld.utils.common import nest_dict_counter
+
+if TYPE_CHECKING:
+    from aworld.core.task import Task
+
+@dataclass
+class ContextUsage:
+    total_context_length: int = 128000
+    used_context_length: int = 0
+    def __init__(self, total_context_length: int = 128000, used_context_length: int = 0):
+        self.total_context_length = total_context_length
+        self.used_context_length = used_context_length
+
+class Context(InheritanceSingleton):
+    """Single instance, can use construction or `instance` static method to create or get `Context` instance.
+
+    Examples:
+        >>> context = Context()
+    or
+        >>> context = Context.instance()
+    """
+
+    def __init__(self,
+                 user: str = None,
+                 task_id: str = None,
+                 trace_id: str = None,
+                 session: Session = None,
+                 engine: str = None,
+                 **kwargs):
+
+        super().__init__()
+        self._user = user
+        self._init(task_id=task_id, trace_id=trace_id, session=session, engine=engine, **kwargs)
+
+    def _init(self, *, task_id: str = None, trace_id: str = None, session: Session = None, engine: str = None):
+        self._task_id = task_id
+        self._task = None
+        self._engine = engine
+        self._trace_id = trace_id
+        self._session: Session = session
+        self.context_info = ConfigDict()
+        self.agent_info = ConfigDict()
+        self.trajectories = OrderedDict()
+        self._token_usage = {
+            "completion_tokens": 0,
+            "prompt_tokens": 0,
+            "total_tokens": 0,
+        }
+        self.state = ContextState()
+        # TODO swarm topology
+        # TODO workspace
+
+        self._swarm = None
+        self._event_manager = None
+
+    def add_token(self, usage: Dict[str, int]):
+        self._token_usage = nest_dict_counter(self._token_usage, usage)
+
+    def reset(self, **kwargs):
+        self._init(**kwargs)
+
+    def set_task(self, task: 'Task'):
+        self._task = task
+
+    def get_task(self) -> 'Task':
+        return self._task
+
+    @property
+    def trace_id(self):
+        return self._trace_id
+
+    @trace_id.setter
+    def trace_id(self, trace_id):
+        self._trace_id = trace_id
+
+    @property
+    def token_usage(self):
+        return self._token_usage
+
+    @property
+    def engine(self):
+        return self._engine
+
+    @engine.setter
+    def engine(self, engine: str):
+        self._engine = engine
+
+    @property
+    def user(self):
+        return self._user
+
+    @user.setter
+    def user(self, user):
+        if user is not None:
+            self._user = user
+
+    @property
+    def task_id(self):
+        return self._task_id
+
+    @task_id.setter
+    def task_id(self, task_id):
+        if task_id is not None:
+            self._task_id = task_id
+
+    @property
+    def session_id(self):
+        if self.session:
+            return self.session.session_id
+        else:
+            return None
+
+    @property
+    def session(self):
+        return self._session
+
+    @session.setter
+    def session(self, session: Session):
+        self._session = session
+
+    @property
+    def swarm(self):
+        return self._swarm
+
+    @swarm.setter
+    def swarm(self, swarm: 'Swarm'):
+        self._swarm = swarm
+
+    @property
+    def event_manager(self):
+        return self._event_manager
+
+    @event_manager.setter
+    def event_manager(self, event_manager: 'EventManager'):
+        self._event_manager = event_manager
+
+    @property
+    def record_path(self):
+        return "."
+
+    @property
+    def is_task(self):
+        return True
+
+    @property
+    def enable_visible(self):
+        return False
+
+    @property
+    def enable_failover(self):
+        return False
+
+    @property
+    def enable_cluster(self):
+        return False
+
+    def get_state(self, key: str, default: Any = None) -> Any:
+        return self.state.get(key, default)
+    
+    def set_state(self, key: str, value: Any):
+        self.state[key] = value
+
+
+
+@dataclass
+class AgentContext:
+    """Agent context containing both configuration and runtime state.
+    
+    AgentContext is the core context management class in the AWorld architecture, used to store and manage
+    the complete state information of an Agent, including configuration data and runtime state. Its main functions are:
+    
+    1. **State Restoration**: Save all state information during Agent execution, supporting Agent state restoration and recovery
+    2. **Configuration Management**: Store Agent's immutable configuration information (such as agent_id, system_prompt, etc.)
+    3. **Runtime State Tracking**: Manage Agent's mutable state during execution (such as messages, step, tools, etc.)
+    4. **LLM Prompt Management**: Manage and maintain the complete prompt context required for LLM calls, including system prompts, historical messages, etc.
+    5. **LLM Call Intervention**: Provide complete control over the LLM call process through Hook and ContextProcessor
+    
+    ## Lifecycle
+    The lifecycle of AgentContext is completely consistent with the Agent instance:
+    - **Creation**: Created during Agent initialization, containing initial configuration
+    - **Runtime**: Continuously update runtime state during Agent execution
+    - **Destruction**: Destroyed along with Agent instance destruction
+    ```
+    ┌─────────────────────── AWorld Runner ─────────────────────────┐
+    |  ┌──────────────────── Agent Execution ────────────────────┐  │
+    │  │  ┌────────────── Step 1 ─────────────┐ ┌── Step 2 ──┐   │  │
+    │  │  │  [LLM Call]     [Tool Call(s)]    │
+    │  │  │  [       Context Update      ]    │
+    ```
+    
+    ## Field Classification
+    - **Immutable Configuration Fields**: agent_id, agent_name, agent_desc, system_prompt, 
+      agent_prompt, tool_names, context_rule
+    - **Mutable Runtime Fields**: tools, step, messages, context_usage, llm_output
+    
+    ## LLM Call Intervention Mechanism
+    AgentContext implements complete control over LLM calls through the following mechanisms:
+    
+    1. **Hook System**:
+       - pre_llm_call_hook: Context preprocessing before LLM call
+       - post_llm_call_hook: Result post-processing after LLM call
+       - pre_tool_call_hook: Context adjustment before tool call
+       - post_tool_call_hook: State update after tool call
+    
+    2. **PromptProcessor**:
+       - Prompt Optimization: Optimize prompt content based on context length limitations
+       - Message Compression: Intelligently compress historical messages to fit model context window
+       - Context Rules: Apply context_rule for customized context processing
+    
+    ## Usage Scenarios
+    1. **Agent Initialization**: Create AgentContext containing configuration information
+    2. **LLM Call Control**: Pass as info parameter in policy(), async_policy() methods to control LLM behavior
+    3. **Hook Callbacks**: Access and modify LLM call context in various Hooks, use PromptProcessor for prompt optimization and context processing
+    4. **State Recovery**: Recover Agent's complete state from persistent storage
+    """
+    
+    # ===== Immutable Configuration Fields =====
+    agent_id: str = None
+    agent_name: str = None
+    agent_desc: str = None
+    system_prompt: str = None
+    agent_prompt: str = None
+    tool_names: List[str] = None
+    model_config: ModelConfig = None
+    context_rule: ContextRuleConfig = None
+    _context: Context = None
+    state: ContextState = None
+
+    # ===== Mutable Configuration Fields =====
+    tools: List[str] = None
+    step: int = 0
+    messages: List[Dict[str, Any]] = None
+    context_usage: ContextUsage = None
+    llm_output: ModelResponse = None
+
+    def __init__(self, 
+                 agent_id: str = None,
+                 agent_name: str = None,
+                 agent_desc: str = None,
+                 system_prompt: str = None,
+                 agent_prompt: str = None,
+                 tool_names: List[str] = None,
+                 model_config: ModelConfig = None,
+                 context_rule: ContextRuleConfig = None,
+                 tools: List[str] = None,
+                 step: int = 0,
+                 messages: List[Dict[str, Any]] = None,
+                 context_usage: ContextUsage = None,
+                 llm_output: ModelResponse = None,
+                 context: Context = None,
+                 parent_state: ContextState = None,
+                 **kwargs):
+        # Configuration fields
+        self.agent_id = agent_id
+        self.agent_name = agent_name
+        self.agent_desc = agent_desc
+        self.system_prompt = system_prompt
+        self.agent_prompt = agent_prompt
+        self.tool_names = tool_names if tool_names is not None else []
+        self.model_config = model_config
+        self.context_rule = context_rule
+        
+        # Runtime state fields
+        self.tools = tools if tools is not None else []
+        self.step = step
+        self.messages = messages if messages is not None else []
+        self.context_usage = context_usage if context_usage is not None else ContextUsage()
+        self.llm_output = llm_output
+
+        # Initialize Context with runner(session) level context
+        self._context = context
+        # Initialize ContextState with parent state (Context's state)
+        # If context_state is provided, use it as parent; otherwise will be set later
+        self.state = ContextState(parent_state=parent_state)
+
+        # Additional fields for backward compatibility
+        self._init(**kwargs)
+
+    def _init(self, **kwargs):
+        self._task_id = kwargs.get('task_id')
+    
+    def set_model_config(self, model_config: ModelConfig):
+        self.model_config = model_config
+
+    def set_messages(self, messages: List[Dict[str, Any]]):
+        self.messages = messages
+
+    def set_tools(self, tools: List[str]):
+        self.tools = tools
+    
+    def set_llm_output(self, llm_output: ModelResponse):
+        self.llm_output = llm_output
+
+    def increment_step(self) -> int:
+        self.step += 1
+        return self.step
+    
+    def set_step(self, step: int):
+        self.step = step
+    
+    def get_step(self) -> int:
+        return self.step
+    
+    def update_context_usage(self, used_context_length: int = None, total_context_length: int = None):
+        if used_context_length is not None:
+            self.context_usage.used_context_length = used_context_length
+        if total_context_length is not None:
+            self.context_usage.total_context_length = total_context_length
+
+    def get_context_usage_ratio(self) -> float:
+        """Get context usage ratio"""
+        if self.context_usage.total_context_length <= 0:
+            return 0.0
+        return self.context_usage.used_context_length / self.context_usage.total_context_length
+
+    def set_parent_state(self, parent_state: ContextState):
+        self.state._parent_state = parent_state
+
+    def get_state(self, key: str, default: Any = None) -> Any:
+        return self.state.get(key, default)
+    
+    def set_state(self, key: str, value: Any):
+        self.state[key] = value
+
+    def get_task(self) -> 'Task':
+        return self._context.get_task()
+
+    def get_session(self) -> Session:
+        return self._context.get_session()
+
+    def get_engine(self) -> str:
+        return self._context.get_engine()
+    
+    def get_user(self) -> str:
+        return self._context.get_user()

aworld/core/context/context_state.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+# coding: utf-8
+
+"""
+Context State Management System
+Provides hierarchical state management with parent-child state inheritance
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+
+class ContextState:
+    """
+    Hierarchical state management class
+    
+    Supports the following features:
+    - Parent-child state inheritance: Child states can access parent state data
+    - State priority: Local state takes precedence over parent state
+    - State isolation: Write operations only affect local state
+    - Flexible lookup: Supports multi-level state lookup
+    """
+    
+    def __init__(self, parent_state: Optional['ContextState'] = None):
+        """
+        Initialize ContextState
+        
+        Args:
+            parent_state: Parent state object for implementing state inheritance
+        """
+        self._data: Dict[str, Any] = {}
+        self._parent_state: Optional['ContextState'] = parent_state
+    
+    def __getitem__(self, key: str) -> Any:
+        """Get state value with parent state inheritance support"""
+        if key in self._data:
+            return self._data[key]
+        elif self._parent_state is not None:
+            return self._parent_state[key]
+        else:
+            raise KeyError(f"Key '{key}' not found in state")
+    
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Set state value, only writes to local state"""
+        self._data[key] = value
+    
+    def __delitem__(self, key: str) -> None:
+        """Delete state value, only deletes from local state"""
+        if key in self._data:
+            del self._data[key]
+        else:
+            raise KeyError(f"Key '{key}' not found in local state")
+    
+    def __contains__(self, key: str) -> bool:
+        """Check if contains specified key, including parent state"""
+        return key in self._data or (self._parent_state is not None and key in self._parent_state)
+    
+    def __len__(self) -> int:
+        """Return the count of all accessible states (including parent state)"""
+        keys = set(self._data.keys())
+        if self._parent_state is not None:
+            keys.update(self._parent_state.keys())
+        return len(keys)
+    
+    def __iter__(self):
+        """Iterate over all accessible keys (including parent state)"""
+        keys = set(self._data.keys())
+        if self._parent_state is not None:
+            keys.update(self._parent_state.keys())
+        return iter(keys)
+    
+    def get(self, key: str, default: Any = None) -> Any:
+        """
+        Get state value, return default value if not exists
+        
+        Args:
+            key: The key to get
+            default: Default value
+            
+        Returns:
+            Value corresponding to key or default value
+        """
+        try:
+            return self[key]
+        except KeyError:
+            return default
+    
+    def update(self, other: Union[Dict[str, Any], 'ContextState']) -> None:
+        """
+        Batch update state
+        
+        Args:
+            other: Data to update, can be dict or another ContextState
+        """
+        if isinstance(other, dict):
+            self._data.update(other)
+        elif isinstance(other, ContextState):
+            self._data.update(other._data)
+        else:
+            raise TypeError("update() argument must be dict or ContextState")
+    
+    def pop(self, key: str, default: Any = None) -> Any:
+        """
+        Delete and return value of specified key, only operates on local state
+        
+        Args:
+            key: The key to delete
+            default: Default value to return if key doesn't exist
+            
+        Returns:
+            The deleted value or default value
+        """
+        return self._data.pop(key, default)
+    
+    def clear(self) -> None:
+        """Clear local state (does not affect parent state)"""
+        self._data.clear()
+    
+    def keys(self) -> List[str]:
+        """Return list of all accessible keys (including parent state)"""
+        keys = set(self._data.keys())
+        if self._parent_state is not None:
+            keys.update(self._parent_state.keys())
+        return list(keys)
+    
+    def values(self) -> List[Any]:
+        """Return list of all accessible values (including parent state)"""
+        return [self[key] for key in self.keys()]
+    
+    def items(self) -> List[tuple]:
+        """Return list of all accessible (key, value) pairs (including parent state)"""
+        return [(key, self[key]) for key in self.keys()]
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert to dictionary, including inherited parent state
+        
+        Returns:
+            Dictionary containing all accessible states
+        """
+        result = {}
+        if self._parent_state is not None:
+            result.update(self._parent_state.to_dict())
+        result.update(self._data)
+        return result
+    
+    def local_dict(self) -> Dict[str, Any]:
+        """
+        Get local state dictionary (excluding parent state)
+        
+        Returns:
+            Dictionary containing only local state
+        """
+        return self._data.copy()
+    
+    def set_parent(self, parent_state: Optional['ContextState']) -> None:
+        """
+        Set parent state
+        
+        Args:
+            parent_state: New parent state object
+        """
+        self._parent_state = parent_state
+    
+    def get_parent(self) -> Optional['ContextState']:
+        """
+        Get parent state object
+        
+        Returns:
+            Parent state object or None
+        """
+        return self._parent_state
+    
+    def has_parent(self) -> bool:
+        """
+        Check if has parent state
+        
+        Returns:
+            True if has parent state, False otherwise
+        """
+        return self._parent_state is not None
+    
+    def __repr__(self) -> str:
+        """Return string representation of state"""
+        local_count = len(self._data)
+        total_count = len(self)
+        parent_info = f" (with parent: {total_count - local_count} inherited)" if self._parent_state else ""
+        return f"ContextState(local: {local_count}{parent_info})"
+    
+    def __str__(self) -> str:
+        """Return detailed string representation of state"""
+        return f"ContextState({self.to_dict()})"

aworld/core/context/session.py

@@ -0,0 +1,13 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import time
+import uuid
+from typing import List
+
+from pydantic import BaseModel
+
+
+class Session(BaseModel):
+    session_id: str = str(uuid.uuid1().hex)
+    last_update_time: float = time.time()
+    trajectories: List = []