compliance-scanner-agent/compliance-agent/src/llm/triage.rs

use std::sync::Arc;

use compliance_core::models::{Finding, FindingStatus};

use crate::llm::LlmClient;
use crate::pipeline::orchestrator::GraphContext;

const TRIAGE_SYSTEM_PROMPT: &str = r#"You are a security finding triage expert. Analyze the following security finding and determine:
1. Is this a true positive? (yes/no)
2. Confidence score (0-10, where 10 is highest confidence this is a real issue)
3. Brief remediation suggestion (1-2 sentences)

Respond in JSON format:
{"true_positive": true/false, "confidence": N, "remediation": "..."}"#;

pub async fn triage_findings(
    llm: &Arc<LlmClient>,
    findings: &mut Vec<Finding>,
    graph_context: Option<&GraphContext>,
) -> usize {
    let mut passed = 0;

    for finding in findings.iter_mut() {
        let mut user_prompt = format!(
            "Scanner: {}\nRule: {}\nSeverity: {}\nTitle: {}\nDescription: {}\nFile: {}\nLine: {}\nCode: {}",
            finding.scanner,
            finding.rule_id.as_deref().unwrap_or("N/A"),
            finding.severity,
            finding.title,
            finding.description,
            finding.file_path.as_deref().unwrap_or("N/A"),
            finding.line_number.map(|n| n.to_string()).unwrap_or_else(|| "N/A".to_string()),
            finding.code_snippet.as_deref().unwrap_or("N/A"),
        );

        // Enrich with graph context if available
        if let Some(ctx) = graph_context {
            if let Some(impact) = ctx
                .impacts
                .iter()
                .find(|i| i.finding_id == finding.fingerprint)
            {
                user_prompt.push_str(&format!(
                    "\n\n--- Code Graph Context ---\n\
                     Blast radius: {} nodes affected\n\
                     Entry points affected: {}\n\
                     Direct callers: {}\n\
                     Communities affected: {}\n\
                     Call chains: {}",
                    impact.blast_radius,
                    if impact.affected_entry_points.is_empty() {
                        "none".to_string()
                    } else {
                        impact.affected_entry_points.join(", ")
                    },
                    if impact.direct_callers.is_empty() {
                        "none".to_string()
                    } else {
                        impact.direct_callers.join(", ")
                    },
                    impact.affected_communities.len(),
                    impact.call_chains.len(),
                ));
            }
        }

        match llm
            .chat(TRIAGE_SYSTEM_PROMPT, &user_prompt, Some(0.1))
            .await
        {
            Ok(response) => {
                // Strip markdown code fences if present (e.g. ```json ... ```)
                let cleaned = response.trim();
                let cleaned = if cleaned.starts_with("```") {
                    let inner = cleaned
                        .trim_start_matches("```json")
                        .trim_start_matches("```")
                        .trim_end_matches("```")
                        .trim();
                    inner
                } else {
                    cleaned
                };
                if let Ok(result) = serde_json::from_str::<TriageResult>(cleaned) {
                    finding.confidence = Some(result.confidence);
                    if let Some(remediation) = result.remediation {
                        finding.remediation = Some(remediation);
                    }

                    if result.confidence >= 3.0 {
                        finding.status = FindingStatus::Triaged;
                        passed += 1;
                    } else {
                        finding.status = FindingStatus::FalsePositive;
                    }
                } else {
                    // If LLM response doesn't parse, keep the finding
                    finding.status = FindingStatus::Triaged;
                    passed += 1;
                    tracing::warn!(
                        "Failed to parse triage response for {}: {response}",
                        finding.fingerprint
                    );
                }
            }
            Err(e) => {
                // On LLM error, keep the finding
                tracing::warn!("LLM triage failed for {}: {e}", finding.fingerprint);
                finding.status = FindingStatus::Triaged;
                passed += 1;
            }
        }
    }

    // Remove false positives
    findings.retain(|f| f.status != FindingStatus::FalsePositive);
    passed
}

#[derive(serde::Deserialize)]
struct TriageResult {
    #[serde(default)]
    #[allow(dead_code)]
    true_positive: bool,
    #[serde(default)]
    confidence: f64,
    remediation: Option<String>,
}