cloudiful_redactor/redactor/

mod.rs

use std::sync::Arc;

use crate::{
    Finding, InputKind, LlmConfig, RedactionArtifact, RedactionPolicy, RedactionResult,
    RedactionRules, RedactionSession, RedactorError, RestoreResult, restore_patch_with_session,
    restore_text_with_session,
};

mod detection;
mod session;
mod stats;

use detection::detect_internal;
use session::SessionRedactorExt;
use stats::stats_for;

#[derive(Debug, Clone, Default)]
pub struct RedactorBuilder {
    llm: Option<LlmConfig>,
    policy: RedactionPolicy,
}

impl RedactorBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_llm(mut self, config: LlmConfig) -> Self {
        self.llm = Some(config);
        self
    }

    pub fn with_person_detection(mut self, enabled: bool) -> Self {
        self.policy.rules.person = enabled;
        self
    }

    pub fn with_redaction_rules(mut self, rules: RedactionRules) -> Self {
        self.policy.rules = rules;
        self
    }

    pub fn with_redaction_policy(mut self, policy: RedactionPolicy) -> Self {
        self.policy = policy;
        self
    }

    pub fn build(self) -> Redactor {
        let custom_strings = crate::detect::CompiledCustomStrings::new(&self.policy.custom_strings);
        Redactor {
            llm: self.llm,
            compiled_policy: Arc::new(CompiledPolicy {
                policy: self.policy,
                custom_strings,
            }),
        }
    }
}

#[derive(Debug, Clone)]
pub struct Redactor {
    pub(super) llm: Option<LlmConfig>,
    compiled_policy: Arc<CompiledPolicy>,
}

#[derive(Debug)]
struct CompiledPolicy {
    policy: RedactionPolicy,
    custom_strings: crate::detect::CompiledCustomStrings,
}

#[derive(Debug, Default)]
pub struct SessionRedactor {
    pub(super) processor: crate::replace::ReplacementProcessor,
}

impl Redactor {
    pub fn policy(&self) -> &RedactionPolicy {
        &self.compiled_policy.policy
    }

    pub fn redact(&self, text: &str) -> Result<RedactionResult, RedactorError> {
        self.redact_with_input_kind(text, InputKind::Text)
    }

    pub fn redact_with_input_kind(
        &self,
        text: &str,
        input_kind: InputKind,
    ) -> Result<RedactionResult, RedactorError> {
        let artifact = self.redact_artifact_with_input_kind(text, input_kind)?;
        Ok(artifact.result)
    }

    pub fn redact_with_source_path(
        &self,
        text: &str,
        source_path: &str,
    ) -> Result<RedactionResult, RedactorError> {
        let artifact = self.redact_artifact_with_input_kind_and_source(
            text,
            InputKind::Text,
            Some(source_path),
        )?;
        Ok(artifact.result)
    }

    pub fn redact_artifact(&self, text: &str) -> Result<RedactionArtifact, RedactorError> {
        self.redact_artifact_with_input_kind(text, InputKind::Text)
    }

    pub fn redact_artifact_with_input_kind(
        &self,
        text: &str,
        input_kind: InputKind,
    ) -> Result<RedactionArtifact, RedactorError> {
        self.redact_artifact_with_input_kind_and_source(text, input_kind, None)
    }

    pub fn redact_artifact_with_input_kind_and_source(
        &self,
        text: &str,
        input_kind: InputKind,
        source_path: Option<&str>,
    ) -> Result<RedactionArtifact, RedactorError> {
        self.redact_artifact_with_input_kind_source_and_prior_session(
            text,
            input_kind,
            source_path,
            None,
            None,
        )
    }

