Custom Modules¶

AgentSociety 2 supports creating and registering custom agent and environment modules, allowing you to extend the platform with your own simulation components.

Overview¶

The custom module system allows you to:

Create custom agent classes with specialized behaviors
Create custom environment modules with domain-specific tools
Automatically discover and register modules through API
Test custom modules with auto-generated test scripts
Seamlessly integrate with existing AgentSociety framework
Persist requirements, design, and validation artifacts through the Progressive Disclosure workflow

Directory Structure¶

Custom modules are placed in the custom/ directory within your workspace:

workspace/
├── custom/                    # User-created directory
│   ├── agents/                # Custom agent classes
│   │   └── my_agent.py
│   └── envs/                  # Custom environment modules
│       └── my_env.py
└── .agentsociety/             # Auto-generated configuration
    ├── agent_classes/
    ├── env_modules/
    └── custom_env_skill/
        └── runs/

Create Custom Agents¶

All custom agents must inherit from AgentBase. Agents are workspace-bound stateless records: state is restored in restore() and written back in to_workspace(); see Using Agents.

from pathlib import Path
from agentsociety2.agent.base import AgentBase
from datetime import datetime
from typing import Any

class MyAgent(AgentBase):
    """My custom Agent"""

    @classmethod
    def mcp_description(cls) -> str:
        return """MyAgent: A custom agent for specific tasks

    This agent demonstrates custom behavior.
    """

    async def restore(self, workspace_path: Path, service_proxy: Any) -> None:
        """恢复 workspace / 服务 / 技能，再追加自定义状态。"""
        await super().restore(workspace_path, service_proxy)
        self._custom_state: dict = {}

    async def ask(self, message: str, readonly: bool = True, *, t=None) -> str:
        """Respond to questions from the environment"""
        prompt = f"Question: {message}\nPlease answer:"
        response = await self.acompletion([{"role": "user", "content": prompt}])
        return response.choices[0].message.content or ""

    async def step(self, tick: int, t: datetime) -> str:
        """Execute one simulation step"""
        return f"Agent {self.id} executing step {tick}"

    async def to_workspace(self, workspace_path: Path) -> None:
        """把动态状态写回 workspace（AGENT.json 等）。"""
        self.persist_agent_json(tick=None, t=self._current_time)

Required Methods¶

Method	Description
`mcp_description()`	Return module description (class method, recommended to override; `AgentBase`/`EnvBase` has default description)
`ask()`	Answer questions from the environment
`step()`	Execute one simulation step
`to_workspace()`	Write dynamic state back to the workspace (required).
`step()`	Restore custom state after `await super().restore(...)` (recommended override).

Create Custom Environments¶

Custom environments must inherit from EnvBase and use the @tool decorator to register methods:

from agentsociety2.env import EnvBase, tool
from datetime import datetime

class MyEnv(EnvBase):
    """My custom environment"""

    def __init__(self, config=None):
        super().__init__()
        # Initialize your environment state

    @classmethod
    def mcp_description(cls) -> str:
        return """MyEnv: A custom environment

    This environment provides custom tools for agents.
    """

    @tool(readonly=True, kind="observe")
    async def get_state(self, agent_id: int) -> dict:
        """Get current environment state (observation tool)"""
        return {"agent_id": agent_id, "state": "normal"}

    @tool(readonly=False)
    async def do_action(self, agent_id: int, action: str) -> dict:
        """Perform an action (modification tool)"""
        return {"agent_id": agent_id, "action": action, "result": "success"}

    async def step(self, tick: int, t: datetime):
        """Environment step"""
        self.t = t

    async def to_workspace(self, workspace_path=None) -> None:
        """把动态状态写入 workspace；覆盖此方法可支持 resume，详见 :doc:`env_modules`。"""
        if workspace_path is not None:
            self._bind_workspace(workspace_path)

    async def restore(self, workspace_path) -> bool:
        """从 workspace 恢复动态状态；resume 时调用，默认返回 False。"""
        self._bind_workspace(workspace_path)
        return False

reality compatibility constraints¶

The generated custom environment module must still follow the actual compatibility constraints of the current repository:

