Skip to content

[LLM] Build A/B testing framework for prompts and models #430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gelluisaac merged 2 commits into from
Jun 29, 2026
+155 −236
Merged

[LLM] Build A/B testing framework for prompts and models #430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d3101d5
Merge branch 'main' into men
Menjay7 Jun 28, 2026
106f395
Merge branch 'main' into men
Menjay7 Jun 29, 2026
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
236 changes: 46 additions & 190 deletions astroml/db/schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -632,124 +632,26 @@ class ModelVersion(Base):
)


# ---------------------------------------------------------------------------
# A/B Testing Framework
# ---------------------------------------------------------------------------
# A/B Testing Framework
# ---------------------------------------------------------------------------

class Experiment(Base):
"""A/B test experiment for comparing models or prompts."""

__tablename__ = "experiments"

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
name: Mapped[str] = mapped_column(String(128), nullable=False, unique=True)
description: Mapped[Optional[str]] = mapped_column(Text)
experiment_type: Mapped[str] = mapped_column(String(32), nullable=False) # 'model', 'prompt'
status: Mapped[str] = mapped_column(String(32), nullable=False, server_default="draft")
traffic_allocation: Mapped[float] = mapped_column(Numeric, nullable=False, server_default="1.0")
start_at: Mapped[Optional[datetime]] = mapped_column()
end_at: Mapped[Optional[datetime]] = mapped_column()
created_at: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(
nullable=False, server_default=func.now(), onupdate=func.now()
)

# Relationships
variants: Mapped[list[Variant]] = relationship(
back_populates="experiment",
cascade="all, delete-orphan",
)

__table_args__ = (
Index("ix_experiments_type", "experiment_type"),
Index("ix_experiments_status", "status"),
Index("ix_experiments_start_at", "start_at"),
CheckConstraint(
"experiment_type IN ('model', 'prompt')",
name="ck_experiments_type",
),
CheckConstraint(
"status IN ('draft', 'running', 'paused', 'completed', 'archived')",
name="ck_experiments_status",
),
CheckConstraint(
"traffic_allocation >= 0 AND traffic_allocation <= 1",
name="ck_experiments_traffic_allocation",
),
)

# ... keep full definition

class Variant(Base):
"""A variant in an A/B test experiment."""

__tablename__ = "variants"

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
experiment_id: Mapped[int] = mapped_column(
Integer, ForeignKey("experiments.id"), nullable=False
)
name: Mapped[str] = mapped_column(String(128), nullable=False)
description: Mapped[Optional[str]] = mapped_column(Text)
traffic_weight: Mapped[float] = mapped_column(Numeric, nullable=False, server_default="0.5")
is_control: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default="false")
model_version_id: Mapped[Optional[int]] = mapped_column(
Integer, ForeignKey("model_versions.id")
)
config: Mapped[Optional[dict]] = mapped_column(
JSON().with_variant(JSONB(), "postgresql")
) # For prompt variants or model config
created_at: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(
nullable=False, server_default=func.now(), onupdate=func.now()
)

# Relationships
experiment: Mapped[Experiment] = relationship(back_populates="variants")
model_version: Mapped[Optional[ModelVersion]] = relationship()
results: Mapped[list[ExperimentResult]] = relationship(
back_populates="variant",
cascade="all, delete-orphan",
)

__table_args__ = (
UniqueConstraint("experiment_id", "name", name="uq_variants_experiment_name"),
Index("ix_variants_experiment_id", "experiment_id"),
Index("ix_variants_model_version_id", "model_version_id"),
CheckConstraint(
"traffic_weight >= 0 AND traffic_weight <= 1",
name="ck_variants_traffic_weight",
),
)

# ... keep full definition

class ExperimentResult(Base):
"""Individual result from an A/B test experiment."""

__tablename__ = "experiment_results"

id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
variant_id: Mapped[int] = mapped_column(
Integer, ForeignKey("variants.id"), nullable=False
)
user_id: Mapped[Optional[str]] = mapped_column(String(128)) # Optional user identifier
session_id: Mapped[Optional[str]] = mapped_column(String(128)) # For session-based analysis
metrics: Mapped[dict] = mapped_column(
JSON().with_variant(JSONB(), "postgresql"), nullable=False
) # e.g., {"accuracy": 0.95, "latency_ms": 100}
metadata: Mapped[Optional[dict]] = mapped_column(
JSON().with_variant(JSONB(), "postgresql")
) # Additional context
created_at: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())

