1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use engram_storage::Memory;
7
8use crate::config::expand_tilde;
9use crate::error::CoreError;
10use crate::export_handler::memory_to_portable_json;
11use crate::lock_helpers;
12use crate::server::ServerState;
13use crate::tags_normalize::{TagsInput, normalize_tags};
14use crate::timestamp::current_utc_timestamp;
15
16const INSIGHT_MEMORY_TYPE: &str = "insight";
17
18#[derive(Debug, Deserialize)]
19#[serde(tag = "type", rename_all = "snake_case")]
20pub enum TrainerMessage {
21 Progress {
22 stage: String,
23 percent: f64,
24 },
25 Insight {
26 id: String,
27 context: String,
28 action: String,
29 result: String,
30 insight_type: String,
31 tags: Option<String>,
32 source_ids: Option<String>,
33 },
34 Recommendation {
35 target_id: String,
36 action: String,
37 reasoning: String,
38 },
39 Metric {
40 name: String,
41 value: f64,
42 },
43 Artifact {
44 path: String,
45 size_bytes: u64,
46 },
47 Complete {
48 insights_generated: u64,
49 duration_secs: f64,
50 },
51}
52
53#[derive(Deserialize)]
54struct DeleteParams {
55 id: String,
56}
57
58pub async fn handle_generate(state: &Arc<ServerState>, _params: Value) -> Result<Value, CoreError> {
59 let trainer_binary = resolve_trainer_binary(&state.config.trainer.trainer_binary);
60 validate_trainer_exists(&trainer_binary)?;
61 let timeout_secs = state.config.trainer.trainer_timeout_secs;
62 let models_path = expand_tilde(&state.config.trainer.models_path);
63 let output = spawn_trainer(
64 &trainer_binary,
65 &state.database_path,
66 &models_path,
67 timeout_secs,
68 )
69 .await?;
70 let insights = parse_trainer_output(&output)?;
71 let inserted_count = insert_generated_insights(state, &insights).await?;
72 Ok(json!({
73 "generated": inserted_count,
74 "messages_parsed": insights.len(),
75 }))
76}
77
78pub async fn handle_list(state: &Arc<ServerState>, _params: Value) -> Result<Value, CoreError> {
79 let state_clone = Arc::clone(state);
80 tokio::task::spawn_blocking(move || {
81 let database = lock_helpers::lock_db(&state_clone);
82 let insights = query_active_insights(&database)?;
83 let count = insights.len();
84 let serialized: Vec<Value> = insights.iter().map(memory_to_portable_json).collect();
85 Ok::<Value, CoreError>(json!({ "insights": serialized, "count": count }))
86 })
87 .await
88 .map_err(|error| CoreError::SocketError(error.to_string()))?
89}
90
91pub async fn handle_delete(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
92 let parsed: DeleteParams = serde_json::from_value(params)
93 .map_err(|error| CoreError::DispatchError(error.to_string()))?;
94 let insight_id = parsed.id;
95 let state_clone = Arc::clone(state);
96 tokio::task::spawn_blocking(move || {
97 let database = lock_helpers::lock_db(&state_clone);
98 let memory = database.get_memory(&insight_id)?;
99 if memory.memory_type != INSIGHT_MEMORY_TYPE {
100 return Err(CoreError::DispatchError(format!(
101 "memory {insight_id} is not an insight"
102 )));
103 }
104 database.delete_memory(&insight_id)?;
105 Ok::<Value, CoreError>(json!({ "deleted": insight_id }))
106 })
107 .await
108 .map_err(|error| CoreError::SocketError(error.to_string()))?
109}
110
111fn resolve_trainer_binary(configured_path: &str) -> String {
112 expand_tilde(configured_path)
113}
114
115fn validate_trainer_exists(binary_path: &str) -> Result<(), CoreError> {
116 if which_binary(binary_path).is_none() {
117 return Err(CoreError::TrainerFailed(binary_path.to_string()));
118 }
119 Ok(())
120}
121
122fn which_binary(binary_path: &str) -> Option<std::path::PathBuf> {
123 let path = std::path::Path::new(binary_path);
124 if path.is_absolute() {
125 return if path.exists() {
126 Some(path.to_path_buf())
127 } else {
128 None
129 };
130 }
131 std::env::var_os("PATH").and_then(|paths| {
132 std::env::split_paths(&paths)
133 .map(|directory| directory.join(binary_path))
134 .find(|candidate| candidate.exists())
135 })
136}
137
138async fn spawn_trainer(
139 binary_path: &str,
140 database_path: &str,
141 models_path: &str,
142 timeout_secs: u64,
143) -> Result<String, CoreError> {
144 let child = tokio::process::Command::new(binary_path)
145 .arg("--database")
146 .arg(database_path)
147 .arg("--models-path")
148 .arg(models_path)
149 .stdout(std::process::Stdio::piped())
150 .stderr(std::process::Stdio::piped())
151 .spawn()
152 .map_err(|error| CoreError::TrainerFailed(format!("{binary_path}: {error}")))?;
153
154 let timeout_duration = std::time::Duration::from_secs(timeout_secs);
155 let wait_result = tokio::time::timeout(timeout_duration, child.wait_with_output()).await;
156
157 match wait_result {
158 Ok(Ok(output)) => {
159 if !output.status.success() {
160 let stderr = String::from_utf8_lossy(&output.stderr);
161 return Err(CoreError::TrainerFailed(format!(
162 "{binary_path} exited with {}: {stderr}",
163 output.status
164 )));
165 }
166 String::from_utf8(output.stdout).map_err(|error| {
167 CoreError::TrainerMalformedOutput(format!("non-utf8 stdout: {error}"))
168 })
169 }
170 Ok(Err(error)) => Err(CoreError::TrainerFailed(format!("{binary_path}: {error}"))),
171 Err(_) => Err(CoreError::TrainerTimeout),
172 }
173}
174
175pub fn parse_trainer_output(output: &str) -> Result<Vec<TrainerMessage>, CoreError> {
176 let mut messages = Vec::new();
177 for (line_number, line) in output.lines().enumerate() {
178 let trimmed = line.trim();
179 if trimmed.is_empty() {
180 continue;
181 }
182 let message: TrainerMessage = serde_json::from_str(trimmed).map_err(|error| {
183 CoreError::TrainerMalformedOutput(format!("line {}: {error}", line_number + 1))
184 })?;
185 messages.push(message);
186 }
187 Ok(messages)
188}
189
190async fn insert_generated_insights(
191 state: &Arc<ServerState>,
192 messages: &[TrainerMessage],
193) -> Result<u64, CoreError> {
194 let mut insights: Vec<Memory> = Vec::new();
195 for message in messages {
196 if let Some(memory) = build_insight_memory(message)? {
197 insights.push(memory);
198 }
199 }
200 let count = insights.len() as u64;
201 if insights.is_empty() {
202 return Ok(0);
203 }
204 let state_clone = Arc::clone(state);
205 tokio::task::spawn_blocking(move || {
206 let database = lock_helpers::lock_db(&state_clone);
207 database.replace_insights(&insights)?;
208 Ok::<u64, CoreError>(count)
209 })
210 .await
211 .map_err(|error| CoreError::SocketError(error.to_string()))?
212}
213
214fn build_insight_memory(message: &TrainerMessage) -> Result<Option<Memory>, CoreError> {
215 match message {
216 TrainerMessage::Insight {
217 id,
218 context,
219 action,
220 result,
221 insight_type,
222 tags,
223 source_ids,
224 } => {
225 let timestamp = current_utc_timestamp();
226 let normalized_tags = normalize_tags(tags.clone().map(TagsInput::Encoded));
227 Ok(Some(Memory {
228 id: id.clone(),
229 memory_type: INSIGHT_MEMORY_TYPE.to_string(),
230 context: context.clone(),
231 action: action.clone(),
232 result: result.clone(),
233 score: 0.0,
234 embedding_context: None,
235 embedding_action: None,
236 embedding_result: None,
237 indexed: false,
238 tags: normalized_tags,
239 project: None,
240 parent_id: None,
241 source_ids: source_ids.clone(),
242 insight_type: Some(insight_type.clone()),
243 created_at: timestamp.clone(),
244 updated_at: timestamp,
245 used_count: 0,
246 last_used_at: None,
247 superseded_by: None,
248 }))
249 }
250 _ => Ok(None),
251 }
252}
253
254fn query_active_insights(database: &engram_storage::Database) -> Result<Vec<Memory>, CoreError> {
255 let mut statement = database
256 .connection()
257 .prepare(
258 "SELECT * FROM memories WHERE memory_type = ?1 \
259 AND superseded_by IS NULL",
260 )
261 .map_err(|error| CoreError::Storage(engram_storage::StorageError::Sqlite(error)))?;
262 let rows = statement
263 .query_map([INSIGHT_MEMORY_TYPE], engram_storage::row_to_memory)
264 .map_err(|error| CoreError::Storage(engram_storage::StorageError::Sqlite(error)))?;
265 let mut insights = Vec::new();
266 for row in rows {
267 insights.push(
268 row.map_err(|error| CoreError::Storage(engram_storage::StorageError::Sqlite(error)))?,
269 );
270 }
271 Ok(insights)
272}