147 lines
5.6 KiB
Python
147 lines
5.6 KiB
Python
import uuid
|
|
from datetime import datetime, timezone
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import (
|
|
JSON,
|
|
Column,
|
|
DateTime,
|
|
Float,
|
|
ForeignKey,
|
|
Integer,
|
|
String,
|
|
Text,
|
|
UniqueConstraint,
|
|
)
|
|
from sqlalchemy.dialects.postgresql import UUID
|
|
from sqlalchemy.orm import relationship
|
|
|
|
from app.db import Base
|
|
|
|
|
|
def now_utc() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
class Sample(Base):
|
|
__tablename__ = "samples"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
repo = Column(String(255), nullable=False)
|
|
commit_sha = Column(String(40), nullable=False)
|
|
language = Column(String(50), nullable=False)
|
|
diff = Column(Text, nullable=False)
|
|
before_context = Column(JSON, nullable=True)
|
|
after_context = Column(JSON, nullable=True)
|
|
created_at = Column(DateTime, default=now_utc, nullable=False)
|
|
|
|
defects = relationship("Defect", back_populates="sample", cascade="all, delete-orphan")
|
|
runs = relationship("ExperimentRun", back_populates="sample")
|
|
|
|
|
|
class Defect(Base):
|
|
__tablename__ = "defects"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
sample_id = Column(UUID(as_uuid=True), ForeignKey("samples.id"), nullable=False)
|
|
defect_type = Column(String(100), nullable=False)
|
|
language = Column(String(50), nullable=False)
|
|
line_start = Column(Integer, nullable=True)
|
|
line_end = Column(Integer, nullable=True)
|
|
description = Column(Text, nullable=True)
|
|
reference_fix = Column(Text, nullable=True)
|
|
created_at = Column(DateTime, default=now_utc, nullable=False)
|
|
|
|
sample = relationship("Sample", back_populates="defects")
|
|
|
|
|
|
class PromptTemplate(Base):
|
|
__tablename__ = "prompt_templates"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
strategy_id = Column(String(100), nullable=False)
|
|
level = Column(String(20), nullable=False)
|
|
created_at = Column(DateTime, default=now_utc, nullable=False)
|
|
|
|
versions = relationship(
|
|
"PromptTemplateVersion",
|
|
back_populates="template",
|
|
cascade="all, delete-orphan",
|
|
order_by="PromptTemplateVersion.version_number",
|
|
)
|
|
|
|
__table_args__ = (UniqueConstraint("strategy_id", "level", name="uix_strategy_level"),)
|
|
|
|
|
|
class PromptTemplateVersion(Base):
|
|
__tablename__ = "prompt_template_versions"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
template_id = Column(UUID(as_uuid=True), ForeignKey("prompt_templates.id"), nullable=False)
|
|
version_number = Column(Integer, nullable=False)
|
|
body = Column(Text, nullable=False)
|
|
variables_schema = Column(JSON, nullable=False, default=dict)
|
|
created_at = Column(DateTime, default=now_utc, nullable=False)
|
|
|
|
template = relationship("PromptTemplate", back_populates="versions")
|
|
runs = relationship("ExperimentRun", back_populates="template_version")
|
|
|
|
|
|
class Experiment(Base):
|
|
__tablename__ = "experiments"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
name = Column(String(255), nullable=False)
|
|
models = Column(JSON, nullable=False)
|
|
levels = Column(JSON, nullable=False)
|
|
sample_ids = Column(JSON, nullable=False)
|
|
repeats = Column(Integer, nullable=False, default=3)
|
|
sampling_params = Column(JSON, nullable=False, default=dict)
|
|
status = Column(String(20), default="pending", nullable=False)
|
|
created_at = Column(DateTime, default=now_utc, nullable=False)
|
|
|
|
runs = relationship("ExperimentRun", back_populates="experiment", cascade="all, delete-orphan")
|
|
|
|
|
|
class ExperimentRun(Base):
|
|
__tablename__ = "experiment_runs"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
run_id = Column(String(64), unique=True, nullable=False, default=lambda: str(uuid.uuid4()))
|
|
experiment_id = Column(UUID(as_uuid=True), ForeignKey("experiments.id"), nullable=False)
|
|
sample_id = Column(UUID(as_uuid=True), ForeignKey("samples.id"), nullable=False)
|
|
model_id = Column(String(100), nullable=False)
|
|
template_version_id = Column(UUID(as_uuid=True), ForeignKey("prompt_template_versions.id"), nullable=False)
|
|
repeat_index = Column(Integer, nullable=False)
|
|
status = Column(String(20), default="pending", nullable=False)
|
|
retry_count = Column(Integer, default=0, nullable=False)
|
|
sampling_params = Column(JSON, nullable=False, default=dict)
|
|
started_at = Column(DateTime, nullable=True)
|
|
completed_at = Column(DateTime, nullable=True)
|
|
created_at = Column(DateTime, default=now_utc, nullable=False)
|
|
|
|
experiment = relationship("Experiment", back_populates="runs")
|
|
sample = relationship("Sample", back_populates="runs")
|
|
template_version = relationship("PromptTemplateVersion", back_populates="runs")
|
|
result = relationship("Result", back_populates="run", uselist=False, cascade="all, delete-orphan")
|
|
|
|
|
|
class Result(Base):
|
|
__tablename__ = "results"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
run_id = Column(UUID(as_uuid=True), ForeignKey("experiment_runs.id"), nullable=False, unique=True)
|
|
raw_output = Column(Text, nullable=True)
|
|
token_usage = Column(JSON, nullable=True)
|
|
latency_ms = Column(Float, nullable=True)
|
|
parsed_findings = Column(JSON, nullable=True)
|
|
detection_rate = Column(Float, nullable=True)
|
|
false_positive_rate = Column(Float, nullable=True)
|
|
coverage_rate = Column(Float, nullable=True)
|
|
stability_score = Column(Float, nullable=True)
|
|
likert_score = Column(Integer, nullable=True)
|
|
created_at = Column(DateTime, default=now_utc, nullable=False)
|
|
updated_at = Column(DateTime, default=now_utc, onupdate=now_utc, nullable=False)
|
|
|
|
run = relationship("ExperimentRun", back_populates="result")
|