Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion abevalflow/db/__init__.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,29 @@
"""Database models and engine for A/B evaluation results persistence."""

from abevalflow.db.engine import Session, get_engine, init_db
from abevalflow.db.models import Base, EvaluationRun, Trial
from abevalflow.db.models import (
Base,
CertificationRow,
EvaluationRun,
GateResultRow,
MCPCheckerRun,
MCPCheckerTask,
ObservabilityMetricsRow,
ScorecardRow,
SecurityScan,
Trial,
)

__all__ = [
"Base",
"CertificationRow",
"EvaluationRun",
"GateResultRow",
"MCPCheckerRun",
"MCPCheckerTask",
"ObservabilityMetricsRow",
"ScorecardRow",
"SecurityScan",
"Session",
"Trial",
"get_engine",
Expand Down
170 changes: 168 additions & 2 deletions abevalflow/db/models.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,14 @@
"""SQLAlchemy models for evaluation results persistence.

Three tables:
Tables:
- ``evaluation_runs``: one row per pipeline run (flattened summary for fast queries)
- ``trials``: one row per trial (drill-down into individual outcomes)
- ``security_scans``: one row per security scan (independent of evaluation runs)
- ``mcpchecker_results`` / ``mcpchecker_tasks``: MCP server evaluation results
- ``scorecards``: one row per pipeline run (unified gate verdict)
- ``gate_results``: one row per gate per run (normalized gate details)
- ``certifications``: one row per certification level per run (max 3)
- ``observability_metrics``: token/timing metrics per run
"""

from __future__ import annotations
Expand All @@ -12,14 +17,17 @@
from datetime import UTC, datetime

from sqlalchemy import (
BigInteger,
Boolean,
DateTime,
Float,
ForeignKey,
Index,
Integer,
Numeric,
String,
Text,
UniqueConstraint,
Uuid,
)
from sqlalchemy.dialects import postgresql
Expand Down Expand Up @@ -56,7 +64,7 @@ class EvaluationRun(Base):
eval_engine: Mapped[str] = mapped_column(String(10), nullable=False, default="harbor")

# Summary
recommendation: Mapped[str] = mapped_column(String(10), nullable=False)
recommendation: Mapped[str] = mapped_column(String(20), nullable=False)
uplift: Mapped[float] = mapped_column(Float, nullable=False)
mean_reward_gap: Mapped[float | None] = mapped_column(Float)
ttest_p_value: Mapped[float | None] = mapped_column(Float)
Expand Down Expand Up @@ -253,3 +261,161 @@ def __repr__(self) -> str:
f"<MCPCheckerTask {self.task_id!r} status={self.status!r} "
f"judge={self.llm_judge_passed}/{self.llm_judge_total}>"
)


class ScorecardRow(Base):
"""One row per pipeline run with unified gate verdict.

scorecard_json duplicates data normalized into gate_results and certifications.
Kept intentionally for fast single-row API responses without JOINs and for
preserving the exact original structure for debugging. Revisit after 6 months.
"""

__tablename__ = "scorecards"

id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4)
pipeline_run_id: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
submission_name: Mapped[str] = mapped_column(String(255), nullable=False)
eval_engine: Mapped[str] = mapped_column(String(50), nullable=False)

recommendation: Mapped[str] = mapped_column(String(20), nullable=False)
recommendation_reason: Mapped[str] = mapped_column(Text, nullable=False)
combination_mode: Mapped[str] = mapped_column(String(20), nullable=False)

gates_passed: Mapped[int] = mapped_column(Integer, nullable=False)
gates_failed: Mapped[int] = mapped_column(Integer, nullable=False)
blocking_gates_passed: Mapped[int] = mapped_column(Integer, nullable=False)
blocking_gates_failed: Mapped[int] = mapped_column(Integer, nullable=False)

highest_certification: Mapped[str] = mapped_column(String(20), nullable=False)
scorecard_json: Mapped[dict] = mapped_column(_JsonVariant, nullable=False)

created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=_utcnow)

gate_results: Mapped[list[GateResultRow]] = relationship(back_populates="scorecard", cascade="all, delete-orphan")
certifications: Mapped[list[CertificationRow]] = relationship(
back_populates="scorecard", cascade="all, delete-orphan"
)

__table_args__ = (
Index("ix_scorecards_submission_name", "submission_name"),
Index("ix_scorecards_submission_created", "submission_name", "created_at"),
Index("ix_scorecards_highest_certification", "highest_certification"),
)