File must be located at custom/envs/*.py
The class definition must be located directly in this file, it cannot just be re-exported
Register key to continue using class_name
There is at least one legal @tool
step() must be present
Parameterless instantiation should be supported by default cls()
If the module needs observation capability, provide an @tool(readonly=True, kind='observe') observation tool
Providing an informative mcp_description() is recommended; if not overridden, the base class default description is displayed
If the module has dynamic in-memory state and should support interrupted runs, override to_workspace and restore. When --resume is used, the module should be able to restore state from state/ENV_STATE.json; see “State persistence and resume” in Environment Modules.

Note

The scanner will skip files containing examples/ in the path (the example is for reference only and does not participate in registration).

@tool Decorator¶

The @tool decorator registers methods as tools accessible to agents:

Parameter	Description
`readonly=True`	Tool does not modify environment state
`readonly=False`	Tool can modify environment state
`kind="observe"`	Observation tool (single agent_id parameter, readonly=True)
`kind="statistics"`	Statistics tool (no parameters, readonly=True)
`kind=None`	Regular tool (any parameters, can be readonly=False)

Registering Custom Modules¶

After creating custom modules, register them using the API:

Scan and Register

curl -X POST http://localhost:8001/api/v1/custom/scan \
  -H "Content-Type: application/json" \
  -d '{"workspace_path": "/path/to/workspace"}'

List Registered Modules

curl http://localhost:8001/api/v1/custom/list

Test Custom Modules

curl -X POST http://localhost:8001/api/v1/custom/test \
  -H "Content-Type: application/json" \
  -d '{"workspace_path": "/path/to/workspace"}'

Create or restore workflow run

curl -X POST http://localhost:8001/api/v1/custom/workflow/runs \
  -H "Content-Type: application/json" \
  -d '{"workspace_path": "/path/to/workspace", "user_request": "create a resource env"}'

Validate workflow run

curl -X POST http://localhost:8001/api/v1/custom/workflow/runs/<run_id>/validate \
  -H "Content-Type: application/json" \
  -d '{"module_path": "custom/envs/my_env.py", "class_name": "MyEnv"}'

API Endpoints¶

Endpoint	Method	Description
`/api/v1/custom/scan`	POST	Scan and register custom modules
`/api/v1/custom/test`	POST	Test custom modules
`/api/v1/custom/clean`	POST	Clean custom module configuration
`/api/v1/custom/list`	GET	List registered custom modules
`/api/v1/custom/status`	GET	Get module status overview
`/api/v1/custom/workflow/runs`	POST	Create or restore a custom environment workflow run
`/api/v1/custom/workflow/runs/{run_id}/validate`	POST	Perform end-to-end verification of scanner/tester/registry

Examples¶

Example agents and environments can be found in the custom/ directory:

custom/agents/examples/simple_agent.py - Basic agent example
custom/agents/examples/advanced_agent.py - Agent with memory and emotions
custom/envs/examples/simple_env.py - Counter environment
custom/envs/examples/advanced_env.py - Resource management environment

These examples demonstrate best practices for creating custom modules.

Configuration¶

Set the WORKSPACE_PATH environment variable to point to your workspace:

export WORKSPACE_PATH=/path/to/workspace

Or add to your .env file:

WORKSPACE_PATH=/path/to/workspace

This setting tells the system where to find the custom/ directory.

Best Practices¶

Naming Conventions

Agent class names should end with Agent
Environment class names should end with Env
File names should use lowercase letters and underscores: my_agent.py

Error Handling

Return meaningful error messages
Keep necessary logs on critical paths to facilitate review and problem location.

State Management

Use restore() to restore state from the workspace and to_workspace() to write dynamic state back.
When a custom Agent overrides restore(), it should first call await super().restore(...) to bind the workspace, skill runtime, and service handles.
Record important state changes in replay
Keep state serializable (JSON compatible)

Tool Design

Use kind="observe" for read-only observations
Use kind="statistics" for aggregated data
Use kind=None and readonly=False for operations
After generation, priority is given to locating the failure reason through validation_report.json in the workflow product.