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