Skip to main content

agent_diva_agent/
context.rs

1//! Context builder for assembling prompts
2
3use crate::skills::SkillsLoader;
4use agent_diva_core::memory::MemoryManager;
5use agent_diva_core::soul::SoulStateStore;
6use agent_diva_providers::Message;
7use agent_diva_tools::sanitize::truncate_tool_result;
8use std::path::Path;
9use std::path::PathBuf;
10
11const DEFAULT_AGENT_NAME: &str = "agent-diva";
12const DEFAULT_AGENT_EMOJI: &str = "🐈";
13const DEFAULT_AGENT_ROLE: &str = "helpful AI assistant";
14
15/// Runtime controls for soul prompt injection.
16#[derive(Debug, Clone)]
17pub struct SoulContextSettings {
18    pub enabled: bool,
19    pub max_chars: usize,
20    pub bootstrap_once: bool,
21}
22
23impl Default for SoulContextSettings {
24    fn default() -> Self {
25        Self {
26            enabled: true,
27            max_chars: 4000,
28            bootstrap_once: true,
29        }
30    }
31}
32
33/// Builds the context for LLM requests
34pub struct ContextBuilder {
35    workspace: PathBuf,
36    skills_loader: SkillsLoader,
37    memory_manager: MemoryManager,
38    soul_settings: SoulContextSettings,
39}
40
41impl ContextBuilder {
42    /// Create a new context builder
43    pub fn new(workspace: PathBuf) -> Self {
44        let skills_loader = SkillsLoader::new(&workspace, None);
45        let memory_manager = MemoryManager::new(&workspace);
46        Self {
47            workspace,
48            skills_loader,
49            memory_manager,
50            soul_settings: SoulContextSettings::default(),
51        }
52    }
53
54    /// Create a new context builder with skills
55    pub fn with_skills(workspace: PathBuf, builtin_skills_dir: Option<PathBuf>) -> Self {
56        let skills_loader = SkillsLoader::new(&workspace, builtin_skills_dir);
57        let memory_manager = MemoryManager::new(&workspace);
58        Self {
59            workspace,
60            skills_loader,
61            memory_manager,
62            soul_settings: SoulContextSettings::default(),
63        }
64    }
65
66    /// Override soul context settings.
67    pub fn set_soul_settings(&mut self, settings: SoulContextSettings) {
68        self.soul_settings = settings;
69    }
70
71    /// Build system prompt from workspace files and memory
72    pub fn build_system_prompt(&self) -> String {
73        let workspace_path = self.workspace.display();
74        let now = chrono::Local::now().format("%Y-%m-%d %H:%M (%A)");
75        let identity_header = self.load_identity_header();
76
77        let mut prompt = format!(
78            r#"{identity_header}
79
80You have access to tools that allow you to:
81- Read, write, and edit files
82- Execute shell commands
83- Search the web and fetch web pages
84- Send messages to users on chat channels
85- Schedule reminders and recurring jobs (cron)
86
87## Current Time
88{now}
89
90## Workspace
91Your workspace is at: {workspace_path}
92- Memory files: {workspace_path}/memory/MEMORY.md
93- Memory history log: {workspace_path}/memory/HISTORY.md"#
94        );
95
96        if self.soul_settings.enabled {
97            self.append_soul_sections(&mut prompt);
98        }
99
100        // Skills - progressive loading
101        // 1) Always-loaded skills (full content)
102        let always_skills = self.skills_loader.get_always_skills();
103        if !always_skills.is_empty() {
104            let always_content = self.skills_loader.load_skills_for_context(&always_skills);
105            if !always_content.is_empty() {
106                prompt.push_str("\n\n## Active Skills\n");
107                prompt.push_str(&always_content);
108            }
109        }
110
111        // 2) Available skills summary
112        let skills_summary = self.skills_loader.build_skills_summary();
113        if !skills_summary.is_empty() {
114            prompt.push_str("\n\n## Skills\n");
115            prompt.push_str(
116                "The following skills extend your capabilities. To use a skill, read its SKILL.md file using the read_file tool.\n",
117            );
118            prompt
119                .push_str("Skills with available=\"false\" need dependencies installed first.\n\n");
120            prompt.push_str(&skills_summary);
121        }
122
123        // Inject long-term memory if available
124        let memory_context = self.memory_manager.get_memory_context();
125        if !memory_context.is_empty() {
126            prompt.push_str("\n\n");
127            prompt.push_str(&memory_context);
128        }
129
130        prompt.push_str(
131            r#"
132
133IMPORTANT: When responding to direct questions or conversations, reply directly with your text response.
134Only use the 'message' tool when you need to send a message to a specific chat channel.
135For normal conversation, just respond with text - do not call the message tool.
136When a user asks to create a reminder, timer, or recurring schedule, use the 'cron' tool instead of saying the feature is unavailable.
137
138Always be helpful, accurate, and concise. When using tools, explain what you're doing."#,
139        );
140
141        prompt.push_str(&format!(
142            "\nWhen remembering something, write to {}/memory/MEMORY.md",
143            workspace_path
144        ));
145
146        prompt
147    }
148
149    fn append_soul_sections(&self, prompt: &mut String) {
150        let sections = [
151            ("AGENTS.md", "Agent Rules"),
152            ("SOUL.md", "Soul"),
153            ("IDENTITY.md", "Identity"),
154            ("USER.md", "User Profile"),
155        ];
156
157        for (rel, title) in sections {
158            if let Some(content) = self.read_soul_file(rel) {
159                self.append_section(prompt, title, &content);
160            }
161        }
162
163        if self.should_include_bootstrap() {
164            if let Some(content) = self.read_soul_file("BOOTSTRAP.md") {
165                let _ = SoulStateStore::new(&self.workspace).mark_bootstrap_seeded();
166                self.append_section(prompt, "Bootstrap", &content);
167            }
168        }
169    }
170
171    fn should_include_bootstrap(&self) -> bool {
172        if !self.soul_settings.bootstrap_once {
173            return true;
174        }
175        let store = SoulStateStore::new(&self.workspace);
176        !store.is_bootstrap_completed()
177    }
178
179    fn read_soul_file(&self, rel: &str) -> Option<String> {
180        let path = self.workspace.join(rel);
181        read_trimmed_markdown(&path, self.soul_settings.max_chars)
182    }
183
184    fn append_section(&self, prompt: &mut String, title: &str, content: &str) {
185        prompt.push_str("\n\n## ");
186        prompt.push_str(title);
187        prompt.push('\n');
188        prompt.push_str(content);
189    }
190
191    fn load_identity_header(&self) -> String {
192        let Some(content) = self.read_soul_file("IDENTITY.md") else {
193            return default_identity_header();
194        };
195
196        let name = parse_identity_field(&content, &["name", "agent", "assistant"])
197            .unwrap_or_else(|| DEFAULT_AGENT_NAME.to_string());
198        let emoji = parse_identity_field(&content, &["emoji", "icon", "signature"])
199            .unwrap_or_else(|| DEFAULT_AGENT_EMOJI.to_string());
200        let role = parse_identity_field(&content, &["role", "nature", "type"])
201            .unwrap_or_else(|| DEFAULT_AGENT_ROLE.to_string());
202        let voice = parse_identity_field(&content, &["voice", "style", "vibe"]);
203
204        let mut header = format!("# {} {}\n\nYou are {}, a {}.", name, emoji, name, role);
205        if let Some(voice) = voice {
206            header.push_str(" Preferred communication style: ");
207            header.push_str(&voice);
208            header.push('.');
209        }
210        header
211    }
212
213    /// Build the complete message list for an LLM call
214    pub fn build_messages(
215        &self,
216        history: Vec<agent_diva_core::session::ChatMessage>,
217        current_message: String,
218        channel: Option<&str>,
219        chat_id: Option<&str>,
220    ) -> Vec<Message> {
221        let mut messages = Vec::new();
222
223        // System prompt
224        let mut system_prompt = self.build_system_prompt();
225        if let (Some(ch), Some(id)) = (channel, chat_id) {
226            system_prompt.push_str(&format!(
227                "\n\n## Current Session\nChannel: {}\nChat ID: {}",
228                ch, id
229            ));
230        }
231        messages.push(Message::system(system_prompt));
232
233        // History - convert from ChatMessage to Message
234        for msg in history {
235            let message = match msg.role.as_str() {
236                "user" => Message::user(&msg.content),
237                "assistant" => {
238                    let mut m = Message::assistant(&msg.content);
239                    // Restore tool_calls from session history
240                    if let Some(ref tc_values) = msg.tool_calls {
241                        if let Ok(calls) =
242                            serde_json::from_value::<Vec<agent_diva_providers::ToolCallRequest>>(
243                                serde_json::Value::Array(tc_values.clone()),
244                            )
245                        {
246                            m.tool_calls = Some(calls);
247                        }
248                    }
249                    if let Some(reasoning) = msg.reasoning_content {
250                        m.reasoning_content = Some(reasoning);
251                    }
252                    if let Some(thinking_blocks) = msg.thinking_blocks {
253                        m.thinking_blocks = Some(thinking_blocks);
254                    }
255                    m
256                }
257                "tool" => {
258                    let tool_call_id = msg.tool_call_id.unwrap_or_default();
259                    let mut m = Message::tool(msg.content, tool_call_id);
260                    m.name = msg.name;
261                    m
262                }
263                _ => continue,
264            };
265            messages.push(message);
266        }
267
268        // Current message
269        messages.push(Message::user(current_message));
270
271        messages
272    }
273
274    /// Add a tool result to the message list
275    ///
276    /// Large tool results are truncated to prevent oversized API requests
277    /// that could cause 400 errors from LLM providers.
278    pub fn add_tool_result(
279        &self,
280        messages: &mut Vec<Message>,
281        tool_call_id: String,
282        _tool_name: String,
283        result: String,
284    ) {
285        // Truncate large tool results to prevent API errors
286        let truncated_result = truncate_tool_result(&result);
287        messages.push(Message::tool(truncated_result, tool_call_id));
288    }
289
290    /// Add an assistant message with optional tool calls
291    pub fn add_assistant_message(
292        &self,
293        messages: &mut Vec<Message>,
294        content: Option<String>,
295        tool_calls: Option<Vec<agent_diva_providers::ToolCallRequest>>,
296        reasoning_content: Option<String>,
297        thinking_blocks: Option<Vec<serde_json::Value>>,
298    ) {
299        let mut msg = Message::assistant(content.unwrap_or_default());
300        if let Some(calls) = tool_calls {
301            msg.tool_calls = Some(calls);
302        }
303        if let Some(reasoning) = reasoning_content {
304            msg.reasoning_content = Some(reasoning);
305        }
306        if let Some(blocks) = thinking_blocks {
307            msg.thinking_blocks = Some(blocks);
308        }
309        messages.push(msg);
310    }
311}
312
313impl Default for ContextBuilder {
314    fn default() -> Self {
315        Self::new(PathBuf::from("."))
316    }
317}
318
319fn read_trimmed_markdown(path: &Path, max_chars: usize) -> Option<String> {
320    let content = std::fs::read_to_string(path).ok()?;
321    let trimmed = content.trim();
322    if trimmed.is_empty() {
323        return None;
324    }
325
326    if trimmed.chars().count() <= max_chars {
327        return Some(trimmed.to_string());
328    }
329
330    let mut out = String::new();
331    for (idx, ch) in trimmed.chars().enumerate() {
332        if idx >= max_chars.saturating_sub(3) {
333            break;
334        }
335        out.push(ch);
336    }
337    out.push_str("...");
338    Some(out)
339}
340
341fn parse_identity_field(content: &str, keys: &[&str]) -> Option<String> {
342    for line in content.lines() {
343        let line = line.trim().trim_start_matches(&['-', '*'][..]).trim();
344        if line.is_empty() {
345            continue;
346        }
347
348        // Split on ASCII or full-width colon to avoid manual byte indices.
349        let (prefix, value_part) = match line.split_once(':').or_else(|| line.split_once('īŧš')) {
350            Some((p, v)) => (p.trim(), v.trim()),
351            None => continue,
352        };
353
354        for key in keys {
355            if prefix.eq_ignore_ascii_case(key) && !value_part.is_empty() {
356                return Some(value_part.to_string());
357            }
358        }
359    }
360    None
361}
362
363fn default_identity_header() -> String {
364    format!(
365        "# {} {}\n\nYou are {}, a {}.",
366        DEFAULT_AGENT_NAME, DEFAULT_AGENT_EMOJI, DEFAULT_AGENT_NAME, DEFAULT_AGENT_ROLE
367    )
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373    use std::fs;
374    use tempfile::TempDir;
375
376    #[test]
377    fn test_build_system_prompt() {
378        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
379        let prompt = builder.build_system_prompt();
380        assert!(prompt.contains("agent-diva"));
381        assert!(prompt.contains("/tmp/test"));
382    }
383
384    #[test]
385    fn test_build_messages() {
386        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
387        let messages =
388            builder.build_messages(vec![], "Hello".to_string(), Some("cli"), Some("test"));
389        assert_eq!(messages.len(), 2); // system + user
390        assert_eq!(messages[0].role, "system");
391        assert_eq!(messages[1].role, "user");
392        assert_eq!(messages[1].content, "Hello");
393    }
394
395    #[test]
396    fn test_build_system_prompt_includes_skills_sections() {
397        let workspace = TempDir::new().unwrap();
398        let skills_dir = workspace.path().join("skills");
399        fs::create_dir_all(skills_dir.join("always-skill")).unwrap();
400        fs::write(
401            skills_dir.join("always-skill").join("SKILL.md"),
402            "---\nname: always-skill\ndescription: Always loaded\nmetadata: '{\"nanobot\":{\"always\":true}}'\n---\n\n# Always skill body\n",
403        )
404        .unwrap();
405
406        let builder = ContextBuilder::with_skills(workspace.path().to_path_buf(), None);
407        let prompt = builder.build_system_prompt();
408
409        assert!(prompt.contains("## Active Skills"));
410        assert!(prompt.contains("## Skills"));
411        assert!(prompt.contains("<skills>"));
412    }
413
414    #[test]
415    fn test_add_tool_result() {
416        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
417        let mut messages = vec![Message::user("test")];
418        builder.add_tool_result(
419            &mut messages,
420            "call_123".to_string(),
421            "read_file".to_string(),
422            "file content".to_string(),
423        );
424        assert_eq!(messages.len(), 2);
425        assert_eq!(messages[1].role, "tool");
426    }
427
428    #[test]
429    fn test_add_assistant_message() {
430        let builder = ContextBuilder::new(PathBuf::from("/tmp/test"));
431        let mut messages = vec![Message::user("test")];
432
433        // Test with reasoning content
434        builder.add_assistant_message(
435            &mut messages,
436            Some("response".to_string()),
437            None,
438            Some("reasoning".to_string()),
439            None,
440        );
441
442        assert_eq!(messages.len(), 2);
443        assert_eq!(messages[1].role, "assistant");
444        assert_eq!(messages[1].content, "response");
445        assert_eq!(messages[1].reasoning_content, Some("reasoning".to_string()));
446    }
447
448    #[test]
449    fn test_build_system_prompt_includes_soul_sections_in_order() {
450        let workspace = TempDir::new().unwrap();
451        fs::write(workspace.path().join("AGENTS.md"), "# Repo Rules").unwrap();
452        fs::write(workspace.path().join("SOUL.md"), "# Core Traits").unwrap();
453        fs::write(workspace.path().join("IDENTITY.md"), "# Identity").unwrap();
454        fs::write(workspace.path().join("USER.md"), "# Preferences").unwrap();
455        fs::write(workspace.path().join("BOOTSTRAP.md"), "# Bootstrap Steps").unwrap();
456
457        let builder = ContextBuilder::new(workspace.path().to_path_buf());
458        let prompt = builder.build_system_prompt();
459
460        let idx_agents = prompt.find("## Agent Rules").unwrap();
461        let idx_soul = prompt.find("## Soul").unwrap();
462        let idx_identity = prompt.find("## Identity").unwrap();
463        let idx_user = prompt.find("## User Profile").unwrap();
464        let idx_bootstrap = prompt.find("## Bootstrap").unwrap();
465
466        assert!(idx_agents < idx_soul);
467        assert!(idx_soul < idx_identity);
468        assert!(idx_identity < idx_user);
469        assert!(idx_user < idx_bootstrap);
470    }
471
472    #[test]
473    fn test_build_system_prompt_skips_bootstrap_when_completed() {
474        let workspace = TempDir::new().unwrap();
475        fs::write(workspace.path().join("BOOTSTRAP.md"), "# Bootstrap Steps").unwrap();
476        let store = SoulStateStore::new(workspace.path());
477        let mut state = agent_diva_core::soul::SoulState::default();
478        state.bootstrap_completed_at = Some(chrono::Utc::now());
479        store.save(&state).unwrap();
480
481        let builder = ContextBuilder::new(workspace.path().to_path_buf());
482        let prompt = builder.build_system_prompt();
483        assert!(!prompt.contains("## Bootstrap"));
484    }
485
486    #[test]
487    fn test_read_trimmed_markdown_respects_char_limit() {
488        let temp = TempDir::new().unwrap();
489        let path = temp.path().join("SOUL.md");
490        fs::write(&path, "abcdefghij").unwrap();
491
492        let got = read_trimmed_markdown(&path, 6).unwrap();
493        assert_eq!(got, "abc...");
494        assert!(got.chars().count() <= 6);
495    }
496
497    #[test]
498    fn test_build_system_prompt_uses_identity_file_for_header() {
499        let workspace = TempDir::new().unwrap();
500        fs::write(
501            workspace.path().join("IDENTITY.md"),
502            "# Identity\n- Name: Nova\n- Emoji: ✨\n- Role: strategic coding partner\n- Style: concise and direct\n",
503        )
504        .unwrap();
505        let builder = ContextBuilder::new(workspace.path().to_path_buf());
506        let prompt = builder.build_system_prompt();
507        assert!(prompt.contains("# Nova ✨"));
508        assert!(prompt.contains("You are Nova, a strategic coding partner."));
509        assert!(prompt.contains("Preferred communication style: concise and direct."));
510    }
511
512    #[test]
513    fn test_build_system_prompt_identity_header_falls_back_to_default() {
514        let workspace = TempDir::new().unwrap();
515        let builder = ContextBuilder::new(workspace.path().to_path_buf());
516        let prompt = builder.build_system_prompt();
517        assert!(prompt.contains("# agent-diva 🐈"));
518        assert!(prompt.contains("You are agent-diva, a helpful AI assistant."));
519    }
520
521    #[test]
522    fn test_build_system_prompt_empty_identity_falls_back_to_default() {
523        let workspace = TempDir::new().unwrap();
524        fs::write(workspace.path().join("IDENTITY.md"), "   \n").unwrap();
525        let builder = ContextBuilder::new(workspace.path().to_path_buf());
526        let prompt = builder.build_system_prompt();
527        assert!(prompt.contains("# agent-diva 🐈"));
528    }
529
530    #[test]
531    fn test_build_system_prompt_long_identity_is_trimmed_by_max_chars() {
532        let workspace = TempDir::new().unwrap();
533        let long_name = "N".repeat(6000);
534        fs::write(
535            workspace.path().join("IDENTITY.md"),
536            format!("- Name: {}\n- Role: helper", long_name),
537        )
538        .unwrap();
539        let mut builder = ContextBuilder::new(workspace.path().to_path_buf());
540        builder.set_soul_settings(SoulContextSettings {
541            enabled: true,
542            max_chars: 120,
543            bootstrap_once: true,
544        });
545        let prompt = builder.build_system_prompt();
546        assert!(prompt.contains("You are "));
547        assert!(prompt.chars().count() > 120);
548    }
549
550    #[test]
551    fn test_parse_identity_field_handles_markdown_list() {
552        let raw = "- Name: Diva\n- Style: pragmatic";
553        assert_eq!(
554            parse_identity_field(raw, &["name"]).as_deref(),
555            Some("Diva")
556        );
557        assert_eq!(
558            parse_identity_field(raw, &["style"]).as_deref(),
559            Some("pragmatic")
560        );
561    }
562
563    #[test]
564    fn test_parse_identity_field_supports_chinese_voice_line() {
565        let raw = "- Voice: įŽ€æ´ã€åŽžį”¨ã€åäŊœ";
566        assert_eq!(
567            parse_identity_field(raw, &["voice"]).as_deref(),
568            Some("įŽ€æ´ã€åŽžį”¨ã€åäŊœ")
569        );
570    }
571}