agentsociety.storage.database

Module Contents

Classes

DatabaseConfig

Database configuration class supporting multiple database types.

DatabaseWriter

Functions

_create_async_engine_from_config

_create_tables

Create tables using SQLAlchemy

Data

__all__

API

agentsociety.storage.database.__all__

[‘DatabaseWriter’, ‘DatabaseConfig’]

class agentsociety.storage.database.DatabaseConfig

Bases: pydantic.BaseModel

Database configuration class supporting multiple database types.

enabled: bool

‘Field(…)’

Whether database storage is enabled

db_type: Literal[postgresql, sqlite]

‘Field(…)’

Database type

pg_dsn: Optional[str]

‘Field(…)’

Database connection string (PostgreSQL)

validate_config()
get_dsn(sqlite_path: pathlib.Path)

Create async SQLAlchemy engine based on configuration

agentsociety.storage.database._create_async_engine_from_config(config: agentsociety.storage.database.DatabaseConfig, sqlite_path: pathlib.Path)
async agentsociety.storage.database._create_tables(exp_id: str, config: agentsociety.storage.database.DatabaseConfig, sqlite_path: pathlib.Path)

Create tables using SQLAlchemy

class agentsociety.storage.database.DatabaseWriter(tenant_id: str, exp_id: str, config: agentsociety.storage.database.DatabaseConfig, home_dir: str)

Initialization

Initialize database writer.

  • Args:

    • tenant_id (str): Tenant ID.

    • exp_id (str): Experiment ID.

    • config (DatabaseConfig): Database configuration.

    • home_dir (str): Home directory. sqlite will be stored in home_dir/sqlite.db

async init()

Initialize database tables

_init_tables()

Initialize table object cache

async _create_tables()

Create tables

_get_insert_func()

Get insert function based on database type

property exp_info_file

Experiment info file path

property storage_path

Storage path

async read_dialogs(day: Optional[int] = None, speaker: Optional[str] = None, dialog_type: Optional[int] = None, start_t: Optional[float] = None, end_t: Optional[float] = None, limit: Optional[int] = None, offset: Optional[int] = None, order_by: str = 'created_at', order_direction: str = 'asc') → List[Dict[str, Any]]

Read dialog records with filtering and pagination.

  • Args:

    • day (Optional[int]): Filter by day.

    • speaker (Optional[str]): Filter by speaker.

    • dialog_type (Optional[int]): Filter by dialog type.

    • start_t (Optional[float]): Filter by start time.

    • end_t (Optional[float]): Filter by end time.

    • limit (Optional[int]): Limit number of records.

    • offset (Optional[int]): Offset for pagination.

    • order_by (str): Column to order by.

    • order_direction (str): Order direction (‘asc’ or ‘desc’).

  • Returns:

    • List[Dict[str, Any]]: List of dialog records.

async read_statuses(day: Optional[int] = None, agent_id: Optional[int] = None, start_t: Optional[float] = None, end_t: Optional[float] = None, action: Optional[str] = None, limit: Optional[int] = None, offset: Optional[int] = None, order_by: str = 'created_at', order_direction: str = 'asc') → List[Dict[str, Any]]

Read status records with filtering and pagination.

  • Args:

    • day (Optional[int]): Filter by day.

    • agent_id (Optional[int]): Filter by agent ID.

    • start_t (Optional[float]): Filter by start time.

    • end_t (Optional[float]): Filter by end time.

    • action (Optional[str]): Filter by action.

    • limit (Optional[int]): Limit number of records.

    • offset (Optional[int]): Offset for pagination.

    • order_by (str): Column to order by.

    • order_direction (str): Order direction (‘asc’ or ‘desc’).

  • Returns:

    • List[Dict[str, Any]]: List of status records.

async read_surveys(day: Optional[int] = None, survey_id: Optional[str] = None, start_t: Optional[float] = None, end_t: Optional[float] = None, limit: Optional[int] = None, offset: Optional[int] = None, order_by: str = 'created_at', order_direction: str = 'asc') → List[Dict[str, Any]]

