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