vent_ai_platform/app/db/repositories/evaluation.py

from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.db.models.evaluation import (
    EvaluationIndicator,
    EvaluationRun,
    EvaluationScore,
)
from app.db.repositories.base import create_model


def list_indicators(
    db: Session,
    *,
    mine_id: str | None = None,
) -> list[EvaluationIndicator]:
    statement = select(EvaluationIndicator).where(EvaluationIndicator.enabled.is_(True))
    if mine_id:
        statement = statement.where(
            (EvaluationIndicator.mine_id == mine_id)
            | (EvaluationIndicator.mine_id.is_(None))
        )
    return list(
        db.scalars(
            statement.order_by(
                EvaluationIndicator.dimension,
                EvaluationIndicator.sort_order,
                EvaluationIndicator.code,
            )
        ).all()
    )


def list_runs(
    db: Session,
    mine_id: str,
    *,
    limit: int = 50,
) -> list[EvaluationRun]:
    return list(
        db.scalars(
            select(EvaluationRun)
            .where(EvaluationRun.mine_id == mine_id)
            .order_by(EvaluationRun.created_at.desc())
            .limit(limit)
        ).all()
    )


def get_run(db: Session, run_id: str) -> EvaluationRun | None:
    return db.get(EvaluationRun, run_id)


def list_scores(db: Session, run_id: str) -> list[EvaluationScore]:
    return list(
        db.scalars(
            select(EvaluationScore)
            .where(EvaluationScore.run_id == run_id)
            .order_by(EvaluationScore.weighted_score.desc())
        ).all()
    )


def create_run(db: Session, data: dict) -> EvaluationRun:
    return create_model(db, EvaluationRun, data)


def update_run(db: Session, run: EvaluationRun, data: dict) -> EvaluationRun:
    for key, value in data.items():
        setattr(run, key, value)
    db.add(run)
    db.commit()
    db.refresh(run)
    return run


def create_score(db: Session, data: dict) -> EvaluationScore:
    return create_model(db, EvaluationScore, data)