Read survey records with filtering and pagination.

  • Args:

    • day (Optional[int]): Filter by day.

    • survey_id (Optional[str]): Filter by survey ID.

    • start_t (Optional[float]): Filter by start time.

    • end_t (Optional[float]): Filter by end time.

    • limit (Optional[int]): Limit number of records.

    • offset (Optional[int]): Offset for pagination.

    • order_by (str): Column to order by.

    • order_direction (str): Order direction (‘asc’ or ‘desc’).

  • Returns:

    • List[Dict[str, Any]]: List of survey records.

async read_profiles() → List[Dict[str, Any]]

Read all agent profiles.

  • Returns:

    • List[Dict[str, Any]]: List of agent profile records.

async read_global_prompts(day: Optional[int] = None, start_t: Optional[float] = None, end_t: Optional[float] = None, limit: Optional[int] = None, offset: Optional[int] = None, order_by: str = 'created_at', order_direction: str = 'asc') → List[Dict[str, Any]]

Read global prompt records with filtering and pagination.

  • Args:

    • day (Optional[int]): Filter by day.

    • start_t (Optional[float]): Filter by start time.

    • end_t (Optional[float]): Filter by end time.

    • limit (Optional[int]): Limit number of records.

    • offset (Optional[int]): Offset for pagination.

    • order_by (str): Column to order by.

    • order_direction (str): Order direction (‘asc’ or ‘desc’).

  • Returns:

    • List[Dict[str, Any]]: List of global prompt records.

async read_task_results(agent_id: Optional[int] = None, limit: Optional[int] = None, offset: Optional[int] = None, order_by: str = 'created_at', order_direction: str = 'asc') → List[Dict[str, Any]]

Read task result records with filtering and pagination.

  • Args:

    • agent_id (Optional[int]): Filter by agent ID.

    • limit (Optional[int]): Limit number of records.

    • offset (Optional[int]): Offset for pagination.

    • order_by (str): Column to order by.

    • order_direction (str): Order direction (‘asc’ or ‘desc’).

  • Returns:

    • List[Dict[str, Any]]: List of task result records.

async read_metrics(key: Optional[str] = None, step: Optional[int] = None, start_step: Optional[int] = None, end_step: Optional[int] = None, limit: Optional[int] = None, offset: Optional[int] = None, order_by: str = 'step', order_direction: str = 'asc') → List[Dict[str, Any]]

Read metric records with filtering and pagination.

  • Args:

    • key (Optional[str]): Filter by metric key.

    • step (Optional[int]): Filter by specific step.

    • start_step (Optional[int]): Filter by start step.

    • end_step (Optional[int]): Filter by end step.

    • limit (Optional[int]): Limit number of records.

    • offset (Optional[int]): Offset for pagination.

    • order_by (str): Column to order by.

    • order_direction (str): Order direction (‘asc’ or ‘desc’).

  • Returns:

    • List[Dict[str, Any]]: List of metric records.

async get_statistics() → Dict[str, Any]

Get comprehensive statistics for the experiment.

  • Returns:

    • Dict[str, Any]: Statistics including counts, time ranges, and summaries.

async write_dialogs(rows: list[agentsociety.storage.type.StorageDialog])
async write_statuses(rows: list[agentsociety.storage.type.StorageStatus])
async write_profiles(rows: list[agentsociety.storage.type.StorageProfile])
async write_surveys(rows: list[agentsociety.storage.type.StorageSurvey])
async write_global_prompt(prompt_info: agentsociety.storage.type.StorageGlobalPrompt)
async write_task_result(rows: list[agentsociety.storage.type.StorageTaskResult])
async log_metric(metrics: list[tuple[str, float, int]])

Batch insert metric data.

Args: metrics: List of tuples (key, value, step)

async update_exp_info(exp_info: agentsociety.storage.type.StorageExpInfo)
async fetch_pending_dialogs()

Fetch all unprocessed pending dialogs from the database.

  • Returns:

    • list[StoragePendingDialog]: List of pending dialogs.

async mark_dialogs_as_processed(pending_ids: list[int])

Mark specified dialogs as processed.

  • Args:

    • pending_ids (list[int]): List of pending dialog IDs to mark as processed.

async fetch_pending_surveys()

Fetch all unprocessed pending surveys from the database.

  • Returns:

    • list[StoragePendingSurvey]: List of pending surveys.

async mark_surveys_as_processed(pending_ids: list[int])

Mark specified surveys as processed.

  • Args:

    • pending_ids (list[int]): List of pending survey IDs to mark as processed.

async close()

Close database connection