# Relationships
variant: Mapped[Variant] = relationship(back_populates="results")

__table_args__ = (
Index("ix_experiment_results_variant_id", "variant_id"),
Index("ix_experiment_results_user_id", "user_id"),
Index("ix_experiment_results_session_id", "session_id"),
Index("ix_experiment_results_created_at", "created_at"),
)
# ... keep full definition


# ---------------------------------------------------------------------------
Expand All @@ -758,103 +660,61 @@ class ExperimentResult(Base):

class GoldenDataset(Base):
"""Golden dataset for model evaluation and benchmarking."""

__tablename__ = "golden_datasets"
# ... keep full definition

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
name: Mapped[str] = mapped_column(String(128), nullable=False, unique=True)
description: Mapped[Optional[str]] = mapped_column(Text)
dataset_type: Mapped[str] = mapped_column(String(32), nullable=False) # 'classification', 'regression', 'anomaly_detection', etc.
task_type: Mapped[str] = mapped_column(String(32), nullable=False)
version: Mapped[str] = mapped_column(String(32), nullable=False)
source: Mapped[Optional[str]] = mapped_column(String(256)) # Data source identifier
size: Mapped[int] = mapped_column(Integer, nullable=False, server_default="0") # Number of entries
status: Mapped[str] = mapped_column(String(32), nullable=False, server_default="draft")
quality_score: Mapped[Optional[float]] = mapped_column(Numeric) # Overall quality metric (0-1)
metadata: Mapped[Optional[dict]] = mapped_column(
JSON().with_variant(JSONB(), "postgresql")
) # Additional dataset metadata
created_at: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(
nullable=False, server_default=func.now(), onupdate=func.now()
)
class GoldenDatasetEntry(Base):
"""Individual entry in a golden dataset with ground truth labels."""
__tablename__ = "golden_dataset_entries"
# ... keep full definition

# Relationships
entries: Mapped[list[GoldenDatasetEntry]] = relationship(
back_populates="dataset",
cascade="all, delete-orphan",
)

__table_args__ = (
Index("ix_golden_datasets_type", "dataset_type"),
Index("ix_golden_datasets_task_type", "task_type"),
Index("ix_golden_datasets_version", "version"),
Index("ix_golden_datasets_status", "status"),
UniqueConstraint("name", "version", name="uq_golden_datasets_name_version"),
CheckConstraint(
"dataset_type IN ('classification', 'regression', 'anomaly_detection', 'clustering', 'custom')",
name="ck_golden_datasets_type",
),
CheckConstraint(
"status IN ('draft', 'review', 'approved', 'archived')",
name="ck_golden_datasets_status",
),
CheckConstraint(
"quality_score IS NULL OR (quality_score >= 0 AND quality_score <= 1)",
name="ck_golden_datasets_quality_score",
),
)
# ---------------------------------------------------------------------------
# Ledger Processing
# ---------------------------------------------------------------------------

class ProcessedLedger(Base):
"""Tracking table for processed ledgers during backfill to ensure idempotency."""
__tablename__ = "processed_ledgers"
# ... keep full definition

class GoldenDatasetEntry(Base):
"""Individual entry in a golden dataset with ground truth labels."""
# ---------------------------------------------------------------------------

__tablename__ = "golden_dataset_entries"
class Experiment(Base):
"""A/B test experiment for comparing models or prompts."""
__tablename__ = "experiments"
# ... (fields and constraints from men branch)
# keep full definition

id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
dataset_id: Mapped[int] = mapped_column(
Integer, ForeignKey("golden_datasets.id"), nullable=False
)
input_data: Mapped[dict] = mapped_column(
JSON().with_variant(JSONB(), "postgresql"), nullable=False
) # Model input features
output_data: Mapped[dict] = mapped_column(
JSON().with_variant(JSONB(), "postgresql"), nullable=False
) # Ground truth labels
metadata: Mapped[Optional[dict]] = mapped_column(
JSON().with_variant(JSONB(), "postgresql")
) # Entry-specific metadata
difficulty: Mapped[Optional[float]] = mapped_column(Numeric) # Difficulty score (0-1)
confidence: Mapped[Optional[float]] = mapped_column(Numeric) # Label confidence (0-1)
created_at: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
class Variant(Base):
"""A variant in an A/B test experiment."""
__tablename__ = "variants"
# ... (fields and constraints from men branch)
# keep full definition

