systemprompt_evaluation/repository/
runs.rs

1//! Repository for evaluation runs.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use sqlx::PgPool;
7use std::str::FromStr;
8use std::sync::Arc;
9use systemprompt_database::DbPool;
10use systemprompt_identifiers::{EvalRubricId, EvalRunId, UserId};
11
12use crate::error::{EvaluationError, Result};
13use crate::models::{EvalRun, EvalRunKind, EvalRunStatus, NewRunParams, TriggerSource};
14
15#[derive(Debug, Clone)]
16pub struct EvalRunRepository {
17    pool: Arc<PgPool>,
18}
19
20impl EvalRunRepository {
21    pub fn new(db: &DbPool) -> Result<Self> {
22        Ok(Self {
23            pool: db.write_pool_arc()?,
24        })
25    }
26
27    pub async fn create(&self, params: &NewRunParams) -> Result<EvalRunId> {
28        let id = EvalRunId::generate();
29        sqlx::query!(
30            r#"
31            INSERT INTO eval_runs (
32                id, kind, judge_provider, judge_model, sample_size,
33                created_by, rubric_id, trigger_source
34            )
35            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
36            "#,
37            id.as_str(),
38            params.kind.as_str(),
39            params.judge_provider,
40            params.judge_model,
41            params.sample_size,
42            params.created_by.as_str(),
43            params.rubric_id.as_ref().map(EvalRubricId::as_str),
44            params.trigger_source.as_str()
45        )
46        .execute(self.pool.as_ref())
47        .await?;
48        Ok(id)
49    }
50
51    pub async fn get(&self, id: &EvalRunId) -> Result<EvalRun> {
52        let row = sqlx::query!(
53            r#"
54            SELECT id, kind, status, judge_provider, judge_model, sample_size,
55                   scored_count, failed_count, cost_microdollars, created_by,
56                   created_at, completed_at, error_message, rubric_id, trigger_source
57            FROM eval_runs
58            WHERE id = $1
59            "#,
60            id.as_str()
61        )
62        .fetch_optional(self.pool.as_ref())
63        .await?
64        .ok_or_else(|| EvaluationError::RunNotFound(id.as_str().to_owned()))?;
65
66        Ok(EvalRun {
67            id: EvalRunId::new(row.id),
68            kind: EvalRunKind::from_str(&row.kind).map_err(EvaluationError::JudgeParse)?,
69            status: match row.status.as_str() {
70                "completed" => EvalRunStatus::Completed,
71                "failed" => EvalRunStatus::Failed,
72                _ => EvalRunStatus::Running,
73            },
74            judge_provider: row.judge_provider,
75            judge_model: row.judge_model,
76            sample_size: row.sample_size,
77            scored_count: row.scored_count,
78            failed_count: row.failed_count,
79            cost_microdollars: row.cost_microdollars,
80            created_by: UserId::new(row.created_by),
81            created_at: row.created_at,
82            completed_at: row.completed_at,
83            error_message: row.error_message,
84            rubric_id: row.rubric_id.map(EvalRubricId::new),
85            trigger_source: match row.trigger_source.as_str() {
86                "scheduled" => TriggerSource::Scheduled,
87                "cli" => TriggerSource::Cli,
88                _ => TriggerSource::Manual,
89            },
90        })
91    }
92
93    pub async fn list_recent(&self, limit: i64) -> Result<Vec<EvalRun>> {
94        let ids = sqlx::query_scalar!(
95            "SELECT id FROM eval_runs ORDER BY created_at DESC LIMIT $1",
96            limit
97        )
98        .fetch_all(self.pool.as_ref())
99        .await?;
100
101        let mut runs = Vec::with_capacity(ids.len());
102        for id in ids {
103            runs.push(self.get(&EvalRunId::new(id)).await?);
104        }
105        Ok(runs)
106    }
107
108    pub async fn record_scored(&self, id: &EvalRunId, failed: bool, cost: i64) -> Result<()> {
109        sqlx::query!(
110            r#"
111            UPDATE eval_runs
112            SET scored_count = scored_count + 1,
113                failed_count = failed_count + CASE WHEN $2 THEN 1 ELSE 0 END,
114                cost_microdollars = cost_microdollars + $3
115            WHERE id = $1
116            "#,
117            id.as_str(),
118            failed,
119            cost
120        )
121        .execute(self.pool.as_ref())
122        .await?;
123        Ok(())
124    }
125
126    pub async fn complete(&self, id: &EvalRunId) -> Result<()> {
127        self.finish(id, EvalRunStatus::Completed, None).await
128    }
129
130    pub async fn fail(&self, id: &EvalRunId, error_message: &str) -> Result<()> {
131        self.finish(id, EvalRunStatus::Failed, Some(error_message))
132            .await
133    }
134
135    async fn finish(
136        &self,
137        id: &EvalRunId,
138        status: EvalRunStatus,
139        error_message: Option<&str>,
140    ) -> Result<()> {
141        sqlx::query!(
142            r#"
143            UPDATE eval_runs
144            SET status = $2, error_message = $3, completed_at = CURRENT_TIMESTAMP
145            WHERE id = $1
146            "#,
147            id.as_str(),
148            status.as_str(),
149            error_message
150        )
151        .execute(self.pool.as_ref())
152        .await?;
153        Ok(())
154    }
155}
systemprompt_evaluation/repository/runs.rs

systemprompt_evaluation/repository/
runs.rs
