Skip to main content

systemprompt_evaluation/repository/
rubrics.rs

1//! Repository for evaluation rubrics.
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::sync::Arc;
8use systemprompt_database::DbPool;
9use systemprompt_identifiers::EvalRubricId;
10
11use crate::error::{EvaluationError, Result};
12use crate::models::{Rubric, RubricDimension};
13
14#[derive(Debug, Clone)]
15pub struct EvalRubricRepository {
16    pool: Arc<PgPool>,
17}
18
19impl EvalRubricRepository {
20    pub fn new(db: &DbPool) -> Result<Self> {
21        Ok(Self {
22            pool: db.write_pool_arc()?,
23        })
24    }
25
26    pub async fn upsert(&self, rubric: &Rubric) -> Result<()> {
27        let dimensions = serde_json::to_value(&rubric.dimensions)?;
28        sqlx::query!(
29            r#"
30            INSERT INTO eval_rubrics (id, name, dimensions, pass_threshold, prompt_template, enabled)
31            VALUES ($1, $2, $3, $4, $5, $6)
32            ON CONFLICT (name) DO UPDATE
33            SET dimensions = EXCLUDED.dimensions,
34                pass_threshold = EXCLUDED.pass_threshold,
35                prompt_template = EXCLUDED.prompt_template,
36                enabled = EXCLUDED.enabled
37            "#,
38            rubric.id.as_str(),
39            rubric.name,
40            dimensions,
41            rubric.pass_threshold,
42            rubric.prompt_template.as_deref(),
43            rubric.enabled
44        )
45        .execute(self.pool.as_ref())
46        .await?;
47        Ok(())
48    }
49
50    pub async fn get_by_name(&self, name: &str) -> Result<Rubric> {
51        let row = sqlx::query!(
52            r#"
53            SELECT id, name, dimensions, pass_threshold, prompt_template, enabled
54            FROM eval_rubrics
55            WHERE name = $1
56            "#,
57            name
58        )
59        .fetch_optional(self.pool.as_ref())
60        .await?
61        .ok_or_else(|| EvaluationError::RubricNotFound(name.to_owned()))?;
62
63        let dimensions: Vec<RubricDimension> = serde_json::from_value(row.dimensions)?;
64        Ok(Rubric {
65            id: EvalRubricId::new(row.id),
66            name: row.name,
67            dimensions,
68            pass_threshold: row.pass_threshold,
69            prompt_template: row.prompt_template,
70            enabled: row.enabled,
71        })
72    }
73}