def __repr__(self) -> str:
return f"<ScorecardRow {self.submission_name!r} recommendation={self.recommendation!r}>"


class GateResultRow(Base):
"""One row per gate per run for normalized gate queries."""

__tablename__ = "gate_results"

id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4)
scorecard_id: Mapped[uuid.UUID] = mapped_column(
Uuid, ForeignKey("scorecards.id", ondelete="CASCADE"), nullable=False
)

gate_name: Mapped[str] = mapped_column(String(100), nullable=False)
gate_type: Mapped[str] = mapped_column(String(20), nullable=False)
policy_key: Mapped[str | None] = mapped_column(String(100))
passed: Mapped[bool] = mapped_column(Boolean, nullable=False)
score: Mapped[float] = mapped_column(Float, nullable=False)
mode: Mapped[str] = mapped_column(String(20), nullable=False)
threshold: Mapped[float | None] = mapped_column(Float)
findings_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
message: Mapped[str | None] = mapped_column(Text)
details_json: Mapped[dict | None] = mapped_column(_JsonVariant)
duration_ms: Mapped[int | None] = mapped_column(Integer)
prompt_tokens: Mapped[int | None] = mapped_column(Integer)
completion_tokens: Mapped[int | None] = mapped_column(Integer)
evaluated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))

scorecard: Mapped[ScorecardRow] = relationship(back_populates="gate_results")

__table_args__ = (
Index("ix_gate_results_scorecard_id", "scorecard_id"),
Index("ix_gate_results_scorecard_policy", "scorecard_id", "policy_key"),
Index("ix_gate_results_gate_passed", "gate_name", "passed"),
Index("ix_gate_results_policy_key", "policy_key"),
)

def __repr__(self) -> str:
return f"<GateResultRow {self.gate_name!r} passed={self.passed!r} score={self.score!r}>"


class CertificationRow(Base):
"""One row per certification level per run (max 3 per scorecard)."""

__tablename__ = "certifications"

id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4)
scorecard_id: Mapped[uuid.UUID] = mapped_column(
Uuid, ForeignKey("scorecards.id", ondelete="CASCADE"), nullable=False
)

level: Mapped[str] = mapped_column(String(20), nullable=False)
passed: Mapped[bool] = mapped_column(Boolean, nullable=False)
checks_total: Mapped[int] = mapped_column(Integer, nullable=False)
checks_passed: Mapped[int] = mapped_column(Integer, nullable=False)
checks_failed: Mapped[int] = mapped_column(Integer, nullable=False)
failed_checks: Mapped[list | None] = mapped_column(_JsonVariant)
details_json: Mapped[dict | None] = mapped_column(_JsonVariant)

scorecard: Mapped[ScorecardRow] = relationship(back_populates="certifications")

__table_args__ = (
Index("ix_certifications_scorecard_id", "scorecard_id"),
Index("ix_certifications_level_passed", "level", "passed"),
UniqueConstraint("scorecard_id", "level", name="uq_certifications_scorecard_level"),
)

def __repr__(self) -> str:
return f"<CertificationRow {self.level!r} passed={self.passed!r}>"


class ObservabilityMetricsRow(Base):
"""Token/timing metrics per pipeline run.

Schema-only in Phase 0. Populated by MetricsContext in Phase A.
No FK to scorecards — metrics may outlive scorecard lifecycle.
"""

__tablename__ = "observability_metrics"

id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4)
pipeline_run_id: Mapped[str] = mapped_column(String(255), nullable=False)
submission_name: Mapped[str] = mapped_column(String(255), nullable=False)

model_name: Mapped[str | None] = mapped_column(String(100))
pipeline_duration_ms: Mapped[int | None] = mapped_column(Integer)
prepare_duration_ms: Mapped[int | None] = mapped_column(Integer)
test_duration_ms: Mapped[int | None] = mapped_column(Integer)
evaluate_duration_ms: Mapped[int | None] = mapped_column(Integer)
analyze_duration_ms: Mapped[int | None] = mapped_column(Integer)
store_duration_ms: Mapped[int | None] = mapped_column(Integer)

total_prompt_tokens: Mapped[int | None] = mapped_column(BigInteger)
total_completion_tokens: Mapped[int | None] = mapped_column(BigInteger)
total_tokens: Mapped[int | None] = mapped_column(BigInteger)
estimated_cost_usd: Mapped[float | None] = mapped_column(Numeric(12, 6))