class ExperimentResult(Base):
"""Individual result from an A/B test experiment."""
__tablename__ = "experiment_results"
# ... (fields and constraints from men branch)
# keep full definition

# Relationships
dataset: Mapped[GoldenDataset] = relationship(back_populates="entries")

__table_args__ = (
Index("ix_golden_dataset_entries_dataset_id", "dataset_id"),
Index("ix_golden_dataset_entries_difficulty", "difficulty"),
Index("ix_golden_dataset_entries_confidence", "confidence"),
CheckConstraint(
"difficulty IS NULL OR (difficulty >= 0 AND difficulty <= 1)",
name="ck_golden_dataset_entries_difficulty",
),
CheckConstraint(
"confidence IS NULL OR (confidence >= 0 AND confidence <= 1)",
name="ck_golden_dataset_entries_confidence",
),
=======
# Ledger Processing
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------

class ProcessedLedger(Base):
"""Tracking table for processed ledgers during backfill to ensure idempotency."""
__tablename__ = "processed_ledgers"
# ... (fields and constraints from main branch)
# keep full definition

class ProcessedLedger(Base):
"""Tracking table for processed ledgers during backfill to ensure idempotency."""

__tablename__ = "processed_ledgers"

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
ledger_sequence: Mapped[int] = mapped_column(Integer, unique=True, nullable=False)
ledger_sequence: Mapped[int] = mapped_column(Integer, unique=True, nullable=False)
source: Mapped[str] = mapped_column(
String(256),
nullable=False,
Expand All @@ -864,14 +724,10 @@ class ProcessedLedger(Base):
nullable=False,
server_default=func.now(),
)
status: Mapped[
Literal["pending", "processing", "completed", "failed"]
] = mapped_column(
String(16),
nullable=False,
server_default="pending",
)
String(32),
status: Mapped[
Literal["pending", "processing", "completed", "failed"]
] = mapped_column(
String(16),
nullable=False,
server_default="pending",
)
Expand All @@ -890,5 +746,5 @@ class ProcessedLedger(Base):
Index("ix_processed_ledgers_status", "status"),
Index("ix_processed_ledgers_source", "source"),
)
>>>>>>> 0ce0bb2e1acdc9414b4d060a86e5547ae2e7dbf9

)
47 changes: 43 additions & 4 deletions astroml/tracking/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,45 @@
from .ab_testing import ABTestingFramework
from .golden_dataset import GoldenDatasetGenerator
from .mlflow_tracker import MLflowTracker
from .model_registry import ModelRegistry
# ---------------------------------------------------------------------------
# A/B Testing Framework
# ---------------------------------------------------------------------------

class Experiment(Base):
"""A/B test experiment for comparing models or prompts."""
__tablename__ = "experiments"
# ... keep full definition

class Variant(Base):
"""A variant in an A/B test experiment."""
__tablename__ = "variants"
# ... keep full definition

class ExperimentResult(Base):
"""Individual result from an A/B test experiment."""
__tablename__ = "experiment_results"
# ... keep full definition


# ---------------------------------------------------------------------------
# Golden Dataset Framework
# ---------------------------------------------------------------------------

class GoldenDataset(Base):
"""Golden dataset for model evaluation and benchmarking."""
__tablename__ = "golden_datasets"
# ... keep full definition

class GoldenDatasetEntry(Base):
"""Individual entry in a golden dataset with ground truth labels."""
__tablename__ = "golden_dataset_entries"
# ... keep full definition


# ---------------------------------------------------------------------------
# Ledger Processing
# ---------------------------------------------------------------------------

class ProcessedLedger(Base):
"""Tracking table for processed ledgers during backfill to ensure idempotency."""
__tablename__ = "processed_ledgers"
# ... keep full definition

__all__ = ["MLflowTracker", "ModelRegistry", "ABTestingFramework", "GoldenDatasetGenerator"]
2 changes: 2 additions & 0 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,10 @@ scipy>=1.11.0
scikit-learn>=1.3.0
pandas>=2.0
polars>=1.0
scipy>=1.10

# ── Database / configuration ───────────────────────────────────────────────

sqlalchemy>=2.0
alembic>=1.12
psycopg2-binary>=2.9
Expand Down
Loading
Loading