engram_core/
judge_handler.rs

1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use engram_judge::{CombinedJudge, JudgeInput};
7
8use crate::error::CoreError;
9use crate::server::ServerState;
10use crate::timestamp::{current_utc_timestamp, parse_timestamp_to_epoch};
11
12const MAX_MEMORY_ID_LENGTH: usize = 100;
13
14#[derive(Deserialize)]
15struct JudgeParams {
16    memory_id: String,
17    query: Option<String>,
18    score: Option<f32>,
19}
20
21pub async fn handle(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
22    let parsed: JudgeParams = serde_json::from_value(params)
23        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
24    if parsed.memory_id.is_empty() {
25        return Err(CoreError::DispatchError(
26            "memory_id must not be empty".into(),
27        ));
28    }
29    if parsed.memory_id.len() > MAX_MEMORY_ID_LENGTH {
30        return Err(CoreError::DispatchError(format!(
31            "memory_id exceeds maximum length of {MAX_MEMORY_ID_LENGTH} chars"
32        )));
33    }
34    if let Some(explicit_score) = parsed.score {
35        return apply_explicit_score(state, &parsed.memory_id, explicit_score).await;
36    }
37    let query = parsed.query.unwrap_or_default();
38    compute_and_apply_score(state, &parsed.memory_id, &query).await
39}
40
41async fn apply_explicit_score(
42    state: &Arc<ServerState>,
43    memory_id: &str,
44    score: f32,
45) -> Result<Value, CoreError> {
46    let id_owned = memory_id.to_string();
47    let state_clone = Arc::clone(state);
48    tokio::task::spawn_blocking(move || {
49        let database = state_clone.database.lock().unwrap();
50        database.set_memory_score(&id_owned, score)?;
51        let timestamp = current_utc_timestamp();
52        database.mark_judged(&id_owned, &timestamp)?;
53        Ok::<(), CoreError>(())
54    })
55    .await
56    .map_err(|error| CoreError::SocketError(error.to_string()))??;
57    Ok(json!({
58        "score": score,
59        "reason": "explicit score",
60        "degraded": false,
61    }))
62}
63
64async fn compute_and_apply_score(
65    state: &Arc<ServerState>,
66    memory_id: &str,
67    query: &str,
68) -> Result<Value, CoreError> {
69    let id_owned = memory_id.to_string();
70    let query_owned = query.to_string();
71    let state_clone = Arc::clone(state);
72    let result = tokio::task::spawn_blocking(move || {
73        let database = state_clone.database.lock().unwrap();
74        let memory = database.get_memory(&id_owned)?;
75        let judge_input = build_judge_input(&memory);
76        let judge = match state_clone.text_generator.as_deref() {
77            Some(generator) => {
78                CombinedJudge::with_llm(generator as &dyn engram_llm_client::TextGenerator)
79            }
80            None => CombinedJudge::heuristic_only(),
81        };
82        let judge_score = judge.score(&query_owned, &judge_input);
83        database.set_memory_score(&id_owned, judge_score.score)?;
84        let timestamp = current_utc_timestamp();
85        database.mark_judged(&id_owned, &timestamp)?;
86        Ok::<_, CoreError>(json!({
87            "score": judge_score.score,
88            "reason": judge_score.reason,
89            "degraded": judge_score.degraded,
90        }))
91    })
92    .await
93    .map_err(|error| CoreError::SocketError(error.to_string()))??;
94    Ok(result)
95}
96
97fn build_judge_input(memory: &engram_storage::Memory) -> JudgeInput {
98    let days_since_update = compute_days_since_update(&memory.updated_at);
99    JudgeInput {
100        context: memory.context.clone(),
101        action: memory.action.clone(),
102        result: memory.result.clone(),
103        days_since_update,
104        used_count: memory.used_count.max(0) as u64,
105    }
106}
107
108fn compute_days_since_update(updated_at: &str) -> f64 {
109    let now = std::time::SystemTime::now()
110        .duration_since(std::time::UNIX_EPOCH)
111        .unwrap_or_default()
112        .as_secs();
113    let parsed = parse_timestamp_to_epoch(updated_at).unwrap_or(now);
114    let elapsed_seconds = now.saturating_sub(parsed);
115    elapsed_seconds as f64 / 86400.0
116}
engram_core/judge_handler.rs

engram_core/
judge_handler.rs
