存储模块

本模块提供实验数据的存储与回放功能。

ReplayWriter

class agentsociety2.storage.ReplayWriter(db_path, *, enabled=True)

基类：ReplaySink

Back-compat alias over ReplaySink.

The legacy public API (agentsociety2.storage, examples) constructed a SQLite writer with a .db file path and called await init() before use. That writer is gone; this alias keeps the symbol importable and writes the same JSONL format, mapping the legacy db_path to a replay directory (a trailing .db is stripped) so old call sites keep working without SQLite.

__init__(db_path, *, enabled=True)

ReplaySink

class agentsociety2.storage.ReplaySink(replay_dir, *, enabled=True)

基类：object

Per-process, distributed, append-only JSONL replay writer.

One instance per writer process (driver / env router actor / each agent Ray task). Holds its own open shard fds; per-shard lock files make concurrent appends to the same shard file safe for rows of any size.

__init__(replay_dir, *, enabled=True)

async init()

No-op (kept for API compatibility with the legacy writer).

async write(table, data)

Append one row to table.

async write_batch(table, rows)

Append multiple rows to table (one flock per shard).

async register_table(schema)

Record a table's schema in _schema.json (idempotent).

async register_dataset(spec, columns)

Record a dataset + its columns in _schema.json (idempotent).

async close()

Close all open shard fds.

ReplayDatasetSpec

class agentsociety2.storage.ReplayDatasetSpec(dataset_id, table_name, module_name, kind, title='', description='', entity_key=None, step_key=None, time_key=None, default_order=<factory>, capabilities=<factory>, version=1)

Semantic metadata for a replay dataset backed by a JSONL table shard set.

dataset_id: str

table_name: str

module_name: str

kind: Literal['entity_snapshot', 'entity_static', 'env_snapshot', 'event_stream', 'metric_series']

title: str = ''

description: str = ''

entity_key: str | None = None

step_key: str | None = None

time_key: str | None = None

default_order: list[str]

capabilities: list[str]

version: int = 1

__init__(dataset_id, table_name, module_name, kind, title='', description='', entity_key=None, step_key=None, time_key=None, default_order=<factory>, capabilities=<factory>, version=1)

ReplayReader

class agentsociety2.storage.ReplayReader(replay_dir)

Read replay _schema.json + sharded JSONL rows via DuckDB.

__init__(replay_dir)

close()

load_dataset_catalog()

Return dataset metadata using the backend's legacy dict shape.

get_dataset_by_id(dataset_id)

find_dataset_by_capability(capability, *, kind=None)

normalize_row(dataset, row)

fetch_dataset_rows(dataset, *, order_by=None, desc=False, step=None, entity_id=None, start_step=None, end_step=None, max_step=None, columns=None, latest_per_entity=False, limit=None, offset=0)

count_dataset_rows(dataset, *, step=None, entity_id=None, start_step=None, end_step=None, max_step=None, latest_per_entity=False)

query_dataset_rows(dataset, *, page, page_size, order_by=None, desc=False, step=None, entity_id=None, start_step=None, end_step=None, max_step=None, columns=None, latest_per_entity=False)

distinct_values(dataset, column, *, order=True)

count_distinct(dataset, column)

min_value(dataset, column)

max_value(dataset, column)

ColumnDef

class agentsociety2.storage.ColumnDef(name, type, nullable=True, default=None, title=None, description=None, logical_type=None, analysis_role=None, unit=None, enum_values=None, example=None, tags=<factory>)

表列定义（SQLite）。

字段说明：

• name：列名。

• type：SQLite 列类型字符串，取值见 ColumnType。

• nullable：是否允许 NULL。

• default：默认值表达式，例如 CURRENT_TIMESTAMP。

• title：可选，人类可读的列标题。

• description：可选，语义描述，供回放、导出与分析使用。

• logical_type：可选，逻辑类型，例如 geo.lng、money。

• analysis_role：可选，分析角色，例如 measure。

• unit：可选，单位字符串，供分析与报告使用。

• enum_values：可选，离散列的枚举值列表。

