Skip to main content

assay_core/storage/
judge_cache.rs

1use super::now_rfc3339ish;
2use crate::storage::Store;
3use rusqlite::params;
4
5#[derive(Clone)]
6pub struct JudgeCache {
7    store: Store,
8}
9
10impl JudgeCache {
11    pub fn new(store: Store) -> Self {
12        Self { store }
13    }
14
15    pub fn get(&self, key: &str) -> anyhow::Result<Option<serde_json::Value>> {
16        let conn = self.store.conn.lock().unwrap();
17        let mut stmt = conn.prepare("SELECT payload_json FROM judge_cache WHERE key=?1")?;
18        let mut rows = stmt.query(params![key])?;
19        if let Some(row) = rows.next()? {
20            let s: String = row.get(0)?;
21            let val: serde_json::Value = serde_json::from_str(&s)?;
22            Ok(Some(val))
23        } else {
24            Ok(None)
25        }
26    }
27
28    pub fn put(
29        &self,
30        key: &str,
31        provider: &str,
32        model: &str,
33        rubric_id: &str,
34        rubric_version: &str,
35        payload: &serde_json::Value,
36    ) -> anyhow::Result<()> {
37        let conn = self.store.conn.lock().unwrap();
38        let payload_json = serde_json::to_string(payload)?;
39        let created_at = now_rfc3339ish();
40
41        conn.execute(
42            "INSERT INTO judge_cache(
43                key, provider, model, rubric_id, rubric_version, created_at, payload_json
44            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
45            ON CONFLICT(key) DO UPDATE SET
46                payload_json=excluded.payload_json,
47                created_at=excluded.created_at",
48            params![
49                key,
50                provider,
51                model,
52                rubric_id,
53                rubric_version,
54                created_at,
55                payload_json
56            ],
57        )?;
58        Ok(())
59    }
60}