Skip to main content

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