sail_scripts/kernel/src/token_budget.rs

//! Dynamic Token Budget Calculator
//!
//! Computes `max_new_tokens` dynamically based on question complexity.
//! This is the kernel-level implementation of the "smart token budget" feature.
//!
//! Strategy:
//!   1. If explicit marks are embedded in the question (e.g., "2 marks", "8 marks")
//!      → use a calibrated lookup table
//!   2. Otherwise, estimate complexity from:
//!      - Prompt length (characters)
//!      - Keyword signals ("explain", "list", "compare", "analyze", "define")
//!      - Question type (open/closed, single/multi-part)
//!
//! All pattern matching uses bounded regexes to prevent ReDoS.

use pyo3::prelude::*;
use regex::Regex;
use once_cell::sync::Lazy;

// ── Token budget lookup table ─────────────────────────────────────────────────

/// (min_marks, max_marks) → max_new_tokens
const MARK_BUDGET_TABLE: &[(u32, u32, u32)] = &[
    (1,  2,   150),  // Brief definition / fill-in-the-blank
    (3,  4,   350),  // Short paragraph answer
    (5,  6,   600),  // Medium explanation
    (7,  8,   900),  // Detailed section answer
    (9,  10, 1200),  // Extended analysis
    (11, 15, 1800),  // Essay / multi-part
    (16, 99, 2500),  // Comprehensive long-form
];

/// Default budget when no marks/complexity signals found
const DEFAULT_BUDGET: u32 = 512;

// ── Patterns ──────────────────────────────────────────────────────────────────

/// Detects "X marks", "X mark", "(X)", "[X marks]" in prompt text
static MARKS_PATTERN: Lazy<Regex> = Lazy::new(|| {
    Regex::new(r"(?i)\b(\d{1,2})\s*(?:marks?|pts?|points?)\b|\[(\d{1,2})\s*marks?\]|\((\d{1,2})\s*marks?\)")
        .expect("MARKS_PATTERN regex is valid")
});

/// "Explain", "Analyze", "Compare", "Evaluate" → longer answer
static DEEP_KEYWORDS: Lazy<Regex> = Lazy::new(|| {
    Regex::new(r"(?i)\b(explain|analyze|analyse|compare|contrast|evaluate|discuss|elaborate|describe in detail|critically|justify|argue|examine)\b")
        .expect("DEEP_KEYWORDS regex valid")
});

/// "Define", "List", "Name", "State", "Give" → shorter answer  
static BRIEF_KEYWORDS: Lazy<Regex> = Lazy::new(|| {
    Regex::new(r"(?i)\b(define|list|name|state|give|identify|what is|what are|who is|when did|where is)\b")
        .expect("BRIEF_KEYWORDS regex valid")
});

/// "Step by step", "With examples", "Show working" → extra tokens needed
static EXPANSION_KEYWORDS: Lazy<Regex> = Lazy::new(|| {
    Regex::new(r"(?i)\b(step[\s-]by[\s-]step|with examples?|show (?:your )?work(?:ing)?|in detail|thoroughly|comprehensively)\b")
        .expect("EXPANSION_KEYWORDS regex valid")
});

// ── Core logic ────────────────────────────────────────────────────────────────

fn compute_budget_internal(prompt: &str, explicit_marks: Option<u32>) -> u32 {
    // Priority 1: Explicit marks provided by caller
    if let Some(marks) = explicit_marks {
        return marks_to_tokens(marks);
    }

    // Priority 2: Detect marks from prompt text
    if let Some(detected) = detect_marks_in_prompt(prompt) {
        return marks_to_tokens(detected);
    }

    // Priority 3: Keyword-based heuristic
    let p = prompt;
    let prompt_len = p.len();

    let deep_signals = DEEP_KEYWORDS.find_iter(p).count();
    let brief_signals = BRIEF_KEYWORDS.find_iter(p).count();
    let expansion_signals = EXPANSION_KEYWORDS.find_iter(p).count();

    // Question complexity score (0 = minimal, 100+ = essay)
    let mut score: i32 = 0;

    // Prompt length contributes to complexity
    score += (prompt_len / 50).min(20) as i32; // Max +20 from length

    // Deep keywords strongly suggest longer response
    score += (deep_signals as i32) * 15;

    // Brief keywords suggest concise response
    score -= (brief_signals as i32) * 10;

    // Expansion modifiers bump the budget
    score += (expansion_signals as i32) * 20;

    // Map score to token budget
    let budget = match score {
        s if s <= 0  => 150,
        1..=10       => 250,
        11..=20      => 400,
        21..=35      => 600,
        36..=50      => 900,
        51..=70      => 1200,
        _            => 1800,
    };

    budget
}

