enact_context/

calibrator.rs

//! 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
#[allow(dead_code)]
static SEGMENT_SEQUENCE: AtomicU64 = AtomicU64::new(1000);

#[allow(dead_code)]
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"));
    }
}