llm_calls_count: Mapped[int | None] = mapped_column(Integer)
trials_count: Mapped[int | None] = mapped_column(Integer)
attempt_number: Mapped[int] = mapped_column(Integer, nullable=False, default=1)

created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=_utcnow)

__table_args__ = (
UniqueConstraint("pipeline_run_id", "attempt_number", name="uq_obs_metrics_run_attempt"),
Index("ix_observability_metrics_submission_name", "submission_name"),
Index("ix_observability_metrics_created_at", "created_at"),
Index("ix_observability_metrics_model_name", "model_name"),
)

def __repr__(self) -> str:
return (
f"<ObservabilityMetricsRow {self.submission_name!r} "
f"tokens={self.total_tokens} cost={self.estimated_cost_usd}>"
)
39 changes: 39 additions & 0 deletions alembic.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
[alembic]
script_location = alembic
file_template = %%(rev)s_%%(slug)s
# Override with DATABASE_URL env var (see alembic/env.py). This default
# is a placeholder and will fail without a running local PostgreSQL.
sqlalchemy.url = postgresql+psycopg://localhost/abevalflow_dev

[loggers]
keys = root,sqlalchemy,alembic

[handlers]
keys = console

[formatters]
keys = generic

[logger_root]
level = WARN
handlers = console

[logger_sqlalchemy]
level = WARN
handlers =
qualname = sqlalchemy.engine

[logger_alembic]
level = INFO
handlers =
qualname = alembic

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatter_generic]
format = %(levelname)-5.5s [%(name)s] %(message)s
datefmt = %H:%M:%S
33 changes: 33 additions & 0 deletions alembic/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# Alembic Migrations

## Existing tables

The tables `evaluation_runs`, `trials`, `security_scans`, `mcpchecker_results`, and `mcpchecker_tasks` were created via `Base.metadata.create_all()` before Alembic was introduced. They are not managed by Alembic migrations.

## For existing databases

On a database that already has the pre-Alembic tables, run:

```bash
DATABASE_URL=postgresql+psycopg://... alembic upgrade head
```

This creates only the new observability tables and stamps the `alembic_version` table.

## For fresh databases

On a completely new database, create the pre-Alembic tables first, then run migrations:

```bash
DATABASE_URL=postgresql+psycopg://... python -c "
from abevalflow.db.engine import get_engine, init_db
init_db(get_engine())
"
DATABASE_URL=postgresql+psycopg://... alembic stamp head
```

Alternatively, `store_results.py` calls `init_db()` on every run, which creates all tables (including the new ones) via `Base.metadata.create_all()`. Alembic is needed only for schema changes to existing tables in future.

## Revision ID convention

This project uses sequential string IDs (`001`, `002`, etc.) instead of Alembic's default random hashes. This is intentional for readability. To avoid conflicts when multiple branches add migrations, coordinate revision numbers via the PR process.
40 changes: 40 additions & 0 deletions alembic/env.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
"""Alembic environment configuration."""

import os
from logging.config import fileConfig

from alembic import context
from sqlalchemy import create_engine

from abevalflow.db.models import Base

config = context.config
if config.config_file_name is not None:
fileConfig(config.config_file_name)

target_metadata = Base.metadata

db_url = os.environ.get("DATABASE_URL")
if db_url:
config.set_main_option("sqlalchemy.url", db_url)


def run_migrations_offline() -> None:
url = config.get_main_option("sqlalchemy.url")
context.configure(url=url, target_metadata=target_metadata, literal_binds=True)
with context.begin_transaction():
context.run_migrations()


def run_migrations_online() -> None:
connectable = create_engine(config.get_main_option("sqlalchemy.url"))
with connectable.connect() as connection:
context.configure(connection=connection, target_metadata=target_metadata)
with context.begin_transaction():
context.run_migrations()


if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()
25 changes: 25 additions & 0 deletions alembic/script.py.mako
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
"""${message}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
${imports if imports else ""}

# revision identifiers, used by Alembic.
revision: str = ${repr(up_revision)}
down_revision: Union[str, None] = ${repr(down_revision)}
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}


def upgrade() -> None:
${upgrades if upgrades else "pass"}


def downgrade() -> None:
${downgrades if downgrades else "pass"}
Loading
Loading