/// Look up max_new_tokens from a mark value using the budget table.
fn marks_to_tokens(marks: u32) -> u32 {
    for &(min_m, max_m, tokens) in MARK_BUDGET_TABLE {
        if marks >= min_m && marks <= max_m {
            return tokens;
        }
    }
    DEFAULT_BUDGET
}

/// Try to extract a mark value from the prompt text.
fn detect_marks_in_prompt(prompt: &str) -> Option<u32> {
    // Use only the first match (first marks reference in prompt)
    let cap = MARKS_PATTERN.captures(prompt)?;

    // Try capture groups in order: group 1, 2, 3
    for i in 1..=3 {
        if let Some(m) = cap.get(i) {
            if let Ok(n) = m.as_str().parse::<u32>() {
                return Some(n);
            }
        }
    }
    None
}

// ── Python-exposed function ───────────────────────────────────────────────────

/// Compute the appropriate max_new_tokens for a given prompt.
///
/// Args:
///     prompt (str): The user question/prompt text
///     explicit_marks (int | None): Explicit mark value if known (overrides detection)
///
/// Returns:
///     int: Recommended max_new_tokens value
///
/// Examples:
///     >>> compute_budget("What is photosynthesis? [2 marks]")
///     150
///     >>> compute_budget("Explain with examples how TCP/IP works.", None)
///     600
///     >>> compute_budget("Any question", explicit_marks=10)
///     1200
#[pyfunction]
#[pyo3(signature = (prompt, explicit_marks = None))]
pub fn compute_budget(
    py: Python<'_>,
    prompt: String,
    explicit_marks: Option<u32>,
) -> u32 {
    py.allow_threads(move || compute_budget_internal(&prompt, explicit_marks))
}

// ── Unit Tests ────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_explicit_marks_override() {
        assert_eq!(compute_budget_internal("anything", Some(2)), 150);
        assert_eq!(compute_budget_internal("anything", Some(8)), 900);
        assert_eq!(compute_budget_internal("anything", Some(15)), 1800);
    }

    #[test]
    fn test_marks_detection_in_prompt() {
        assert_eq!(compute_budget_internal("Define photosynthesis. [2 marks]", None), 150);
        assert_eq!(compute_budget_internal("Evaluate the French Revolution (8 marks)", None), 900);
        assert_eq!(compute_budget_internal("Explain 5 marks worth of TCP/IP", None), 600);
    }

    #[test]
    fn test_keyword_heuristics() {
        // Deep keywords should increase budget
        let long_q = compute_budget_internal("Analyze and compare the causes of WW1 and WW2 in detail step-by-step", None);
        let short_q = compute_budget_internal("Define photosynthesis", None);
        assert!(long_q > short_q, "Deep question should get more tokens");
    }

    #[test]
    fn test_empty_prompt() {
        let result = compute_budget_internal("", None);
        assert!(result > 0, "Should return a valid budget even for empty prompts");
    }

    #[test]
    fn test_marks_to_tokens_table() {
        assert_eq!(marks_to_tokens(1), 150);
        assert_eq!(marks_to_tokens(2), 150);
        assert_eq!(marks_to_tokens(3), 350);
        assert_eq!(marks_to_tokens(10), 1200);
        assert_eq!(marks_to_tokens(20), 2500);
    }
}