enya-agent 0.1.0

package/crates/enact-context/src/calibrator.rs

@@ -0,0 +1,535 @@
+//! Prompt Calibrator
+//!
+//! Constructs calibrated prompts for spawned callables by:
+//! - Selecting relevant context segments based on priority
+//! - Applying token budgets for the target model
+//! - Formatting context for optimal model consumption
+//!
+//! @see packages/enact-schemas/src/context.schemas.ts
+
+use crate::budget::ContextBudget;
+use crate::segment::{ContextPriority, ContextSegment, ContextSegmentType};
+use crate::token_counter::TokenCounter;
+use crate::window::ContextWindow;
+use chrono::{DateTime, Utc};
+use enact_core::kernel::ExecutionId;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+/// Global sequence counter for segments
+static SEGMENT_SEQUENCE: AtomicU64 = AtomicU64::new(1000);
+
+fn next_sequence() -> u64 {
+    SEGMENT_SEQUENCE.fetch_add(1, Ordering::SeqCst)
+}
+
+/// Prompt calibration configuration
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct CalibrationConfig {
+    /// Maximum tokens for the calibrated prompt
+    pub max_tokens: usize,
+
+    /// Minimum tokens to reserve for response
+    pub response_reserve: usize,
+
+    /// Priority threshold - only include segments at or above this priority
+    pub min_priority: ContextPriority,
+
+    /// Whether to include system context
+    pub include_system: bool,
+
+    /// Whether to include conversation history
+    pub include_history: bool,
+
+    /// Maximum history messages to include
+    pub max_history_messages: usize,
+
+    /// Whether to include working memory
+    pub include_working_memory: bool,
+
+    /// Whether to include RAG context
+    pub include_rag: bool,
+
+    /// Maximum RAG chunks to include
+    pub max_rag_chunks: usize,
+
+    /// Custom segment filters by type
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub segment_filters: Option<HashMap<String, bool>>,
+}
+
+impl Default for CalibrationConfig {
+    fn default() -> Self {
+        Self {
+            max_tokens: 8000,
+            response_reserve: 2000,
+            min_priority: ContextPriority::Low,
+            include_system: true,
+            include_history: true,
+            max_history_messages: 20,
+            include_working_memory: true,
+            include_rag: true,
+            max_rag_chunks: 5,
+            segment_filters: None,
+        }
+    }
+}
+
+impl CalibrationConfig {
+    /// Create a minimal config for quick tasks
+    pub fn minimal() -> Self {
+        Self {
+            max_tokens: 4000,
+            response_reserve: 1000,
+            min_priority: ContextPriority::High,
+            include_system: true,
+            include_history: false,
+            max_history_messages: 0,
+            include_working_memory: false,
+            include_rag: false,
+            max_rag_chunks: 0,
+            segment_filters: None,
+        }
+    }
+
+    /// Create a full context config for complex tasks
+    pub fn full_context() -> Self {
+        Self {
+            max_tokens: 32000,
+            response_reserve: 4000,
+            min_priority: ContextPriority::Low,
+            include_system: true,
+            include_history: true,
+            max_history_messages: 50,
+            include_working_memory: true,
+            include_rag: true,
+            max_rag_chunks: 10,
+            segment_filters: None,
+        }
+    }
+
+    /// Available tokens after response reserve
+    pub fn available_tokens(&self) -> usize {
+        self.max_tokens.saturating_sub(self.response_reserve)
+    }
+}
+
+/// Result of prompt calibration
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct CalibratedPrompt {
+    /// Execution ID this prompt was calibrated for
+    pub execution_id: ExecutionId,
+
+    /// The calibrated segments in order
+    pub segments: Vec<ContextSegment>,
+
+    /// Total tokens used
+    pub total_tokens: usize,
+
+    /// Tokens available for response
+    pub response_tokens: usize,
+
+    /// Segments that were excluded due to budget
+    pub excluded_count: usize,
+
+    /// Calibration timestamp
+    pub calibrated_at: DateTime<Utc>,
+
+    /// Configuration used
+    pub config: CalibrationConfig,
+}
+
+impl CalibratedPrompt {
+    /// Get the formatted prompt as a string
+    pub fn as_text(&self) -> String {
+        self.segments
+            .iter()
+            .map(|s| s.content.clone())
+            .collect::<Vec<_>>()
+            .join("\n\n")
+    }
+
+    /// Get segments by type
+    pub fn segments_by_type(&self, segment_type: ContextSegmentType) -> Vec<&ContextSegment> {
+        self.segments
+            .iter()
+            .filter(|s| s.segment_type == segment_type)
+            .collect()
+    }
+
+    /// Check if prompt has system context
+    pub fn has_system(&self) -> bool {
+        self.segments
+            .iter()
+            .any(|s| s.segment_type == ContextSegmentType::System)
+    }
+
+    /// Check if prompt has history
+    pub fn has_history(&self) -> bool {
+        self.segments
+            .iter()
+            .any(|s| s.segment_type == ContextSegmentType::History)
+    }
+}
+
+/// Prompt Calibrator - constructs calibrated prompts for spawned callables
+pub struct PromptCalibrator {
+    token_counter: TokenCounter,
+}
+
+impl PromptCalibrator {
+    /// Create a new calibrator
+    pub fn new() -> Self {
+        Self {
+            token_counter: TokenCounter::default(),
+        }
+    }
+
+    /// Calibrate a prompt from a context window
+    pub fn calibrate(
+        &self,
+        window: &ContextWindow,
+        config: &CalibrationConfig,
+    ) -> CalibratedPrompt {
+        let execution_id = window.budget().execution_id.clone();
+        let available = config.available_tokens();
+
+        // Get all segments sorted by priority (highest first)
+        let mut segments = window.segments().to_vec();
+        segments.sort_by(|a, b| b.priority.cmp(&a.priority));
+
+        // Filter and select segments within budget
+        let mut selected: Vec<ContextSegment> = Vec::new();
+        let mut total_tokens = 0;
+        let mut excluded_count = 0;
+        let mut history_count = 0;
+        let mut rag_count = 0;
+
+        for segment in segments {
+            // Check priority threshold
+            if segment.priority < config.min_priority {
+                excluded_count += 1;
+                continue;
+            }
+
+            // Apply type-specific filters
+            match segment.segment_type {
+                ContextSegmentType::System if !config.include_system => {
+                    excluded_count += 1;
+                    continue;
+                }
+                ContextSegmentType::History if !config.include_history => {
+                    excluded_count += 1;
+                    continue;
+                }
+                ContextSegmentType::History if history_count >= config.max_history_messages => {
+                    excluded_count += 1;
+                    continue;
+                }
+                ContextSegmentType::WorkingMemory if !config.include_working_memory => {
+                    excluded_count += 1;
+                    continue;
+                }
+                ContextSegmentType::RagContext if !config.include_rag => {
+                    excluded_count += 1;
+                    continue;
+                }
+                ContextSegmentType::RagContext if rag_count >= config.max_rag_chunks => {
+                    excluded_count += 1;
+                    continue;
+                }
+                _ => {}
+            }
+
+            // Check if segment fits in budget
+            let segment_tokens = segment.token_count;
+            if total_tokens + segment_tokens > available {
+                excluded_count += 1;
+                continue;
+            }
+
+            // Include segment
+            total_tokens += segment_tokens;
+            if segment.segment_type == ContextSegmentType::History {
+                history_count += 1;
+            }
+            if segment.segment_type == ContextSegmentType::RagContext {
+                rag_count += 1;
+            }
+            selected.push(segment);
+        }
+
+        // Re-sort by natural order (system first, then by sequence)
+        selected.sort_by(|a, b| {
+            // System always first
+            if a.segment_type == ContextSegmentType::System
+                && b.segment_type != ContextSegmentType::System
+            {
+                return std::cmp::Ordering::Less;
+            }
+            if b.segment_type == ContextSegmentType::System
+                && a.segment_type != ContextSegmentType::System
+            {
+                return std::cmp::Ordering::Greater;
+            }
+            // Then by sequence
+            a.sequence.cmp(&b.sequence)
+        });
+
+        CalibratedPrompt {
+            execution_id,
+            segments: selected,
+            total_tokens,
+            response_tokens: config.max_tokens.saturating_sub(total_tokens),
+            excluded_count,
+            calibrated_at: Utc::now(),
+            config: config.clone(),
+        }
+    }
+
+    /// Calibrate from raw segments (without a full window)
+    pub fn calibrate_segments(
+        &self,
+        execution_id: ExecutionId,
+        segments: Vec<ContextSegment>,
+        config: &CalibrationConfig,
+    ) -> CalibratedPrompt {
+        // Create a temporary budget and window
+        let budget = ContextBudget::new(execution_id.clone(), config.max_tokens, config.response_reserve);
+        let mut window = ContextWindow::new(budget).expect("valid budget");
+
+        for segment in segments {
+            let _ = window.add_segment(segment);
+        }
+
+        self.calibrate(&window, config)
+    }
+
+    /// Create a child prompt from parent context
+    ///
+    /// This is used when spawning child callables to provide them with
+    /// relevant context from the parent execution.
+    pub fn calibrate_for_child(
+        &self,
+        parent_window: &ContextWindow,
+        child_execution_id: ExecutionId,
+        task_description: &str,
+        config: &CalibrationConfig,
+    ) -> CalibratedPrompt {
+        let available = config.available_tokens();
+
+        // Start with essential context for the child
+        let mut selected: Vec<ContextSegment> = Vec::new();
+        let mut total_tokens = 0;
+
+        // Add task description as system context
+        let task_content = format!(
+            "You are executing a sub-task. Task: {}\n\nParent context follows:",
+            task_description
+        );
+        let task_tokens = self.token_counter.count(&task_content);
+        if task_tokens <= available {
+            let task_segment = ContextSegment::system(task_content, task_tokens);
+            total_tokens += task_tokens;
+            selected.push(task_segment);
+        }
+
+        // Get parent segments sorted by priority
+        let mut parent_segments = parent_window.segments().to_vec();
+        parent_segments.sort_by(|a, b| b.priority.cmp(&a.priority));
+
+        let mut excluded_count = 0;
+
+        // Add relevant parent context
+        for segment in parent_segments {
+            // Skip low priority for child contexts
+            if segment.priority < ContextPriority::Medium {
+                excluded_count += 1;
+                continue;
+            }
+
+            // Check budget
+            let segment_tokens = segment.token_count;
+            if total_tokens + segment_tokens > available {
+                excluded_count += 1;
+                continue;
+            }
+
+            total_tokens += segment_tokens;
+            selected.push(segment);
+        }
+
+        // Sort for natural reading order
+        selected.sort_by(|a, b| {
+            if a.segment_type == ContextSegmentType::System
+                && b.segment_type != ContextSegmentType::System
+            {
+                return std::cmp::Ordering::Less;
+            }
+            if b.segment_type == ContextSegmentType::System
+                && a.segment_type != ContextSegmentType::System
+            {
+                return std::cmp::Ordering::Greater;
+            }
+            a.sequence.cmp(&b.sequence)
+        });
+
+        CalibratedPrompt {
+            execution_id: child_execution_id,
+            segments: selected,
+            total_tokens,
+            response_tokens: config.max_tokens.saturating_sub(total_tokens),
+            excluded_count,
+            calibrated_at: Utc::now(),
+            config: config.clone(),
+        }
+    }
+}
+
+impl Default for PromptCalibrator {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_execution_id() -> ExecutionId {
+        ExecutionId::new()
+    }
+
+    #[test]
+    fn test_calibration_config_defaults() {
+        let config = CalibrationConfig::default();
+        assert_eq!(config.max_tokens, 8000);
+        assert_eq!(config.response_reserve, 2000);
+        assert_eq!(config.available_tokens(), 6000);
+    }
+
+    #[test]
+    fn test_calibration_config_minimal() {
+        let config = CalibrationConfig::minimal();
+        assert!(!config.include_history);
+        assert!(!config.include_working_memory);
+        assert_eq!(config.min_priority, ContextPriority::High);
+    }
+
+    #[test]
+    fn test_calibrate_empty_window() {
+        let calibrator = PromptCalibrator::new();
+        let budget = ContextBudget::preset_default(test_execution_id());
+        let window = ContextWindow::new(budget).unwrap();
+        let config = CalibrationConfig::default();
+
+        let result = calibrator.calibrate(&window, &config);
+        assert_eq!(result.segments.len(), 0);
+        assert_eq!(result.total_tokens, 0);
+    }
+
+    #[test]
+    fn test_calibrate_with_segments() {
+        let calibrator = PromptCalibrator::new();
+        let budget = ContextBudget::preset_default(test_execution_id());
+        let mut window = ContextWindow::new(budget).unwrap();
+
+        window
+            .add_segment(ContextSegment::system("You are a helpful assistant.", 10))
+            .unwrap();
+        window
+            .add_segment(ContextSegment::user_input("Hello!", 5, 1))
+            .unwrap();
+
+        let config = CalibrationConfig::default();
+        let result = calibrator.calibrate(&window, &config);
+
+        assert_eq!(result.segments.len(), 2);
+        assert!(result.total_tokens > 0);
+        assert!(result.has_system());
+    }
+
+    #[test]
+    fn test_calibrate_respects_priority() {
+        let calibrator = PromptCalibrator::new();
+        let budget = ContextBudget::preset_default(test_execution_id());
+        let mut window = ContextWindow::new(budget).unwrap();
+
+        window
+            .add_segment(ContextSegment::system("System prompt", 10))
+            .unwrap();
+        window
+            .add_segment(
+                ContextSegment::new(
+                    ContextSegmentType::History,
+                    "Low priority history".to_string(),
+                    20,
+                    1,
+                )
+                .with_priority(ContextPriority::Low),
+            )
+            .unwrap();
+
+        // With high priority threshold, should exclude low priority
+        let config = CalibrationConfig {
+            min_priority: ContextPriority::High,
+            ..Default::default()
+        };
+        let result = calibrator.calibrate(&window, &config);
+
+        // Only system (critical priority) should be included
+        assert_eq!(result.segments.len(), 1);
+        assert!(result.has_system());
+        assert!(!result.has_history());
+    }
+
+    #[test]
+    fn test_calibrate_for_child() {
+        let calibrator = PromptCalibrator::new();
+        let parent_budget = ContextBudget::preset_default(test_execution_id());
+        let mut parent_window = ContextWindow::new(parent_budget).unwrap();
+
+        parent_window
+            .add_segment(ContextSegment::system("Parent system prompt", 15))
+            .unwrap();
+        parent_window
+            .add_segment(ContextSegment::user_input("Parent user input", 10, 1))
+            .unwrap();
+
+        let child_id = ExecutionId::new();
+        let config = CalibrationConfig::default();
+        let result =
+            calibrator.calibrate_for_child(&parent_window, child_id, "Analyze data", &config);
+
+        // Should have task description and parent context
+        assert!(result.total_tokens > 0);
+        assert!(result
+            .segments
+            .iter()
+            .any(|s| s.content.contains("sub-task")));
+    }
+
+    #[test]
+    fn test_calibrated_prompt_as_text() {
+        let calibrator = PromptCalibrator::new();
+        let budget = ContextBudget::preset_default(test_execution_id());
+        let mut window = ContextWindow::new(budget).unwrap();
+
+        window
+            .add_segment(ContextSegment::system("System", 5))
+            .unwrap();
+        window
+            .add_segment(ContextSegment::user_input("User", 5, 1))
+            .unwrap();
+
+        let config = CalibrationConfig::default();
+        let result = calibrator.calibrate(&window, &config);
+
+        let text = result.as_text();
+        assert!(text.contains("System"));
+        assert!(text.contains("User"));
+    }
+}