File size: 2,297 Bytes
3b0d6b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import declarative_base, sessionmaker
DATABASE_URL = "sqlite:///./agent_arena.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
def create_all():
from models import Task, Agent, Run # noqa: F401
Base.metadata.create_all(bind=engine)
_ensure_run_schema()
def _ensure_run_schema():
expected = {
"functional_success": "INTEGER DEFAULT 0",
"test_pass_rate": "FLOAT",
"alignment_score_llm": "FLOAT",
"alignment_score_human": "FLOAT",
"alignment_rubric_version": "VARCHAR(100)",
"quality_score": "FLOAT",
"agent_duration_s": "FLOAT DEFAULT 0.0",
"end_to_end_duration_s": "FLOAT DEFAULT 0.0",
"timeout_flag": "INTEGER DEFAULT 0",
"tool_failure_flag": "INTEGER DEFAULT 0",
"invalid_patch_flag": "INTEGER DEFAULT 0",
"task_family": "VARCHAR(100)",
"seed": "VARCHAR(100)",
"replicate_index": "INTEGER DEFAULT 0",
"experiment_id": "VARCHAR(100)",
"arm_id": "VARCHAR(100)",
}
inspector = inspect(engine)
columns = {col["name"] for col in inspector.get_columns("runs")}
missing = [(name, ddl) for name, ddl in expected.items() if name not in columns]
if not missing:
return
with engine.begin() as conn:
for name, ddl in missing:
conn.execute(text(f"ALTER TABLE runs ADD COLUMN {name} {ddl}"))
conn.execute(text(
"UPDATE runs SET functional_success = CASE "
"WHEN functional_success IS NULL OR functional_success = 0 THEN "
"CASE WHEN status = 'success' THEN 1 ELSE 0 END "
"ELSE functional_success END"
))
conn.execute(text(
"UPDATE runs SET agent_duration_s = COALESCE(agent_duration_s, duration_s), "
"end_to_end_duration_s = COALESCE(end_to_end_duration_s, duration_s), "
"task_family = COALESCE(task_family, ''), "
"replicate_index = COALESCE(replicate_index, 0)"
))
|