atomr_agents_eval/
annotation.rs

1//! Annotation queue for human review.
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use atomr_agents_core::{Result, RunId, Value};
7use parking_lot::RwLock;
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum Verdict {
13    Pending,
14    Approved,
15    Rejected,
16    NeedsEdit,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct AnnotationItem {
21    pub id: String,
22    pub run_id: RunId,
23    pub prompt: String,
24    pub output: Value,
25    pub verdict: Verdict,
26    pub note: Option<String>,
27    pub created_at_ms: i64,
28}
29
30#[async_trait]
31pub trait AnnotationQueue: Send + Sync + 'static {
32    async fn enqueue(&self, item: AnnotationItem) -> Result<()>;
33    async fn next_pending(&self) -> Result<Option<AnnotationItem>>;
34    async fn submit(&self, id: &str, verdict: Verdict, note: Option<String>) -> Result<()>;
35    async fn list(&self) -> Result<Vec<AnnotationItem>>;
36}
37
38#[derive(Default, Clone)]
39pub struct InMemoryAnnotationQueue {
40    inner: Arc<RwLock<Vec<AnnotationItem>>>,
41}
42
43impl InMemoryAnnotationQueue {
44    pub fn new() -> Self {
45        Self::default()
46    }
47}
48
49#[async_trait]
50impl AnnotationQueue for InMemoryAnnotationQueue {
51    async fn enqueue(&self, item: AnnotationItem) -> Result<()> {
52        self.inner.write().push(item);
53        Ok(())
54    }
55
56    async fn next_pending(&self) -> Result<Option<AnnotationItem>> {
57        Ok(self
58            .inner
59            .read()
60            .iter()
61            .find(|i| matches!(i.verdict, Verdict::Pending))
62            .cloned())
63    }
64
65    async fn submit(&self, id: &str, verdict: Verdict, note: Option<String>) -> Result<()> {
66        let mut g = self.inner.write();
67        if let Some(item) = g.iter_mut().find(|i| i.id == id) {
68            item.verdict = verdict;
69            item.note = note;
70        }
71        Ok(())
72    }
73
74    async fn list(&self) -> Result<Vec<AnnotationItem>> {
75        Ok(self.inner.read().clone())
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    fn item(id: &str) -> AnnotationItem {
84        AnnotationItem {
85            id: id.into(),
86            run_id: RunId::from(format!("run-{id}")),
87            prompt: "hi".into(),
88            output: Value::String("answer".into()),
89            verdict: Verdict::Pending,
90            note: None,
91            created_at_ms: 0,
92        }
93    }
94
95    #[tokio::test]
96    async fn enqueue_then_submit_updates_verdict() {
97        let q = InMemoryAnnotationQueue::new();
98        q.enqueue(item("a")).await.unwrap();
99        q.enqueue(item("b")).await.unwrap();
100        let next = q.next_pending().await.unwrap().unwrap();
101        assert_eq!(next.id, "a");
102        q.submit("a", Verdict::Approved, Some("ok".into())).await.unwrap();
103        let next = q.next_pending().await.unwrap().unwrap();
104        assert_eq!(next.id, "b");
105        let all = q.list().await.unwrap();
106        assert_eq!(all[0].verdict, Verdict::Approved);
107        assert_eq!(all[0].note.as_deref(), Some("ok"));
108    }
109}
atomr_agents_eval/annotation.rs

atomr_agents_eval/
annotation.rs
