Skip to main content

zeph_memory/
scenes.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! `MemScene` consolidation (#2332).
5//!
6//! Groups semantically related semantic-tier messages into stable entity profiles (scenes).
7//! Runs as a separate background loop, decoupled from tier promotion timing.
8
9use 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/// A `MemScene` groups semantically related semantic-tier messages with a stable entity profile.
25///
26/// Scenes are created and updated by [`start_scene_consolidation_loop`] and can be
27/// listed via [`list_scenes`].
28#[derive(Debug, Clone)]
29pub struct MemScene {
30    /// `SQLite` row ID of the scene.
31    pub id: MemSceneId,
32    /// Short human-readable label for the scene (e.g. `"Rust programming"`).
33    pub label: String,
34    /// LLM-generated entity profile summarising the scene's members.
35    pub profile: String,
36    /// Number of messages currently assigned to this scene.
37    pub member_count: u32,
38    /// Unix timestamp when the scene was first created.
39    pub created_at: i64,
40    /// Unix timestamp of the last profile update.
41    pub updated_at: i64,
42}
43
44/// Configuration for the scene consolidation background loop.
45#[derive(Debug, Clone)]
46pub struct SceneConfig {
47    /// Enable or disable the scene consolidation loop.
48    pub enabled: bool,
49    /// Minimum cosine similarity for two messages to be assigned to the same scene.
50    pub similarity_threshold: f32,
51    /// Maximum number of unassigned messages to process per sweep.
52    pub batch_size: usize,
53    /// How often to run a sweep, in seconds.
54    pub sweep_interval_secs: u64,
55}
56
57/// Start the background scene consolidation loop.
58///
59/// Each sweep clusters unassigned semantic-tier messages into `MemScenes`.
60/// Runs independently from the tier promotion loop.
61pub 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    // Skip first tick to avoid running immediately at startup.
74    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/// Stats collected during a single scene consolidation sweep.
110#[derive(Debug, Default)]
111pub struct SceneStats {
112    pub candidates: usize,
113    pub scenes_created: usize,
114    pub messages_assigned: usize,
115}
116
117/// Execute one full scene consolidation sweep.
118///
119/// # Errors
120///
121/// Returns an error if the `SQLite` query fails. LLM and embedding errors are logged but skipped.
122#[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    // Embed all candidates in a single batch call.
146    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    // Cluster by cosine similarity.
161    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    // Try JSON parsing first.
247    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    // Fallback: treat first line as label, rest as profile.
265    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
282/// List all `MemScenes` from the store.
283///
284/// # Errors
285///
286/// Returns an error if the `SQLite` query fails.
287pub 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}