roboticus_agent/

normalization.rs

//! # Normalization
//!
//! Detects malformed tool protocol and narrated next steps in ReAct loop output,
//! and builds retry prompts to guide the model back to correct behavior.

use regex::Regex;
use serde::Serialize;
use std::sync::LazyLock;

/// Maximum number of normalization retries before giving up.
pub const MAX_NORMALIZATION_RETRIES: u8 = 2;

/// Describes the kind of protocol failure detected in model output.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum NormalizationPattern {
    /// Tool call JSON is syntactically malformed.
    MalformedToolCall,
    /// Model narrated what it would do instead of calling a tool.
    NarratedToolUse,
    /// Model produced empty or whitespace-only action.
    EmptyAction,
}

// ── Compiled regex patterns ──────────────────────────────────────────────────

/// Matches an `Action:` or `tool_call` keyword followed by broken JSON.
/// "Broken" means there is at least one `{` but unbalanced braces, or a
/// bare `{` followed by content that is missing a closing `}`.
static RE_ACTION_OR_TOOL_CALL: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?i)(Action\s*:|tool_call)").unwrap());

/// Matches narration patterns where the model describes an intent to use a
/// tool rather than actually using it, only when no `Action:` block is present.
static RE_NARRATED: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(r"(?i)(I would use|I'll run|let me use the|I should call|I need to invoke)\s+\w+")
        .unwrap()
});

/// Matches `Action:` followed by only optional whitespace (empty action).
static RE_EMPTY_ACTION: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?i)Action\s*:\s*$").unwrap());

// ── Public API ────────────────────────────────────────────────────────────────

/// Inspect `content` and return the first `NormalizationPattern` detected, or
/// `None` if the content looks well-formed.
///
/// Detection priority:
/// 1. `EmptyAction`  — content is blank, or `Action:` with nothing after it.
/// 2. `MalformedToolCall` — `Action:` / `tool_call` present but JSON is broken.
/// 3. `NarratedToolUse` — model described intent without an `Action:` block.
pub fn detect_normalization_failure(content: &str) -> Option<NormalizationPattern> {
    // 1. Empty / whitespace-only content
    if content.trim().is_empty() {
        return Some(NormalizationPattern::EmptyAction);
    }

    // 2. `Action:` present but nothing follows it (trailing whitespace only)
    if RE_EMPTY_ACTION.is_match(content) {
        return Some(NormalizationPattern::EmptyAction);
    }

    // 3. `Action:` or `tool_call` present → check whether JSON is well-formed
    if RE_ACTION_OR_TOOL_CALL.is_match(content) {
        if has_broken_json(content) {
            return Some(NormalizationPattern::MalformedToolCall);
        }
        // Keyword present and JSON looks balanced — not a failure
        return None;
    }

    // 4. No `Action:` block at all, but the model narrated tool use
    if RE_NARRATED.is_match(content) {
        return Some(NormalizationPattern::NarratedToolUse);
    }

    None
}

/// Build a system-level retry instruction for the given `pattern`.
///
/// `tool_count` is embedded in the `NarratedToolUse` message so the model
/// knows how many tools are available.
pub fn build_normalization_retry_prompt(
    pattern: &NormalizationPattern,
    tool_count: usize,
) -> String {
    match pattern {
        NormalizationPattern::MalformedToolCall => {
            "Your previous response contained a malformed tool call. \
             Please retry using the correct JSON format:\n\
             Action: tool_name\n\
             Action Input: {\"param\": \"value\"}"
                .to_string()
        }
        NormalizationPattern::NarratedToolUse => {
            format!(
                "You described what you would do instead of doing it. \
                 Use the Action/Action Input format to actually invoke the tool. \
                 You have {tool_count} tools available."
            )
        }
        NormalizationPattern::EmptyAction => "Your previous response was empty. \
             Please provide either a direct answer or use a tool via \
             Action/Action Input format."
            .to_string(),
    }
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Returns `true` when the content contains a `{` character but the curly
/// braces are unbalanced, indicating broken JSON.  Also returns `true` when
/// there are no braces at all after an `Action:` / `tool_call` keyword (the
/// caller has already confirmed such a keyword is present).
fn has_broken_json(content: &str) -> bool {
    let open: i32 = content.chars().filter(|&c| c == '{').count() as i32;
    let close: i32 = content.chars().filter(|&c| c == '}').count() as i32;

    if open == 0 {
        // Keyword present but no JSON object at all — malformed
        return true;
    }

    // Unbalanced braces or a quoted field left open (simple heuristic)
    open != close
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    // 1. Detects malformed JSON tool call (unbalanced braces)
    #[test]
    fn detects_malformed_tool_call_unbalanced_braces() {
        let content = "Action: web_search\nAction Input: {\"query\": \"rust async\"";
        assert_eq!(
            detect_normalization_failure(content),
            Some(NormalizationPattern::MalformedToolCall)
        );
    }

    // 2. Detects malformed tool call — keyword present but no JSON at all
    #[test]
    fn detects_malformed_tool_call_no_json() {
        let content = "Action: web_search\nAction Input: query rust async";
        assert_eq!(
            detect_normalization_failure(content),
            Some(NormalizationPattern::MalformedToolCall)
        );
    }

    // 3. Detects narrated tool use ("I would use web_search to…")
    #[test]
    fn detects_narrated_tool_use() {
        let content = "I would use web_search to find recent articles on the topic.";
        assert_eq!(
            detect_normalization_failure(content),
            Some(NormalizationPattern::NarratedToolUse)
        );
    }

    // 4. Detects empty action — content is whitespace only
    #[test]
    fn detects_empty_action_whitespace_only() {
        let content = "   \n\t  ";
        assert_eq!(
            detect_normalization_failure(content),
            Some(NormalizationPattern::EmptyAction)
        );
    }

    // 5. Normal tool call NOT detected as failure
    #[test]
    fn normal_tool_call_not_detected_as_failure() {
        let content = "Action: web_search\nAction Input: {\"query\": \"rust async\"}";
        assert_eq!(detect_normalization_failure(content), None);
    }

    // 6. Normal text response (no tool) NOT detected as failure
    #[test]
    fn normal_text_response_not_detected_as_failure() {
        let content = "The answer is 42. Rust is a systems programming language.";
        assert_eq!(detect_normalization_failure(content), None);
    }

    // 7. Retry prompt includes tool count for NarratedToolUse
    #[test]
    fn retry_prompt_includes_tool_count() {
        let prompt = build_normalization_retry_prompt(&NormalizationPattern::NarratedToolUse, 7);
        assert!(prompt.contains("7 tools available"));
    }

    // 8. Detects empty Action: line (trailing whitespace after colon)
    #[test]
    fn detects_empty_action_line() {
        let content = "Thought: I should search for this.\nAction:   ";
        assert_eq!(
            detect_normalization_failure(content),
            Some(NormalizationPattern::EmptyAction)
        );
    }
}