Skip to main content

mofa_foundation/agent/context/
prompt.rs

1//! Prompt Context - Specialized context for prompt building
2//!
3//! This module provides prompt-building capabilities as a specialization
4//! of the core context system. It uses RichAgentContext for storage
5//! while providing domain-specific prompt building functionality.
6//!
7//! # Design
8//!
9//! - Uses `RichAgentContext` for context management
10//! - Provides fluent builder API for prompt construction
11//! - Supports bootstrap file loading and memory integration
12//! - Progressive disclosure for skills/tools
13//!
14//! # Example
15//!
16//! ```rust,ignore
17//! use mofa_foundation::agent::context::prompt::PromptContextBuilder;
18//! use mofa_kernel::agent::context::CoreAgentContext;
19//!
20//! let ctx = PromptContextBuilder::new("/path/to/workspace")
21//!     .with_name("MyAgent")
22//!     .build()
23//!     .await?;
24//!
25//! let prompt = ctx.build_system_prompt().await?;
26//! ```
27
28use chrono::Utc;
29use std::path::{Path, PathBuf};
30use tokio::fs;
31
32use super::rich::RichAgentContext;
33use crate::agent::components::memory::FileBasedStorage;
34use mofa_kernel::agent::context::AgentContext;
35use mofa_kernel::agent::error::{AgentError, AgentResult};
36use std::sync::Arc;
37
38/// Agent identity information
39#[derive(Debug, Clone)]
40pub struct AgentIdentity {
41    /// Agent name
42    pub name: String,
43    /// Agent description
44    pub description: String,
45    /// Emoji icon
46    pub icon: Option<String>,
47}
48
49impl Default for AgentIdentity {
50    fn default() -> Self {
51        Self {
52            name: "Agent".to_string(),
53            description: "A helpful AI assistant".to_string(),
54            icon: None,
55        }
56    }
57}
58
59/// Context builder for agent system prompts
60///
61/// This component standardizes how agents build their system prompts by:
62/// - Loading bootstrap files from the workspace
63/// - Injecting memory context
64/// - Managing skill loading and progressive disclosure
65/// - Providing a consistent identity section
66///
67/// Uses `RichAgentContext` internally for context management.
68pub struct PromptContext {
69    /// Base workspace directory
70    workspace: PathBuf,
71    /// Agent identity (name, description)
72    identity: AgentIdentity,
73    /// Bootstrap files to load
74    bootstrap_files: Vec<String>,
75    /// Memory storage (lazy initialized)
76    memory: Option<Arc<FileBasedStorage>>,
77    /// Always-loaded skills/tools
78    always_load: Vec<String>,
79    /// Agent name for display
80    agent_name: String,
81    /// Rich context for extended functionality
82    rich_ctx: RichAgentContext,
83}
84
85impl PromptContext {
86    /// Default bootstrap files to load
87    pub fn default_bootstrap_files() -> Vec<String> {
88        vec![
89            "AGENTS.md".to_string(),
90            "SOUL.md".to_string(),
91            "USER.md".to_string(),
92            "TOOLS.md".to_string(),
93            "IDENTITY.md".to_string(),
94        ]
95    }
96
97    /// Create a new prompt context
98    pub async fn new(workspace: impl AsRef<Path>) -> AgentResult<Self> {
99        let workspace = workspace.as_ref().to_path_buf();
100        let core_ctx = AgentContext::new(format!("prompt-{}", uuid::Uuid::new_v4()));
101        let rich_ctx = RichAgentContext::new(core_ctx);
102
103        Ok(Self {
104            workspace,
105            identity: AgentIdentity::default(),
106            bootstrap_files: Self::default_bootstrap_files(),
107            memory: None,
108            always_load: Vec::new(),
109            agent_name: "agent".to_string(),
110            rich_ctx,
111        })
112    }
113
114    /// Create with custom identity
115    pub async fn with_identity(
116        workspace: impl AsRef<Path>,
117        identity: AgentIdentity,
118    ) -> AgentResult<Self> {
119        let workspace = workspace.as_ref().to_path_buf();
120        let agent_name = identity.name.clone();
121        let core_ctx = AgentContext::new(format!("prompt-{}", uuid::Uuid::new_v4()));
122        let rich_ctx = RichAgentContext::new(core_ctx);
123
124        Ok(Self {
125            workspace,
126            identity,
127            bootstrap_files: Self::default_bootstrap_files(),
128            memory: None,
129            always_load: Vec::new(),
130            agent_name,
131            rich_ctx,
132        })
133    }
134
135    /// Set the bootstrap files to load
136    pub fn with_bootstrap_files(mut self, files: Vec<String>) -> Self {
137        self.bootstrap_files = files;
138        self
139    }
140
141    /// Set skills that should always be loaded
142    pub fn with_always_load(mut self, skills: Vec<String>) -> Self {
143        self.always_load = skills;
144        self
145    }
146
147    /// Initialize memory storage (lazy)
148    async fn init_memory(&mut self) -> AgentResult<()> {
149        if self.memory.is_none() {
150            self.memory = Some(Arc::new(
151                FileBasedStorage::new(&self.workspace).await.map_err(|e| {
152                    AgentError::MemoryError(format!("Failed to init memory: {}", e))
153                })?,
154            ));
155        }
156        Ok(())
157    }
158
159    /// Build the complete system prompt
160    pub async fn build_system_prompt(&mut self) -> AgentResult<String> {
161        let mut parts = Vec::new();
162
163        // 1. Core identity
164        parts.push(self.get_identity_section());
165
166        // 2. Bootstrap files
167        let bootstrap = self.load_bootstrap_files().await?;
168        if !bootstrap.is_empty() {
169            parts.push(bootstrap);
170        }
171
172        // 3. Memory context (lazy init)
173        if let Err(_) = self.init_memory().await {
174            // Memory is optional, continue without it
175        } else if let Some(memory) = &self.memory
176            && let Ok(memory_context) = memory.get_memory_context().await
177            && !memory_context.is_empty()
178        {
179            parts.push(format!("# Memory\n\n{}", memory_context));
180        }
181
182        // 4. Record that we built a prompt (using rich context)
183        self.rich_ctx
184            .record_output(
185                "prompt_builder",
186                serde_json::json!({
187                    "prompt_length": parts.join("\n\n---\n\n").len(),
188                    "bootstrap_files": self.bootstrap_files.len(),
189                }),
190            )
191            .await;
192
193        Ok(parts.join("\n\n---\n\n"))
194    }
195
196    /// Get the identity section of the system prompt
197    fn get_identity_section(&self) -> String {
198        let now = Utc::now().format("%Y-%m-%d %H:%M (%A)");
199        let workspace_path = self.workspace.display();
200        let icon = self.identity.icon.as_deref().unwrap_or("");
201        let description = if self.identity.description.is_empty() {
202            "a helpful AI assistant"
203        } else {
204            &self.identity.description
205        };
206
207        format!(
208            r#"# {} {} {}
209
210You are {}, {}.
211
212## Current Time
213{}
214
215## Workspace
216Your workspace is at: {}
217- Memory files: {}/memory/MEMORY.md
218- Daily notes: {}/memory/YYYY-MM-DD.md
219- Custom skills: {}/skills/{{{{skill-name}}}}/SKILL.md
220
221Always be helpful, accurate, and concise. When using tools, explain what you're doing.
222When remembering something, write to {}/memory/MEMORY.md"#,
223            icon,
224            self.identity.name,
225            description,
226            self.identity.name,
227            description,
228            now,
229            workspace_path,
230            workspace_path,
231            workspace_path,
232            workspace_path,
233            workspace_path
234        )
235    }
236
237    /// Load bootstrap files from workspace
238    async fn load_bootstrap_files(&self) -> AgentResult<String> {
239        let mut parts = Vec::new();
240
241        for filename in &self.bootstrap_files {
242            let file_path = self.workspace.join(filename);
243            if file_path.exists()
244                && let Ok(content) = fs::read_to_string(&file_path).await
245            {
246                parts.push(format!("## {}\n\n{}", filename, content));
247            }
248        }
249
250        Ok(parts.join("\n\n"))
251    }
252
253    /// Get memory storage reference
254    pub async fn memory(&mut self) -> AgentResult<&FileBasedStorage> {
255        self.init_memory().await?;
256        Ok(self.memory.as_ref().unwrap())
257    }
258
259    /// Get the workspace path
260    pub fn workspace(&self) -> &Path {
261        &self.workspace
262    }
263
264    /// Get the rich context for extended functionality
265    pub fn rich_context(&self) -> &RichAgentContext {
266        &self.rich_ctx
267    }
268
269    /// Get the identity
270    pub fn identity(&self) -> &AgentIdentity {
271        &self.identity
272    }
273}
274
275/// Builder for creating PromptContext with fluent API
276pub struct PromptContextBuilder {
277    workspace: PathBuf,
278    identity: AgentIdentity,
279    bootstrap_files: Vec<String>,
280    always_load: Vec<String>,
281}
282
283impl PromptContextBuilder {
284    /// Create a new builder
285    pub fn new(workspace: impl AsRef<Path>) -> Self {
286        Self {
287            workspace: workspace.as_ref().to_path_buf(),
288            identity: AgentIdentity::default(),
289            bootstrap_files: PromptContext::default_bootstrap_files(),
290            always_load: Vec::new(),
291        }
292    }
293
294    /// Set the agent identity
295    pub fn with_identity(mut self, identity: AgentIdentity) -> Self {
296        self.identity = identity;
297        self
298    }
299
300    /// Set the agent name
301    pub fn with_name(mut self, name: impl Into<String>) -> Self {
302        self.identity.name = name.into();
303        self
304    }
305
306    /// Set the agent description
307    pub fn with_description(mut self, description: impl Into<String>) -> Self {
308        self.identity.description = description.into();
309        self
310    }
311
312    /// Set the agent icon
313    pub fn with_icon(mut self, icon: impl Into<String>) -> Self {
314        self.identity.icon = Some(icon.into());
315        self
316    }
317
318    /// Set bootstrap files
319    pub fn with_bootstrap_files(mut self, files: Vec<String>) -> Self {
320        self.bootstrap_files = files;
321        self
322    }
323
324    /// Set always-loaded skills
325    pub fn with_always_load(mut self, skills: Vec<String>) -> Self {
326        self.always_load = skills;
327        self
328    }
329
330    /// Build the PromptContext
331    pub async fn build(self) -> AgentResult<PromptContext> {
332        let agent_name = self.identity.name.clone();
333        PromptContext::with_identity(&self.workspace, self.identity)
334            .await
335            .map(|mut ctx| {
336                ctx.bootstrap_files = self.bootstrap_files;
337                ctx.always_load = self.always_load;
338                ctx.agent_name = agent_name;
339                ctx
340            })
341    }
342}
343
344impl Clone for PromptContext {
345    fn clone(&self) -> Self {
346        Self {
347            workspace: self.workspace.clone(),
348            identity: self.identity.clone(),
349            bootstrap_files: self.bootstrap_files.clone(),
350            memory: self.memory.clone(),
351            always_load: self.always_load.clone(),
352            agent_name: self.agent_name.clone(),
353            rich_ctx: RichAgentContext::new(self.rich_ctx.inner().clone()),
354        }
355    }
356}