1use std::sync::Arc;
10use std::time::Duration;
11
12use tokio_util::sync::CancellationToken;
13use tracing::Instrument as _;
14use zeph_llm::any::AnyProvider;
15use zeph_llm::provider::LlmProvider as _;
16
17use crate::error::MemoryError;
18use crate::store::SqliteStore;
19use crate::sweep_helpers::{
20 EmbedBatchOutcome, cluster_by_cosine_similarity, embed_batch_with_validation,
21};
22use crate::types::{MemSceneId, MessageId};
23
24#[derive(Debug, Clone)]
29pub struct MemScene {
30 pub id: MemSceneId,
32 pub label: String,
34 pub profile: String,
36 pub member_count: u32,
38 pub created_at: i64,
40 pub updated_at: i64,
42}
43
44#[derive(Debug, Clone)]
46pub struct SceneConfig {
47 pub enabled: bool,
49 pub similarity_threshold: f32,
51 pub batch_size: usize,
53 pub sweep_interval_secs: u64,
55}
56
57pub async fn start_scene_consolidation_loop(
62 store: Arc<SqliteStore>,
63 provider: AnyProvider,
64 config: SceneConfig,
65 cancel: CancellationToken,
66) {
67 if !config.enabled {
68 tracing::debug!("scene consolidation disabled (tiers.scene_enabled = false)");
69 return;
70 }
71
72 let mut ticker = tokio::time::interval(Duration::from_secs(config.sweep_interval_secs));
73 ticker.tick().await;
75
76 loop {
77 tokio::select! {
78 () = cancel.cancelled() => {
79 tracing::debug!("scene consolidation loop shutting down");
80 return;
81 }
82 _ = ticker.tick() => {}
83 }
84
85 tracing::debug!("scene consolidation: starting sweep");
86 let start = std::time::Instant::now();
87
88 match consolidate_scenes(&store, &provider, &config).await {
89 Ok(stats) => {
90 tracing::info!(
91 candidates = stats.candidates,
92 scenes_created = stats.scenes_created,
93 messages_assigned = stats.messages_assigned,
94 elapsed_ms = start.elapsed().as_millis(),
95 "scene consolidation: sweep complete"
96 );
97 }
98 Err(e) => {
99 tracing::warn!(
100 error = %e,
101 elapsed_ms = start.elapsed().as_millis(),
102 "scene consolidation: sweep failed, will retry"
103 );
104 }
105 }
106 }
107}
108
109#[derive(Debug, Default)]
111pub struct SceneStats {
112 pub candidates: usize,
113 pub scenes_created: usize,
114 pub messages_assigned: usize,
115}
116
117#[tracing::instrument(name = "memory.scenes.consolidation_sweep", skip_all)]
123pub async fn consolidate_scenes(
124 store: &SqliteStore,
125 provider: &AnyProvider,
126 config: &SceneConfig,
127) -> Result<SceneStats, MemoryError> {
128 let candidates = store
129 .find_unscened_semantic_messages(config.batch_size)
130 .await?;
131
132 if candidates.len() < 2 {
133 return Ok(SceneStats::default());
134 }
135
136 let mut stats = SceneStats {
137 candidates: candidates.len(),
138 ..SceneStats::default()
139 };
140
141 if !provider.supports_embeddings() {
142 return Ok(stats);
143 }
144
145 let texts: Vec<&str> = candidates.iter().map(|(_, c)| c.as_str()).collect();
147 let span = tracing::info_span!("memory.scenes.embed_batch", count = texts.len());
148 let vecs = embed_batch_with_validation(provider, &texts, "scene consolidation")
149 .instrument(span)
150 .await;
151 let embedded: Vec<((MessageId, String), Vec<f32>)> = match vecs {
152 EmbedBatchOutcome::Ok(vecs) => candidates.into_iter().zip(vecs).collect(),
153 EmbedBatchOutcome::Skip => return Ok(stats),
154 };
155
156 if embedded.len() < 2 {
157 return Ok(stats);
158 }
159
160 let clusters = cluster_by_cosine_similarity(embedded, config.similarity_threshold);
162
163 for cluster in clusters {
164 if cluster.len() < 2 {
165 continue;
166 }
167
168 let contents: Vec<&str> = cluster.iter().map(|((_, c), _)| c.as_str()).collect();
169 let msg_ids: Vec<MessageId> = cluster.iter().map(|((id, _), _)| *id).collect();
170
171 match generate_scene_label_and_profile(provider, &contents).await {
172 Ok((label, profile)) => {
173 let label = label.chars().take(100).collect::<String>();
174 let profile = profile.chars().take(2000).collect::<String>();
175 match store.insert_mem_scene(&label, &profile, &msg_ids).await {
176 Ok(_scene_id) => {
177 stats.scenes_created += 1;
178 stats.messages_assigned += msg_ids.len();
179 }
180 Err(e) => {
181 tracing::warn!(
182 error = %e,
183 cluster_size = msg_ids.len(),
184 "scene consolidation: failed to insert scene"
185 );
186 }
187 }
188 }
189 Err(e) => {
190 tracing::warn!(
191 error = %e,
192 cluster_size = msg_ids.len(),
193 "scene consolidation: LLM label generation failed, skipping cluster"
194 );
195 }
196 }
197 }
198
199 Ok(stats)
200}
201
202async fn generate_scene_label_and_profile(
203 provider: &AnyProvider,
204 contents: &[&str],
205) -> Result<(String, String), MemoryError> {
206 use zeph_llm::provider::{Message, MessageMetadata, Role};
207
208 let bullet_list: String = contents
209 .iter()
210 .enumerate()
211 .map(|(i, c)| format!("{}. {c}", i + 1))
212 .collect::<Vec<_>>()
213 .join("\n");
214
215 let system_content = "You are a memory scene architect. \
216 Given a set of related semantic facts, generate:\n\
217 1. A short label (5 words max) identifying the core entity or topic.\n\
218 2. A 2-3 sentence entity profile summarizing the key facts.\n\
219 Respond in JSON: {\"label\": \"...\", \"profile\": \"...\"}";
220
221 let user_content =
222 format!("Generate a label and profile for these related facts:\n\n{bullet_list}");
223
224 let messages = vec![
225 Message {
226 role: Role::System,
227 content: system_content.to_owned(),
228 parts: vec![],
229 metadata: MessageMetadata::default(),
230 },
231 Message {
232 role: Role::User,
233 content: user_content,
234 parts: vec![],
235 metadata: MessageMetadata::default(),
236 },
237 ];
238
239 let result =
240 crate::sweep_helpers::llm_call_with_timeout(provider, &messages, "scene LLM call").await?;
241
242 parse_label_profile(&result)
243}
244
245fn parse_label_profile(response: &str) -> Result<(String, String), MemoryError> {
246 if let Ok(val) = serde_json::from_str::<serde_json::Value>(response) {
248 let label = val
249 .get("label")
250 .and_then(|v| v.as_str())
251 .unwrap_or("")
252 .trim()
253 .to_owned();
254 let profile = val
255 .get("profile")
256 .and_then(|v| v.as_str())
257 .unwrap_or("")
258 .trim()
259 .to_owned();
260 if !label.is_empty() && !profile.is_empty() {
261 return Ok((label, profile));
262 }
263 }
264 let trimmed = response.trim();
266 let mut lines = trimmed.splitn(2, '\n');
267 let label = lines.next().unwrap_or("").trim().to_owned();
268 let profile = lines.next().unwrap_or(trimmed).trim().to_owned();
269 if label.is_empty() {
270 return Err(MemoryError::InvalidInput(
271 "scene LLM returned empty label".into(),
272 ));
273 }
274 let profile = if profile.is_empty() {
275 label.clone()
276 } else {
277 profile
278 };
279 Ok((label, profile))
280}
281
282pub async fn list_scenes(store: &SqliteStore) -> Result<Vec<MemScene>, MemoryError> {
288 store.list_mem_scenes().await
289}
290
291#[cfg(test)]
292mod tests {
293 use super::*;
294
295 #[test]
296 fn parse_label_profile_valid_json() {
297 let json = r#"{"label": "Rust Auth JWT", "profile": "The project uses JWT for auth."}"#;
298 let (label, profile) = parse_label_profile(json).unwrap();
299 assert_eq!(label, "Rust Auth JWT");
300 assert_eq!(profile, "The project uses JWT for auth.");
301 }
302
303 #[test]
304 fn parse_label_profile_fallback_lines() {
305 let text = "Rust Auth\nJWT tokens used for authentication. Rate limited at 100 rps.";
306 let (label, profile) = parse_label_profile(text).unwrap();
307 assert_eq!(label, "Rust Auth");
308 assert!(profile.contains("JWT"));
309 }
310
311 #[test]
312 fn parse_label_profile_empty_fails() {
313 assert!(parse_label_profile("").is_err());
314 }
315}