Skip to main content

nexus_memory_agent/
soul.rs

1//! Soul management for unified user identity and cross-project learnings.
2
3use std::fs;
4use std::path::PathBuf;
5use std::sync::Arc;
6
7use anyhow::Context;
8
9use nexus_llm::{ChatMessage, GenerateParams, LlmClient};
10use serde::{Deserialize, Serialize};
11use tracing::{debug, info, warn};
12
13use crate::prompts::{SOUL_EVALUATION_PROMPT, SOUL_NORMALIZATION_PROMPT};
14
15/// Maximum token budget for the soul document.
16const SOUL_MAX_TOKENS: usize = 2048;
17
18/// Path to the unified soul.md file.
19pub fn soul_path() -> PathBuf {
20    dirs::config_dir()
21        .unwrap_or_else(|| PathBuf::from("."))
22        .join("nexus")
23        .join("soul.md")
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27#[serde(rename_all = "PascalCase")]
28pub enum SoulCategory {
29    IdentityPreference,
30    TechnicalLearning,
31    WorkingPattern,
32    AgentNote,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct SoulCandidate {
37    pub content: String,
38    pub source_project: String,
39    pub observation_count: u32,
40    pub category: String,
41    pub source_agent: String,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct NormalizedLearning {
46    pub content: String,
47    pub category: SoulCategory,
48    pub confidence: f32,
49    pub observation_count: u32,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize)]
53struct NormalizationResponse {
54    pub normalized: Vec<NormalizedLearning>,
55    pub discarded_count: usize,
56}
57
58pub struct SoulBuilder {
59    llm: Arc<dyn LlmClient>,
60}
61
62impl SoulBuilder {
63    pub fn new(llm: Arc<dyn LlmClient>) -> Self {
64        Self { llm }
65    }
66
67    /// Read the current soul document.
68    /// Returns Ok(String) on success, Ok(empty) if not found, Err for I/O errors.
69    pub fn read_current_soul(&self) -> anyhow::Result<String> {
70        let path = soul_path();
71        match fs::read_to_string(&path) {
72            Ok(soul) => Ok(soul),
73            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(String::new()),
74            Err(e) => Err(e).with_context(|| format!("Failed to read {}", path.display())),
75        }
76    }
77
78    /// Step 3: Normalization Gate
79    pub async fn normalize_candidates(
80        &self,
81        candidates: &[SoulCandidate],
82    ) -> anyhow::Result<Vec<NormalizedLearning>> {
83        if candidates.is_empty() {
84            return Ok(Vec::new());
85        }
86
87        let candidates_json = serde_json::to_string_pretty(candidates)?;
88        let system_prompt = SOUL_NORMALIZATION_PROMPT;
89        let user_prompt = format!(
90            "Normalize the following project-specific candidates for the user's Soul:\n\n{}",
91            candidates_json
92        );
93
94        let messages = vec![
95            ChatMessage::system(system_prompt),
96            ChatMessage::user(user_prompt),
97        ];
98
99        let params = GenerateParams {
100            messages,
101            json_mode: true,
102            ..Default::default()
103        };
104
105        let response = self.llm.generate(params).await?;
106
107        let content = response.content.trim();
108        let clean_response = if let Some(start) = content.find('{') {
109            if let Some(end) = content.rfind('}') {
110                if end > start {
111                    &content[start..=end]
112                } else {
113                    content
114                }
115            } else {
116                content
117            }
118        } else {
119            content
120        };
121
122        match serde_json::from_str::<NormalizationResponse>(clean_response) {
123            Ok(res) => {
124                let filtered = res
125                    .normalized
126                    .into_iter()
127                    .filter(|l| l.confidence >= 0.70)
128                    .collect();
129                Ok(filtered)
130            }
131            Err(e) => {
132                warn!(
133                    "Soul normalization parse failed; response length: {} chars, error: {}",
134                    clean_response.len(),
135                    e
136                );
137                Err(e).context("Failed to parse soul normalization response")
138            }
139        }
140    }
141
142    /// Step 4: Evaluation Gate & Rebuild
143    pub async fn evaluate_and_merge(
144        &self,
145        current_soul: &str,
146        normalized: &[NormalizedLearning],
147    ) -> anyhow::Result<String> {
148        if normalized.is_empty() {
149            return Ok(current_soul.to_string());
150        }
151
152        let normalized_json = serde_json::to_string_pretty(normalized)?;
153        let system_prompt = SOUL_EVALUATION_PROMPT;
154        let user_prompt = format!(
155            "CURRENT SOUL:\n{}\n\nNEW CANDIDATES:\n{}",
156            current_soul, normalized_json
157        );
158
159        let messages = vec![
160            ChatMessage::system(system_prompt),
161            ChatMessage::user(user_prompt),
162        ];
163
164        let params = GenerateParams {
165            messages,
166            ..Default::default()
167        };
168
169        let response = self.llm.generate(params).await?;
170        let new_soul = response.content;
171
172        let required_headers = [
173            "# Nexus Soul",
174            "## Identity & Preferences",
175            "## Technical Learnings",
176            "## Working Patterns",
177            "## Agent Notes",
178        ];
179        if new_soul.len() < 50 || required_headers.iter().any(|h| !new_soul.contains(h)) {
180            warn!("LLM returned invalid or incomplete soul document. Keeping existing.");
181            return Ok(current_soul.to_string());
182        }
183
184        Ok(new_soul)
185    }
186
187    /// Full pipeline: normalize -> evaluate -> write
188    pub async fn rebuild_soul(&self, candidates: &[SoulCandidate]) -> anyhow::Result<String> {
189        info!(
190            "Starting Soul rebuild pipeline with {} candidates",
191            candidates.len()
192        );
193
194        let normalized = self.normalize_candidates(candidates).await?;
195        debug!("Normalized into {} general learnings", normalized.len());
196
197        let current = self.read_current_soul()?;
198        let mut new_soul = self.evaluate_and_merge(&current, &normalized).await?;
199
200        if new_soul.len() / 4 > SOUL_MAX_TOKENS {
201            debug!("Soul exceeded budget. Attempting compression.");
202            let messages = vec![
203                ChatMessage::system("You are a text compression engine. Compress the following markdown soul profile while preserving all high-signal patterns. Keep it under 1500 words."),
204                ChatMessage::user(&new_soul),
205            ];
206            let params = GenerateParams {
207                messages,
208                ..Default::default()
209            };
210            if let Ok(res) = self.llm.generate(params).await {
211                let compressed = res.content;
212                let has_header = compressed.contains("# Nexus Soul");
213                let within_budget = compressed.len() / 4 <= SOUL_MAX_TOKENS;
214                if has_header && within_budget {
215                    new_soul = compressed;
216                } else {
217                    debug!(
218                        has_header,
219                        within_budget,
220                        compressed_len = compressed.len(),
221                        "Soul compression output failed validation; keeping pre-compression version"
222                    );
223                }
224            }
225        }
226
227        // Enforce token budget — refuse to write oversized soul even if compression failed
228        if new_soul.len() / 4 > SOUL_MAX_TOKENS {
229            anyhow::bail!(
230                "Refusing to write oversized soul: {} estimated tokens > {}",
231                new_soul.len() / 4,
232                SOUL_MAX_TOKENS
233            );
234        }
235
236        let path = soul_path();
237        if path.exists() {
238            let bak = path.with_extension("md.bak");
239            fs::copy(&path, &bak).with_context(|| {
240                format!(
241                    "Failed to create backup at {}. Aborting soul rebuild.",
242                    bak.display()
243                )
244            })?;
245        }
246
247        if let Some(parent) = path.parent() {
248            fs::create_dir_all(parent).with_context(|| {
249                format!("Failed to create soul directory at {}", parent.display())
250            })?;
251        }
252        nexus_core::fsutil::atomic_write(&path, &new_soul)?;
253
254        info!(
255            "Soul rebuild complete. Wrote {} bytes to {}",
256            new_soul.len(),
257            path.display()
258        );
259        Ok(new_soul)
260    }
261}