mofa_foundation/llm/
context.rs

1//! Agent context builder framework
2//!
3//! This module provides:
4//! - Flexible system prompt building
5//! - Bootstrap file loading from workspace
6//! - Agent identity integration
7//! - Vision message support
8
9use crate::llm::types::{ChatMessage, ContentPart, ImageUrl, MessageContent, Role};
10use anyhow::Result;
11use serde::{Deserialize, Serialize};
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14use tokio::sync::RwLock;
15
16/// Agent identity information
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct AgentIdentity {
19    /// Agent name
20    pub name: String,
21    /// Agent description
22    pub description: String,
23    /// Agent icon (emoji)
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub icon: Option<String>,
26}
27
28impl AgentIdentity {
29    /// Create a new agent identity
30    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
31        Self {
32            name: name.into(),
33            description: description.into(),
34            icon: None,
35        }
36    }
37
38    /// Set the icon
39    pub fn with_icon(mut self, icon: impl Into<String>) -> Self {
40        self.icon = Some(icon.into());
41        self
42    }
43}
44
45/// Default bootstrap files to load
46const DEFAULT_BOOTSTRAP_FILES: &[&str] =
47    &["AGENTS.md", "SOUL.md", "USER.md", "TOOLS.md", "IDENTITY.md"];
48
49/// Skills manager trait for progressive loading
50#[async_trait::async_trait]
51pub trait SkillsManager: Send + Sync {
52    /// Get skills that should always be loaded
53    async fn get_always_skills(&self) -> Vec<String>;
54
55    /// Load skills content for context
56    async fn load_skills_for_context(&self, names: &[String]) -> String;
57
58    /// Get skills summary for display
59    async fn build_skills_summary(&self) -> String;
60}
61
62/// No-op skills manager for when skills aren't needed
63pub struct NoOpSkillsManager;
64
65#[async_trait::async_trait]
66impl SkillsManager for NoOpSkillsManager {
67    async fn get_always_skills(&self) -> Vec<String> {
68        Vec::new()
69    }
70
71    async fn load_skills_for_context(&self, _names: &[String]) -> String {
72        String::new()
73    }
74
75    async fn build_skills_summary(&self) -> String {
76        String::new()
77    }
78}
79
80/// Context builder for agent prompts
81pub struct AgentContextBuilder {
82    /// Workspace path
83    workspace: PathBuf,
84    /// Bootstrap files to load
85    bootstrap_files: Vec<String>,
86    /// Agent identity
87    identity: AgentIdentity,
88    /// Optional skills manager
89    skills: Option<Arc<dyn SkillsManager>>,
90    /// Cached system prompt
91    cached_prompt: Arc<RwLock<Option<String>>>,
92}
93
94impl AgentContextBuilder {
95    /// Create a new context builder
96    pub fn new(workspace: PathBuf) -> Self {
97        Self {
98            workspace,
99            bootstrap_files: DEFAULT_BOOTSTRAP_FILES
100                .iter()
101                .map(|s| s.to_string())
102                .collect(),
103            identity: AgentIdentity {
104                name: "agent".to_string(),
105                description: "AI assistant".to_string(),
106                icon: None,
107            },
108            skills: None,
109            cached_prompt: Arc::new(RwLock::new(None)),
110        }
111    }
112
113    /// Set bootstrap files
114    pub fn with_bootstrap_files(mut self, files: Vec<String>) -> Self {
115        self.bootstrap_files = files;
116        self
117    }
118
119    /// Set agent identity
120    pub fn with_identity(mut self, identity: AgentIdentity) -> Self {
121        self.identity = identity;
122        self
123    }
124
125    /// Set skills manager
126    pub fn with_skills(mut self, skills: Arc<dyn SkillsManager>) -> Self {
127        self.skills = Some(skills);
128        self
129    }
130
131    /// Build system prompt from bootstrap files
132    pub async fn build_system_prompt(&self) -> Result<String> {
133        // Check cache
134        {
135            let cached = self.cached_prompt.read().await;
136            if let Some(prompt) = cached.as_ref() {
137                return Ok(prompt.clone());
138            }
139        }
140
141        let mut parts = Vec::new();
142
143        // Add identity header
144        parts.push(format!("# Agent: {}", self.identity.name));
145        parts.push(format!("{}\n", self.identity.description));
146
147        // Load bootstrap files
148        for filename in &self.bootstrap_files {
149            let path = self.workspace.join(filename);
150            if let Ok(content) = Self::load_file(&path) {
151                parts.push(format!("## {}\n{}", filename, content));
152            }
153        }
154
155        // Add skills section if available
156        if let Some(skills) = &self.skills {
157            let always_skills = skills.get_always_skills().await;
158            if !always_skills.is_empty() {
159                let content = skills.load_skills_for_context(&always_skills).await;
160                if !content.is_empty() {
161                    parts.push(format!("# Active Skills\n\n{}", content));
162                }
163            }
164
165            let summary = skills.build_skills_summary().await;
166            if !summary.is_empty() {
167                parts.push(format!(
168                    r#"# Skills
169
170The following skills extend your capabilities. To use a skill, read its documentation.
171
172{}"#,
173                    summary
174                ));
175            }
176        }
177
178        let prompt = parts.join("\n\n---\n\n");
179
180        // Cache the result
181        let mut cached = self.cached_prompt.write().await;
182        *cached = Some(prompt.clone());
183
184        Ok(prompt)
185    }
186
187    /// Build messages with history and current input
188    pub async fn build_messages(
189        &self,
190        history: Vec<ChatMessage>,
191        current: &str,
192        media: Option<Vec<String>>,
193    ) -> Result<Vec<ChatMessage>> {
194        let mut messages = Vec::new();
195
196        // System prompt
197        let system_prompt = self.build_system_prompt().await?;
198        messages.push(ChatMessage::system(system_prompt));
199
200        // History
201        messages.extend(history);
202
203        // Current message (with optional media)
204        let user_msg = if let Some(media_paths) = media {
205            if !media_paths.is_empty() {
206                Self::build_vision_message(current, &media_paths)?
207            } else {
208                ChatMessage::user(current)
209            }
210        } else {
211            ChatMessage::user(current)
212        };
213
214        messages.push(user_msg);
215
216        Ok(messages)
217    }
218
219    /// Build messages with skill names
220    pub async fn build_messages_with_skills(
221        &self,
222        history: Vec<ChatMessage>,
223        current: &str,
224        media: Option<Vec<String>>,
225        skill_names: Option<&[String]>,
226    ) -> Result<Vec<ChatMessage>> {
227        let mut messages = Vec::new();
228
229        // Build system prompt with optional skills
230        let system_prompt = self.build_system_prompt().await?;
231
232        let final_prompt = if let Some(skills) = &self.skills {
233            if let Some(names) = skill_names {
234                if !names.is_empty() {
235                    let skills_content = skills.load_skills_for_context(names).await;
236                    if !skills_content.is_empty() {
237                        format!(
238                            "{}\n\n# Requested Skills\n\n{}",
239                            system_prompt, skills_content
240                        )
241                    } else {
242                        system_prompt
243                    }
244                } else {
245                    system_prompt
246                }
247            } else {
248                system_prompt
249            }
250        } else {
251            system_prompt
252        };
253
254        messages.push(ChatMessage::system(final_prompt));
255
256        // History
257        messages.extend(history);
258
259        // Current message (with optional media)
260        let user_msg = if let Some(media_paths) = media {
261            if !media_paths.is_empty() {
262                Self::build_vision_message(current, &media_paths)?
263            } else {
264                ChatMessage::user(current)
265            }
266        } else {
267            ChatMessage::user(current)
268        };
269
270        messages.push(user_msg);
271
272        Ok(messages)
273    }
274
275    /// Build a vision message with images
276    fn build_vision_message(text: &str, image_paths: &[String]) -> Result<ChatMessage> {
277        let mut parts = vec![ContentPart::Text {
278            text: text.to_string(),
279        }];
280
281        for path in image_paths {
282            let image_url = Self::encode_image_data_url(Path::new(path))?;
283            parts.push(ContentPart::Image { image_url });
284        }
285
286        Ok(ChatMessage {
287            role: Role::User,
288            content: Some(MessageContent::Parts(parts)),
289            name: None,
290            tool_calls: None,
291            tool_call_id: None,
292        })
293    }
294
295    /// Encode an image file as a data URL
296    fn encode_image_data_url(path: &Path) -> Result<ImageUrl> {
297        use base64::Engine;
298        use base64::engine::general_purpose::STANDARD_NO_PAD;
299        use std::fs;
300
301        let bytes = fs::read(path)?;
302        let mime_type = infer::get_from_path(path)?
303            .ok_or_else(|| anyhow::anyhow!("Unknown MIME type for: {:?}", path))?
304            .mime_type()
305            .to_string();
306
307        let base64 = STANDARD_NO_PAD.encode(&bytes);
308        let url = format!("data:{};base64,{}", mime_type, base64);
309
310        Ok(ImageUrl { url, detail: None })
311    }
312
313    /// Load a file's content
314    fn load_file(path: &Path) -> Result<String> {
315        std::fs::read_to_string(path)
316            .map_err(|e| anyhow::anyhow!("Failed to read {:?}: {}", path, e))
317    }
318
319    /// Get the workspace path
320    pub fn workspace(&self) -> &Path {
321        &self.workspace
322    }
323
324    /// Get the agent identity
325    pub fn identity(&self) -> &AgentIdentity {
326        &self.identity
327    }
328
329    /// Clear the cached prompt
330    pub async fn clear_cache(&self) {
331        let mut cached = self.cached_prompt.write().await;
332        *cached = None;
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339
340    #[test]
341    fn test_agent_identity_new() {
342        let identity = AgentIdentity::new("test", "Test agent");
343        assert_eq!(identity.name, "test");
344        assert_eq!(identity.description, "Test agent");
345        assert!(identity.icon.is_none());
346    }
347
348    #[test]
349    fn test_agent_identity_with_icon() {
350        let identity = AgentIdentity::new("test", "Test agent").with_icon("🤖");
351        assert_eq!(identity.icon, Some("🤖".to_string()));
352    }
353
354    #[tokio::test]
355    async fn test_context_builder_new() {
356        let workspace = std::env::temp_dir();
357        let builder = AgentContextBuilder::new(workspace.clone());
358
359        assert_eq!(builder.workspace(), &workspace);
360        assert_eq!(builder.identity().name, "agent");
361    }
362
363    #[tokio::test]
364    async fn test_context_builder_with_identity() {
365        let workspace = std::env::temp_dir();
366        let identity = AgentIdentity::new("custom", "Custom agent");
367        let builder = AgentContextBuilder::new(workspace).with_identity(identity);
368
369        assert_eq!(builder.identity().name, "custom");
370    }
371}
mofa_foundation/llm/context.rs

mofa_foundation/llm/
context.rs
