Skip to main content

kaizen/core_loop/
cases.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2use crate::core::event::SessionRecord;
3use crate::core_loop::{CaseRecord, CaseRef, CaseStatus};
4use crate::store::Store;
5use anyhow::{Result, anyhow};
6use rusqlite::{OptionalExtension, params};
7
8pub fn create_case(
9    store: &Store,
10    session: &SessionRecord,
11    source_key: &str,
12    reason: &str,
13    label: Option<String>,
14    now_ms: u64,
15) -> Result<CaseRecord> {
16    let rec = record(session, source_key, reason, label, now_ms);
17    store.conn().execute(
18        "INSERT OR IGNORE INTO cases
19         (id, source_key, session_id, reason, label, status, prompt_fingerprint, metadata_json, created_at_ms)
20         VALUES (?1, ?2, ?3, ?4, ?5, 'open', ?6, ?7, ?8)",
21        params![rec.id, rec.source_key, rec.session_id, rec.reason, rec.label, rec.prompt_fingerprint, rec.metadata_json, rec.created_at_ms as i64],
22    )?;
23    get_by_source(store, source_key)
24}
25
26pub fn add_ref(store: &Store, case_id: &str, ref_kind: &str, ref_key: &str) -> Result<()> {
27    store.conn().execute(
28        "INSERT OR IGNORE INTO case_refs (case_id, ref_kind, ref_key) VALUES (?1, ?2, ?3)",
29        params![case_id, ref_kind, ref_key],
30    )?;
31    Ok(())
32}
33
34pub fn list(store: &Store, status: Option<CaseStatus>) -> Result<Vec<CaseRecord>> {
35    let sql = "SELECT id, source_key, session_id, reason, label, status, prompt_fingerprint, metadata_json, created_at_ms FROM cases WHERE (?1 IS NULL OR status = ?1) ORDER BY created_at_ms DESC";
36    let mut stmt = store.conn().prepare(sql)?;
37    let rows = stmt.query_map(params![status.map(|s| s.as_str().to_string())], row)?;
38    rows.map(|r| r.map_err(anyhow::Error::from)).collect()
39}
40
41pub fn get(store: &Store, id: &str) -> Result<CaseRecord> {
42    let sql = "SELECT id, source_key, session_id, reason, label, status, prompt_fingerprint, metadata_json, created_at_ms FROM cases WHERE id = ?1";
43    store
44        .conn()
45        .query_row(sql, params![id], row)
46        .optional()?
47        .ok_or_else(|| anyhow!("case not found: {id}"))
48}
49
50pub fn refs(store: &Store, case_id: &str) -> Result<Vec<CaseRef>> {
51    let mut stmt = store.conn().prepare("SELECT case_id, ref_kind, ref_key FROM case_refs WHERE case_id = ?1 ORDER BY ref_kind, ref_key")?;
52    let rows = stmt.query_map(params![case_id], |r| {
53        Ok(CaseRef {
54            case_id: r.get(0)?,
55            ref_kind: r.get(1)?,
56            ref_key: r.get(2)?,
57        })
58    })?;
59    rows.map(|r| r.map_err(anyhow::Error::from)).collect()
60}
61
62pub fn archive(store: &Store, id: &str) -> Result<()> {
63    store.conn().execute(
64        "UPDATE cases SET status = 'archived' WHERE id = ?1",
65        params![id],
66    )?;
67    Ok(())
68}
69
70pub fn mine(store: &Store, since_ms: u64, now_ms: u64) -> Result<Vec<CaseRecord>> {
71    let mut out = eval_cases(store, since_ms, now_ms)?;
72    out.extend(feedback_cases(store, since_ms, now_ms)?);
73    Ok(out)
74}
75
76fn eval_cases(store: &Store, since_ms: u64, now_ms: u64) -> Result<Vec<CaseRecord>> {
77    store
78        .list_evals_in_window(since_ms, now_ms)?
79        .into_iter()
80        .filter(|r| r.flagged || r.score < 0.4)
81        .map(|r| {
82            from_session(
83                store,
84                &r.session_id,
85                &format!("eval:{}", r.rubric_id),
86                "low_eval",
87                now_ms,
88            )
89        })
90        .collect()
91}
92
93fn feedback_cases(store: &Store, since_ms: u64, now_ms: u64) -> Result<Vec<CaseRecord>> {
94    store
95        .list_feedback_in_window(since_ms, now_ms)?
96        .into_iter()
97        .filter(|r| {
98            r.label.as_ref().is_some_and(|l| {
99                matches!(
100                    l,
101                    crate::feedback::types::FeedbackLabel::Bad
102                        | crate::feedback::types::FeedbackLabel::Regression
103                )
104            })
105        })
106        .map(|r| from_session(store, &r.session_id, "feedback:bad", "bad_feedback", now_ms))
107        .collect()
108}
109
110fn from_session(
111    store: &Store,
112    id: &str,
113    prefix: &str,
114    reason: &str,
115    now_ms: u64,
116) -> Result<CaseRecord> {
117    let s = store
118        .get_session(id)?
119        .ok_or_else(|| anyhow!("session not found: {id}"))?;
120    let key = format!("{prefix}:{id}");
121    let rec = create_case(store, &s, &key, reason, Some(reason.into()), now_ms)?;
122    add_ref(store, &rec.id, "session", id)?;
123    Ok(rec)
124}
125
126fn record(
127    s: &SessionRecord,
128    source_key: &str,
129    reason: &str,
130    label: Option<String>,
131    now_ms: u64,
132) -> CaseRecord {
133    CaseRecord {
134        id: uuid::Uuid::now_v7().to_string(),
135        source_key: source_key.into(),
136        session_id: s.id.clone(),
137        reason: reason.into(),
138        label,
139        status: CaseStatus::Open,
140        prompt_fingerprint: s.prompt_fingerprint.clone(),
141        metadata_json: "{}".into(),
142        created_at_ms: now_ms,
143    }
144}
145
146fn get_by_source(store: &Store, source_key: &str) -> Result<CaseRecord> {
147    let sql = "SELECT id, source_key, session_id, reason, label, status, prompt_fingerprint, metadata_json, created_at_ms FROM cases WHERE source_key = ?1";
148    store
149        .conn()
150        .query_row(sql, params![source_key], row)
151        .map_err(Into::into)
152}
153
154fn row(r: &rusqlite::Row<'_>) -> rusqlite::Result<CaseRecord> {
155    Ok(CaseRecord {
156        id: r.get(0)?,
157        source_key: r.get(1)?,
158        session_id: r.get(2)?,
159        reason: r.get(3)?,
160        label: r.get(4)?,
161        status: status(r.get::<_, String>(5)?.as_str()),
162        prompt_fingerprint: r.get(6)?,
163        metadata_json: r.get(7)?,
164        created_at_ms: r.get::<_, i64>(8)? as u64,
165    })
166}
167
168fn status(s: &str) -> CaseStatus {
169    if s == "archived" {
170        CaseStatus::Archived
171    } else {
172        CaseStatus::Open
173    }
174}