Skip to main content

engram_core/
status_handler.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use serde_json::{Value, json};
5
6use crate::config::{Config, expand_tilde};
7use crate::error::CoreError;
8use crate::server::ServerState;
9
10pub async fn handle(state: &Arc<ServerState>, _params: Value) -> Result<Value, CoreError> {
11    let state_clone = Arc::clone(state);
12    tokio::task::spawn_blocking(move || {
13        let database = state_clone.database.lock().unwrap();
14        let memory_count = count_memories(&database)?;
15        let indexed_count = count_indexed_memories(&database)?;
16        let pending_judgments = database.get_pending_judgments(usize::MAX)?.len();
17        let index_size = state_clone.indexes.read().unwrap().len();
18        let hints = build_hints(
19            &state_clone.config,
20            &state_clone.database_path,
21            memory_count,
22            pending_judgments,
23        );
24        Ok::<Value, CoreError>(json!({
25            "memory_count": memory_count,
26            "indexed_count": indexed_count,
27            "pending_judgments": pending_judgments,
28            "index_size": index_size,
29            "hints": hints,
30        }))
31    })
32    .await
33    .map_err(|error| CoreError::SocketError(error.to_string()))?
34}
35
36fn build_hints(
37    config: &Config,
38    database_path: &str,
39    memory_count: usize,
40    pending_judgments: usize,
41) -> Vec<String> {
42    let mut hints = Vec::new();
43    let models_path = expand_tilde(&config.trainer.models_path);
44    let has_onnx = directory_contains_onnx(&models_path);
45
46    if memory_count >= 20 && !has_onnx {
47        hints.push(format!(
48            "You have {memory_count} memories. Install trainer for self-learning: \
49             pip install engram-trainer && engram train"
50        ));
51    }
52
53    if memory_count >= 20
54        && has_onnx
55        && onnx_models_older_than_database(database_path, &models_path)
56    {
57        hints.push("Models may be outdated. Re-run: engram train".to_string());
58    }
59
60    if pending_judgments > 10 {
61        hints.push(format!(
62            "{pending_judgments} memories pending judgment. \
63             Use memory_judge to improve search quality"
64        ));
65    }
66
67    let llm_key_set = config
68        .llm
69        .api_key
70        .as_deref()
71        .is_some_and(|key| !key.is_empty());
72    if !llm_key_set {
73        hints.push(
74            "Set ENGRAM_OPENAI_API_KEY for LLM-powered judge, HyDE, and consolidation".to_string(),
75        );
76    }
77
78    hints
79}
80
81fn directory_contains_onnx(path: &str) -> bool {
82    let Ok(entries) = std::fs::read_dir(path) else {
83        return false;
84    };
85    entries
86        .flatten()
87        .any(|entry| entry.path().extension().is_some_and(|ext| ext == "onnx"))
88}
89
90fn onnx_models_older_than_database(database_path: &str, models_path: &str) -> bool {
91    let database_modified = Path::new(database_path)
92        .metadata()
93        .ok()
94        .and_then(|meta| meta.modified().ok());
95
96    let newest_onnx = std::fs::read_dir(models_path)
97        .ok()
98        .into_iter()
99        .flatten()
100        .flatten()
101        .filter(|entry| entry.path().extension().is_some_and(|ext| ext == "onnx"))
102        .filter_map(|entry| entry.metadata().ok()?.modified().ok())
103        .max();
104
105    match (database_modified, newest_onnx) {
106        (Some(db_time), Some(model_time)) => model_time < db_time,
107        _ => false,
108    }
109}
110
111fn count_memories(database: &engram_storage::Database) -> Result<usize, CoreError> {
112    let count: i64 = database
113        .connection()
114        .query_row("SELECT COUNT(*) FROM memories", [], |row| row.get(0))
115        .map_err(|error| CoreError::Storage(engram_storage::StorageError::Sqlite(error)))?;
116    Ok(count as usize)
117}
118
119fn count_indexed_memories(database: &engram_storage::Database) -> Result<usize, CoreError> {
120    let count: i64 = database
121        .connection()
122        .query_row(
123            "SELECT COUNT(*) FROM memories WHERE indexed = TRUE",
124            [],
125            |row| row.get(0),
126        )
127        .map_err(|error| CoreError::Storage(engram_storage::StorageError::Sqlite(error)))?;
128    Ok(count as usize)
129}