Storage Module

This module provides storage and playback functions of experimental data.

ReplayWriter

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

Bases: 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)

ReplayWriter

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

Bases: 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 SQLite table.

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)

ReplayWriter

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>)

Table column definition (SQLite).

Field description:

• name: Column name.

• type: SQLite column type string. See ColumnType for available values.

• nullable: Whether NULL is allowed.

• default: Default value expression, such as CURRENT_TIMESTAMP.

• title: Optional human-readable column title.

• description: Optional semantic description used for replay, export, and analysis.

• logical_type: Optional logical type, such as geo.lng or money.

• analysis_role: Optional analysis role, such as measure.

• unit: Optional unit string used by analysis and reporting.

• enum_values: Optional list of enumeration values for discrete columns.

• example: Optional example value.

• tags: Optional free-form tag list.

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()

Returns:

SQL fragment for column definition.

Return type:

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>)

Database table structure definition (used for dynamic table creation).

Parameters:

• name (str) – Table name.

• columns (List[ColumnDef]) – List of column definitions.

• primary_key (List[str]) – List of primary key column names.

• indexes (List[List[str]]) – A list of index definitions (each item is a list of column names).

name: str

columns: List[ColumnDef]

primary_key: List[str]

indexes: List[List[str]]

to_create_sql()

Returns:

CREATE TABLE SQL statement.

Return type:

str

to_index_sql()

Returns:

CREATE INDEX List of SQL statements.

Return type:

List[str]

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

Distributed Replay Proxy

In large-scale simulations, ReplayProxy only carries replay_dir and the enabled flag. After each agent / env / society process receives the proxy, it lazily loads a local ReplaySink and appends sharded JSONL directly.

ReplayWriter

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.

Compatible data model

The following models are only compatible with reading historical SQLite databases; new experiments no longer write to these agent tables by default.

AgentProfile

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

agent profile information (frame table).

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>)

A snapshot of the agent’s state at a certain step (framework table).

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 conversation record (frame table).

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].