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