systemprompt_evaluation/repository/
rubrics.rs1use 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}