You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
302 lines
15 KiB
302 lines
15 KiB
from __future__ import annotations
|
|
import json
|
|
import os
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
import aiosqlite
|
|
|
|
from app.models.agent import AgentPersona
|
|
from app.models.action import ActionEntry
|
|
from app.models.world import WorldMetrics, WorldState
|
|
from app.models.simulation import SimulationState, SimulationStatus
|
|
|
|
|
|
class SimulationStore:
|
|
def __init__(self, data_dir: str = "data"):
|
|
self.data_dir = data_dir
|
|
self.global_db_path = os.path.join(data_dir, "global.db")
|
|
|
|
async def init(self):
|
|
os.makedirs(self.data_dir, exist_ok=True)
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS simulations (
|
|
id TEXT PRIMARY KEY,
|
|
status TEXT NOT NULL DEFAULT 'generating_world',
|
|
rules_text TEXT NOT NULL DEFAULT '',
|
|
world_name TEXT NOT NULL DEFAULT '',
|
|
agent_count INTEGER NOT NULL DEFAULT 0,
|
|
created_at TEXT NOT NULL,
|
|
is_public INTEGER NOT NULL DEFAULT 0,
|
|
view_count INTEGER NOT NULL DEFAULT 0,
|
|
fork_count INTEGER NOT NULL DEFAULT 0,
|
|
forked_from TEXT
|
|
)
|
|
""")
|
|
await db.commit()
|
|
|
|
def _sim_dir(self, sim_id: str) -> str:
|
|
return os.path.join(self.data_dir, "simulations", sim_id)
|
|
|
|
def _sim_db_path(self, sim_id: str) -> str:
|
|
return os.path.join(self._sim_dir(sim_id), "simulation.db")
|
|
|
|
async def _init_sim_db(self, sim_id: str):
|
|
sim_dir = self._sim_dir(sim_id)
|
|
os.makedirs(sim_dir, exist_ok=True)
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute("CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT)")
|
|
await db.execute("CREATE TABLE IF NOT EXISTS agents (agent_id INTEGER PRIMARY KEY, state TEXT, updated_at TEXT)")
|
|
await db.execute("CREATE TABLE IF NOT EXISTS world_state (round INTEGER PRIMARY KEY, state TEXT)")
|
|
await db.execute("CREATE TABLE IF NOT EXISTS actions (id INTEGER PRIMARY KEY AUTOINCREMENT, round INTEGER, day INTEGER, entry TEXT)")
|
|
await db.execute("CREATE TABLE IF NOT EXISTS narratives (round INTEGER PRIMARY KEY, day INTEGER, time_of_day TEXT, text TEXT)")
|
|
await db.execute("CREATE TABLE IF NOT EXISTS metrics_history (round INTEGER PRIMARY KEY, metrics TEXT)")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_actions_round ON actions(round)")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_narratives_day ON narratives(day)")
|
|
await db.execute("CREATE TABLE IF NOT EXISTS report (id INTEGER PRIMARY KEY CHECK (id = 1), status TEXT DEFAULT 'pending', report_json TEXT, error TEXT)")
|
|
await db.commit()
|
|
|
|
async def create(self, sim_id: str, rules_text: str) -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
await db.execute(
|
|
"INSERT INTO simulations (id, rules_text, created_at) VALUES (?, ?, ?)",
|
|
(sim_id, rules_text, now),
|
|
)
|
|
await db.commit()
|
|
await self._init_sim_db(sim_id)
|
|
|
|
async def update_status(self, sim_id: str, status: str, **kwargs) -> None:
|
|
sets = ["status = ?"]
|
|
vals: list[Any] = [status]
|
|
for k, v in kwargs.items():
|
|
sets.append(f"{k} = ?")
|
|
vals.append(v)
|
|
vals.append(sim_id)
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
await db.execute(f"UPDATE simulations SET {', '.join(sets)} WHERE id = ?", vals)
|
|
await db.commit()
|
|
|
|
async def get_simulation(self, sim_id: str) -> dict | None:
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
cursor = await db.execute("SELECT * FROM simulations WHERE id = ?", (sim_id,))
|
|
row = await cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
async def list_simulations(self, public_only: bool = False, sort: str = "recent") -> list[dict]:
|
|
query = "SELECT * FROM simulations"
|
|
if public_only:
|
|
query += " WHERE is_public = 1"
|
|
order = {"recent": "created_at DESC", "views": "view_count DESC", "forks": "fork_count DESC"}
|
|
query += f" ORDER BY {order.get(sort, 'created_at DESC')}"
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
cursor = await db.execute(query)
|
|
return [dict(row) for row in await cursor.fetchall()]
|
|
|
|
async def set_meta(self, sim_id: str, key: str, value: Any) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)",
|
|
(key, json.dumps(value, default=str)),
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_meta(self, sim_id: str, key: str) -> Any:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
cursor = await db.execute("SELECT value FROM meta WHERE key = ?", (key,))
|
|
row = await cursor.fetchone()
|
|
return json.loads(row[0]) if row else None
|
|
|
|
async def save_agent(self, sim_id: str, agent: AgentPersona) -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT OR REPLACE INTO agents (agent_id, state, updated_at) VALUES (?, ?, ?)",
|
|
(agent.id, agent.model_dump_json(), now),
|
|
)
|
|
await db.commit()
|
|
|
|
async def save_agents_batch(self, sim_id: str, agents: list[AgentPersona]) -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.executemany(
|
|
"INSERT OR REPLACE INTO agents (agent_id, state, updated_at) VALUES (?, ?, ?)",
|
|
[(a.id, a.model_dump_json(), now) for a in agents],
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_agent(self, sim_id: str, agent_id: int) -> AgentPersona | None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
cursor = await db.execute("SELECT state FROM agents WHERE agent_id = ?", (agent_id,))
|
|
row = await cursor.fetchone()
|
|
return AgentPersona.model_validate_json(row[0]) if row else None
|
|
|
|
async def get_all_agents(self, sim_id: str) -> list[AgentPersona]:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
cursor = await db.execute("SELECT state FROM agents ORDER BY agent_id")
|
|
rows = await cursor.fetchall()
|
|
return [AgentPersona.model_validate_json(row[0]) for row in rows]
|
|
|
|
async def save_world_state(self, sim_id: str, round_num: int, state: WorldState) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT OR REPLACE INTO world_state (round, state) VALUES (?, ?)",
|
|
(round_num, state.model_dump_json()),
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_world_state(self, sim_id: str, round_num: int | None = None) -> WorldState | None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
if round_num is not None:
|
|
cursor = await db.execute("SELECT state FROM world_state WHERE round = ?", (round_num,))
|
|
else:
|
|
cursor = await db.execute("SELECT state FROM world_state ORDER BY round DESC LIMIT 1")
|
|
row = await cursor.fetchone()
|
|
return WorldState.model_validate_json(row[0]) if row else None
|
|
|
|
async def save_action(self, sim_id: str, entry: ActionEntry) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT INTO actions (round, day, entry) VALUES (?, ?, ?)",
|
|
(entry.round, entry.day, entry.model_dump_json()),
|
|
)
|
|
await db.commit()
|
|
|
|
async def save_actions_batch(self, sim_id: str, entries: list[ActionEntry]) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.executemany(
|
|
"INSERT INTO actions (round, day, entry) VALUES (?, ?, ?)",
|
|
[(e.round, e.day, e.model_dump_json()) for e in entries],
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_actions(self, sim_id: str, from_round: int = 0, to_round: int | None = None) -> list[ActionEntry]:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
if to_round is not None:
|
|
cursor = await db.execute(
|
|
"SELECT entry FROM actions WHERE round >= ? AND round <= ? ORDER BY id",
|
|
(from_round, to_round),
|
|
)
|
|
else:
|
|
cursor = await db.execute(
|
|
"SELECT entry FROM actions WHERE round >= ? ORDER BY id", (from_round,)
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [ActionEntry.model_validate_json(row[0]) for row in rows]
|
|
|
|
async def save_narrative(self, sim_id: str, round_num: int, day: int, time_of_day: str, text: str) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT OR REPLACE INTO narratives (round, day, time_of_day, text) VALUES (?, ?, ?, ?)",
|
|
(round_num, day, time_of_day, text),
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_narratives(self, sim_id: str, from_day: int = 0, to_day: int | None = None) -> list[dict]:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
if to_day is not None:
|
|
cursor = await db.execute(
|
|
"SELECT * FROM narratives WHERE day >= ? AND day <= ? ORDER BY round",
|
|
(from_day, to_day),
|
|
)
|
|
else:
|
|
cursor = await db.execute(
|
|
"SELECT * FROM narratives WHERE day >= ? ORDER BY round", (from_day,)
|
|
)
|
|
return [dict(row) for row in await cursor.fetchall()]
|
|
|
|
async def save_metrics(self, sim_id: str, round_num: int, metrics: WorldMetrics) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT OR REPLACE INTO metrics_history (round, metrics) VALUES (?, ?)",
|
|
(round_num, metrics.model_dump_json()),
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_metrics_history(self, sim_id: str) -> list[dict]:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
cursor = await db.execute("SELECT round, metrics FROM metrics_history ORDER BY round")
|
|
rows = await cursor.fetchall()
|
|
return [{"round": row[0], **json.loads(row[1])} for row in rows]
|
|
|
|
async def increment_view_count(self, sim_id: str) -> None:
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
await db.execute("UPDATE simulations SET view_count = view_count + 1 WHERE id = ?", (sim_id,))
|
|
await db.commit()
|
|
|
|
async def increment_fork_count(self, sim_id: str) -> None:
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
await db.execute("UPDATE simulations SET fork_count = fork_count + 1 WHERE id = ?", (sim_id,))
|
|
await db.commit()
|
|
|
|
async def publish(self, sim_id: str) -> None:
|
|
async with aiosqlite.connect(self.global_db_path) as db:
|
|
await db.execute("UPDATE simulations SET is_public = 1 WHERE id = ?", (sim_id,))
|
|
await db.commit()
|
|
|
|
async def copy_state_at_round(self, source_sim_id: str, target_sim_id: str, round_num: int) -> None:
|
|
await self._init_sim_db(target_sim_id)
|
|
src_db_path = self._sim_db_path(source_sim_id)
|
|
tgt_db_path = self._sim_db_path(target_sim_id)
|
|
|
|
async with aiosqlite.connect(src_db_path) as src_db:
|
|
async with aiosqlite.connect(tgt_db_path) as tgt_db:
|
|
cursor = await src_db.execute("SELECT key, value FROM meta")
|
|
rows = await cursor.fetchall()
|
|
await tgt_db.executemany("INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)", rows)
|
|
|
|
cursor = await src_db.execute("SELECT agent_id, state, updated_at FROM agents")
|
|
rows = await cursor.fetchall()
|
|
await tgt_db.executemany("INSERT OR REPLACE INTO agents (agent_id, state, updated_at) VALUES (?, ?, ?)", rows)
|
|
|
|
cursor = await src_db.execute("SELECT round, state FROM world_state WHERE round <= ?", (round_num,))
|
|
rows = await cursor.fetchall()
|
|
await tgt_db.executemany("INSERT OR REPLACE INTO world_state (round, state) VALUES (?, ?)", rows)
|
|
|
|
cursor = await src_db.execute("SELECT round, day, entry FROM actions WHERE round <= ?", (round_num,))
|
|
rows = await cursor.fetchall()
|
|
await tgt_db.executemany("INSERT INTO actions (round, day, entry) VALUES (?, ?, ?)", rows)
|
|
|
|
cursor = await src_db.execute("SELECT round, day, time_of_day, text FROM narratives WHERE round <= ?", (round_num,))
|
|
rows = await cursor.fetchall()
|
|
await tgt_db.executemany("INSERT OR REPLACE INTO narratives (round, day, time_of_day, text) VALUES (?, ?, ?, ?)", rows)
|
|
|
|
cursor = await src_db.execute("SELECT round, metrics FROM metrics_history WHERE round <= ?", (round_num,))
|
|
rows = await cursor.fetchall()
|
|
await tgt_db.executemany("INSERT OR REPLACE INTO metrics_history (round, metrics) VALUES (?, ?)", rows)
|
|
|
|
await tgt_db.commit()
|
|
|
|
async def set_report_status(self, sim_id: str, status: str, error: str | None = None) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT OR REPLACE INTO report (id, status, error) VALUES (1, ?, ?)",
|
|
(status, error),
|
|
)
|
|
await db.commit()
|
|
|
|
async def save_report(self, sim_id: str, report: dict) -> None:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
await db.execute(
|
|
"INSERT OR REPLACE INTO report (id, status, report_json) VALUES (1, 'ready', ?)",
|
|
(json.dumps(report, default=str),),
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_report(self, sim_id: str) -> dict | None:
|
|
try:
|
|
async with aiosqlite.connect(self._sim_db_path(sim_id)) as db:
|
|
cursor = await db.execute("SELECT status, report_json, error FROM report WHERE id = 1")
|
|
row = await cursor.fetchone()
|
|
if not row:
|
|
return {"status": "pending"}
|
|
status, report_json, error = row
|
|
if status == "ready" and report_json:
|
|
return {"status": "ready", "report": json.loads(report_json)}
|
|
return {"status": status, "error": error}
|
|
except Exception:
|
|
return {"status": "pending"}
|
|
|