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