    pub fn redact_artifact_with_input_kind_source_and_prior_session(
        &self,
        text: &str,
        input_kind: InputKind,
        source_path: Option<&str>,
        prior_session: Option<&RedactionSession>,
        external_id: Option<&str>,
    ) -> Result<RedactionArtifact, RedactorError> {
        let outcome = detect_internal(self, text, input_kind, source_path);
        let findings = outcome.findings;
        let mut processor =
            crate::replace::ReplacementProcessor::with_prior_session(prior_session, external_id)?;
        let redacted_text = processor.redact_fragment(text, &findings);
        let session = processor.build_session(text, &redacted_text, self.policy());
        let applied_replacements = processor.into_applied_replacements();
        let stats = stats_for(self.llm.is_some(), &findings, outcome.stats);

        Ok(RedactionArtifact {
            result: RedactionResult {
                redacted_text: redacted_text.clone(),
                findings,
                applied_replacements,
                stats,
            },
            session,
        })
    }

    pub fn redact_with_session(&self, text: &str) -> Result<RedactionSession, RedactorError> {
        self.redact_with_session_input_kind(text, InputKind::Text)
    }

    pub fn redact_with_session_input_kind(
        &self,
        text: &str,
        input_kind: InputKind,
    ) -> Result<RedactionSession, RedactorError> {
        Ok(self
            .redact_artifact_with_input_kind(text, input_kind)?
            .session)
    }

    pub fn detect(&self, text: &str) -> Result<Vec<Finding>, RedactorError> {
        self.detect_with_input_kind(text, InputKind::Text)
    }

    pub fn detect_with_input_kind(
        &self,
        text: &str,
        input_kind: InputKind,
    ) -> Result<Vec<Finding>, RedactorError> {
        Ok(detect_internal(self, text, input_kind, None).findings)
    }

    pub fn detect_with_source_path(
        &self,
        text: &str,
        source_path: &str,
    ) -> Result<Vec<Finding>, RedactorError> {
        Ok(detect_internal(self, text, InputKind::Text, Some(source_path)).findings)
    }

    pub fn restore_text(&self, text: &str, session: &RedactionSession) -> RestoreResult {
        restore_text_with_session(text, session)
    }

    pub fn restore_patch(&self, patch: &str, session: &RedactionSession) -> RestoreResult {
        restore_patch_with_session(patch, session)
    }

    pub fn redact_artifact_with_prior_session(
        &self,
        text: &str,
        input_kind: InputKind,
        prior_session: Option<&RedactionSession>,
        external_id: Option<&str>,
    ) -> Result<RedactionArtifact, RedactorError> {
        self.redact_artifact_with_input_kind_source_and_prior_session(
            text,
            input_kind,
            None,
            prior_session,
            external_id,
        )
    }
}

impl SessionRedactor {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_prior_session(
        prior_session: Option<&RedactionSession>,
        external_id: Option<&str>,
    ) -> Result<Self, RedactorError> {
        Ok(Self {
            processor: crate::replace::ReplacementProcessor::with_prior_session(
                prior_session,
                external_id,
            )?,
        })
    }

    pub fn redact_fragment(
        &mut self,
        redactor: &Redactor,
        text: &str,
    ) -> Result<String, RedactorError> {
        self.redact_text_fragment(redactor, text)
    }

    pub fn redact_fragment_with_input_kind(
        &mut self,
        redactor: &Redactor,
        text: &str,
        input_kind: InputKind,
    ) -> Result<String, RedactorError> {
        self.redact_text_fragment_with_input_kind(redactor, text, input_kind)
    }

    pub fn finish_session(
        &self,
        original_text: &str,
        redacted_text: &str,
        policy: &RedactionPolicy,
    ) -> RedactionSession {
        self.build_redaction_session(original_text, redacted_text, policy)
    }

    pub fn has_applied_replacements(&self) -> bool {
        self.processor.has_applied_replacements()
    }

    pub fn build_session(&self, original_text: &str, redacted_text: &str) -> RedactionSession {
        self.build_redaction_session(
            original_text,
            redacted_text,
            &crate::RedactionPolicy::default(),
        )
    }

    pub fn max_token_len(&self) -> usize {
        self.max_replacement_token_len()
    }
}