hirn-engine 0.1.0

Engine for the hirn cognitive memory database
Documentation
use hirn_core::MemoryId;
use hirn_core::revision::RevisionRef;
use hirn_core::types::{Layer, Namespace};

use crate::db::HirnDB;
use crate::diagnostics::QueryDiagnostics;
use crate::ql::context::ThinkResult;
use crate::recall::RecallResult;
use crate::scoring::{ScoreBreakdown, ScoringWeights};

#[derive(Debug, Clone)]
pub struct RetrievalSuppressionSummary {
    pub candidate_count: usize,
    pub threshold_filtered_count: usize,
    pub competitive_inhibition_count: usize,
    pub truncated_by_limit_count: usize,
}

#[derive(Debug, Clone)]
pub struct RetrievedRecordExplanation {
    pub memory_id: MemoryId,
    pub layer: Layer,
    pub revision: Option<RevisionRef>,
    pub composite_score: Option<f32>,
    pub score_breakdown: Option<ScoreBreakdown>,
    pub raw_text_redacted: bool,
    pub ranking_details_redacted: bool,
    pub resource_evidence_count: usize,
}

#[derive(Debug, Clone)]
pub struct RetrievalExplanation {
    pub diagnostics: QueryDiagnostics,
    pub scoring_weights: ScoringWeights,
    pub policy: RetrievalPolicySummary,
    pub suppression: RetrievalSuppressionSummary,
    pub raw_text_redacted_results: usize,
    pub results: Vec<RetrievedRecordExplanation>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RetrievalPolicyScope {
    Unrestricted,
    RequestedNamespace,
    AllowedNamespaces,
    RequestedNamespaceDenied,
}

#[derive(Debug, Clone)]
pub struct RetrievalPolicySummary {
    pub scope: RetrievalPolicyScope,
    pub requested_namespace: Option<Namespace>,
    pub allowed_namespaces: Option<Vec<Namespace>>,
    pub namespace_restriction_applied: bool,
    pub raw_text_redaction_applied: bool,
}

#[derive(Debug, Clone)]
pub struct ThinkExplanation {
    pub retrieval: RetrievalExplanation,
    pub token_budget: usize,
    pub token_count: usize,
    pub records_included_count: usize,
    pub records_excluded_count: usize,
    pub conflict_group_count: usize,
    pub query_time_ms: f64,
}

pub(crate) fn build_retrieval_explanation(
    db: &HirnDB,
    actor_id: &str,
    results: &[RecallResult],
    diagnostics: QueryDiagnostics,
    scoring_weights: ScoringWeights,
    requested_namespace: Option<Namespace>,
    allowed_namespaces: Option<Vec<Namespace>>,
) -> RetrievalExplanation {
    let result_explanations: Vec<_> = results
        .iter()
        .map(|result| {
            let raw_text_redacted = !db.can_read_raw_content(actor_id, &result.record);
            let ranking_details_redacted = raw_text_redacted;

            RetrievedRecordExplanation {
                memory_id: result.record.id(),
                layer: result.record.layer(),
                revision: result.revision,
                composite_score: (!ranking_details_redacted).then_some(result.composite_score),
                score_breakdown: (!ranking_details_redacted).then_some(result.score_breakdown),
                raw_text_redacted,
                ranking_details_redacted,
                resource_evidence_count: result.resource_evidence.len(),
            }
        })
        .collect();

    let raw_text_redacted_results = result_explanations
        .iter()
        .filter(|result| result.raw_text_redacted)
        .count();

    let policy = build_policy_summary(
        requested_namespace,
        allowed_namespaces,
        raw_text_redacted_results,
    );

    RetrievalExplanation {
        suppression: RetrievalSuppressionSummary {
            candidate_count: diagnostics.records_scanned.unwrap_or(results.len()),
            threshold_filtered_count: diagnostics.threshold_filtered_count.unwrap_or(0),
            competitive_inhibition_count: diagnostics.competitive_inhibition_count.unwrap_or(0),
            truncated_by_limit_count: diagnostics.truncated_by_limit_count.unwrap_or(0),
        },
        diagnostics,
        scoring_weights,
        policy,
        raw_text_redacted_results,
        results: result_explanations,
    }
}

fn build_policy_summary(
    requested_namespace: Option<Namespace>,
    allowed_namespaces: Option<Vec<Namespace>>,
    raw_text_redacted_results: usize,
) -> RetrievalPolicySummary {
    let scope = match (requested_namespace, allowed_namespaces.as_ref()) {
        (Some(requested), Some(allowed)) if !allowed.contains(&requested) => {
            RetrievalPolicyScope::RequestedNamespaceDenied
        }
        (Some(_), _) => RetrievalPolicyScope::RequestedNamespace,
        (None, Some(_)) => RetrievalPolicyScope::AllowedNamespaces,
        (None, None) => RetrievalPolicyScope::Unrestricted,
    };

    RetrievalPolicySummary {
        scope,
        requested_namespace,
        allowed_namespaces,
        namespace_restriction_applied: !matches!(scope, RetrievalPolicyScope::Unrestricted),
        raw_text_redaction_applied: raw_text_redacted_results > 0,
    }
}

pub(crate) fn build_think_explanation(
    retrieval: RetrievalExplanation,
    think_result: &ThinkResult,
    token_budget: usize,
) -> ThinkExplanation {
    ThinkExplanation {
        retrieval,
        token_budget,
        token_count: think_result.token_count,
        records_included_count: think_result.records_included.len(),
        records_excluded_count: think_result.records_excluded_count,
        conflict_group_count: think_result.conflict_groups.len(),
        query_time_ms: think_result.query_time_ms,
    }
}

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

    #[test]
    fn policy_summary_reports_allowed_namespace_scope() {
        let policy = build_policy_summary(
            None,
            Some(vec![Namespace::default_ns(), Namespace::shared()]),
            0,
        );

        assert_eq!(policy.scope, RetrievalPolicyScope::AllowedNamespaces);
        assert!(policy.namespace_restriction_applied);
        assert!(!policy.raw_text_redaction_applied);
    }

    #[test]
    fn policy_summary_reports_denied_requested_namespace() {
        let requested = Namespace::new("private-a").unwrap();
        let allowed = Namespace::new("private-b").unwrap();

        let policy = build_policy_summary(Some(requested), Some(vec![allowed]), 1);

        assert_eq!(policy.scope, RetrievalPolicyScope::RequestedNamespaceDenied);
        assert!(policy.namespace_restriction_applied);
        assert!(policy.raw_text_redaction_applied);
        assert_eq!(policy.requested_namespace, Some(requested));
    }
}