• example：可选，示例值。

• tags：可选，自由标签列表。

name: str

type: Literal['INTEGER', 'REAL', 'TEXT', 'BLOB', 'TIMESTAMP', 'JSON']

nullable: bool = True

default: str | None = None

title: str | None = None

description: str | None = None

logical_type: str | None = None

analysis_role: str | None = None

unit: str | None = None

enum_values: list[Any] | None = None

example: Any | None = None

tags: list[str]

to_sql()

返回:

列定义的 SQL 片段。

返回类型:

str

__init__(name, type, nullable=True, default=None, title=None, description=None, logical_type=None, analysis_role=None, unit=None, enum_values=None, example=None, tags=<factory>)

TableSchema

class agentsociety2.storage.TableSchema(name, columns, primary_key=<factory>, indexes=<factory>)

数据库表结构定义（用于动态建表）。

参数:

• name (str) -- 表名。

• columns (List[ColumnDef]) -- 列定义列表。

• primary_key (List[str]) -- 主键列名列表。

• indexes (List[List[str]]) -- 索引定义列表（每项为列名列表）。

name: str

columns: List[ColumnDef]

primary_key: List[str]

indexes: List[List[str]]

to_create_sql()

返回:

CREATE TABLE SQL 语句。

返回类型:

str

to_index_sql()

返回:

CREATE INDEX SQL 语句列表。

返回类型:

List[str]

__init__(name, columns, primary_key=<factory>, indexes=<factory>)

分布式 Replay Proxy

大规模仿真下，ReplayProxy 只携带 replay_dir 与启用标志。各 agent / env / society 进程收到 proxy 后在本地懒加载 ReplaySink，并直接 append sharded JSONL。

ReplayProxy

class agentsociety2.storage.replay_proxy.ReplayProxy(replay_dir=None, enabled=True, _sink=None)

Serializable config for distributed replay writing.

Carries only the replay output directory and an enable flag — no Ray actor handle. Methods lazily build+cache a process-local ReplaySink (from replay_dir) on first use and delegate, so each agent Ray task / env actor that receives a deserialized copy gets its own sink and its own file descriptors. enabled=False makes every method a no-op.

replay_dir: str | None = None

enabled: bool = True

async write(table, data)

async write_batch(table, rows)

async register_table(schema)

async register_dataset(spec, columns)

async close()

__init__(replay_dir=None, enabled=True, _sink=None)

agentsociety2.storage.build_replay_sink(replay_proxy)

Build a per-process ReplaySink from a ReplayProxy.

Returns None when the proxy is disabled or carries no replay_dir.

兼容数据模型

以下模型仅用于兼容读取历史 SQLite 数据库；新实验默认不再写入这些 agent 表。

AgentProfile

class agentsociety2.storage.models.AgentProfile(*, id, name, profile={}, created_at=<factory>)

agent 档案信息（框架表）。

id: int

name: str

profile: Dict[str, Any]

created_at: datetime

__init__(**data)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

model_config = {'from_attributes': True, 'read_from_attributes': True, 'read_with_orm_mode': True, 'registry': PydanticUndefined, 'table': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

AgentStatus

class agentsociety2.storage.models.AgentStatus(*, id, step, t, action=None, status=None, created_at=<factory>)

agent 在某一步的状态快照（框架表）。

id: int

step: int

t: datetime

action: str | None

status: Dict[str, Any] | None

created_at: datetime

__init__(**data)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

model_config = {'from_attributes': True, 'read_from_attributes': True, 'read_with_orm_mode': True, 'registry': PydanticUndefined, 'table': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

AgentDialog

class agentsociety2.storage.models.AgentDialog(*, id=None, agent_id, step, t, type, speaker, content, created_at=<factory>)

agent 对话记录（框架表）。

id: int | None

agent_id: int

step: int

t: datetime

speaker: str

content: str

created_at: datetime

__init__(**data)

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

model_config = {'from_attributes': True, 'read_from_attributes': True, 'read_with_orm_mode': True, 'registry': PydanticUndefined, 'table': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].