cp_validator/

corpus.rs

//! Test query corpus for validator evaluations.
//!
//! Per CP-015 section 4: Test queries are drawn from a canonical query
//! corpus maintained by the network. The corpus is a set of query strings
//! with known-relevant chunk IDs, stored on Arweave and updated through
//! network consensus.

use cp_arweave::ArweaveClient;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use tracing::info;
use uuid::Uuid;

/// A single test query with ground-truth relevance information.
///
/// Per CP-015 section 4: each test query includes the query text,
/// a pre-computed embedding, the set of known-relevant chunk IDs,
/// and graded relevance scores.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestQuery {
    /// The query text.
    pub query_text: String,

    /// Pre-computed query embedding (i16 quantized, variable dimension).
    /// This is the raw i16 vector as stored in Arweave, matching the
    /// network's canonical embedding model.
    pub query_embedding: Vec<i16>,

    /// Ground-truth relevant chunk IDs for this query.
    pub relevant_chunk_ids: Vec<Uuid>,

    /// Graded relevance scores corresponding to `relevant_chunk_ids`.
    /// Values 0-3 where 0 = not relevant, 1 = marginally relevant,
    /// 2 = relevant, 3 = highly relevant.
    /// Must be the same length as `relevant_chunk_ids`.
    pub relevance_grades: Vec<u8>,
}

impl TestQuery {
    /// Convert the i16 query embedding to the format expected by
    /// `GraphStore::search` (f32).
    pub fn query_embedding_f32(&self) -> Vec<f32> {
        self.query_embedding
            .iter()
            .map(|&v| f32::from(v) / 32767.0)
            .collect()
    }

    /// Return the i16 query embedding directly.
    pub fn query_embedding_i16(&self) -> Vec<i16> {
        self.query_embedding.clone()
    }

    /// Build a map from chunk ID to relevance grade for NDCG computation.
    pub fn relevance_grade_map(&self) -> HashMap<Uuid, u8> {
        self.relevant_chunk_ids
            .iter()
            .zip(self.relevance_grades.iter())
            .map(|(id, &grade)| (*id, grade))
            .collect()
    }
}

/// A collection of test queries forming the canonical evaluation corpus.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestCorpus {
    /// The test queries in this corpus.
    pub queries: Vec<TestQuery>,
}

impl TestCorpus {
    /// Create an empty test corpus.
    pub fn new() -> Self {
        Self {
            queries: Vec::new(),
        }
    }

    /// Number of queries in the corpus.
    pub fn len(&self) -> usize {
        self.queries.len()
    }

    /// Whether the corpus is empty.
    pub fn is_empty(&self) -> bool {
        self.queries.is_empty()
    }

    /// Load a test corpus from a file on disk.
    ///
    /// Tries CBOR first, then JSON as fallback.
    pub fn load_from_file(path: &Path) -> Result<Self, crate::ValidatorError> {
        let bytes = std::fs::read(path).map_err(|e| {
            crate::ValidatorError::Corpus(format!(
                "Failed to read corpus file {}: {}",
                path.display(),
                e
            ))
        })?;

        // Try CBOR first, then JSON as fallback
        if let Ok(corpus) = ciborium::from_reader::<TestCorpus, _>(bytes.as_slice()) {
            info!(queries = corpus.queries.len(), path = %path.display(), "Loaded corpus from CBOR");
            return Ok(corpus);
        }

        let corpus: TestCorpus = serde_json::from_slice(&bytes).map_err(|e| {
            crate::ValidatorError::Corpus(format!(
                "Failed to parse corpus file {} (tried CBOR and JSON): {}",
                path.display(),
                e
            ))
        })?;

        info!(queries = corpus.queries.len(), path = %path.display(), "Loaded corpus from JSON");
        Ok(corpus)
    }

    /// Save the corpus to a CBOR file on disk.
    pub fn save_to_file(&self, path: &Path) -> Result<(), crate::ValidatorError> {
        let mut buf = Vec::new();
        ciborium::into_writer(self, &mut buf).map_err(|e| {
            crate::ValidatorError::Corpus(format!("Failed to serialize corpus: {e}"))
        })?;

        std::fs::write(path, &buf).map_err(|e| {
            crate::ValidatorError::Corpus(format!(
                "Failed to write corpus file {}: {}",
                path.display(),
                e
            ))
        })?;

        info!(queries = self.queries.len(), path = %path.display(), "Saved corpus to CBOR");
        Ok(())
    }

    /// Load the canonical test corpus from Arweave.
    ///
    /// Queries Arweave for the latest corpus transaction tagged with
    /// Content-Type: application/x-canon-test-corpus, downloads and
    /// deserializes it.
    pub async fn load_from_arweave(arweave: &ArweaveClient) -> Result<Self, crate::ValidatorError> {
        info!("Loading test corpus from Arweave");

        // Query for corpus transactions
        let tags = vec![cp_arweave::TagFilter::new(
            "Content-Type",
            &["application/x-canon-test-corpus"],
        )];

        let result = arweave
            .query_transactions(&tags, 1, None, cp_arweave::SortOrder::HeightDesc)
            .await
            .map_err(|e| crate::ValidatorError::Arweave(format!("Failed to query corpus: {e}")))?;

        if result.edges.is_empty() {
            return Err(crate::ValidatorError::Corpus(
                "No test corpus found on Arweave".to_string(),
            ));
        }

        let tx_id = &result.edges[0].node.id;
        info!(tx_id = tx_id, "Found corpus transaction");

        let data = arweave.get_transaction_data(tx_id).await.map_err(|e| {
            crate::ValidatorError::Arweave(format!("Failed to download corpus: {e}"))
        })?;

        // Decompress if needed (try zstd first, then raw CBOR)
        let corpus_bytes = match zstd::decode_all(data.as_slice()) {
            Ok(decompressed) => decompressed,
            Err(_) => data,
        };

        let corpus: TestCorpus = ciborium::from_reader(corpus_bytes.as_slice()).map_err(|e| {
            crate::ValidatorError::Corpus(format!("Failed to parse corpus from Arweave: {e}"))
        })?;

        info!(queries = corpus.queries.len(), "Loaded corpus from Arweave");
        Ok(corpus)
    }
}

impl Default for TestCorpus {
    fn default() -> Self {
        Self::new()
    }
}

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

    fn sample_query(i: usize) -> TestQuery {
        TestQuery {
            query_text: format!("What is topic {i}?"),
            query_embedding: vec![i as i16; 10],
            relevant_chunk_ids: vec![
                Uuid::from_bytes([i as u8; 16]),
                Uuid::from_bytes([(i + 1) as u8; 16]),
            ],
            relevance_grades: vec![3, 1],
        }
    }

    #[test]
    fn test_test_query_relevance_grade_map() {
        let query = sample_query(0);
        let map = query.relevance_grade_map();
        assert_eq!(map.len(), 2);
        assert_eq!(*map.get(&Uuid::from_bytes([0u8; 16])).unwrap(), 3);
        assert_eq!(*map.get(&Uuid::from_bytes([1u8; 16])).unwrap(), 1);
    }

    #[test]
    fn test_test_query_embedding_i16() {
        let query = sample_query(5);
        let emb = query.query_embedding_i16();
        assert_eq!(emb, vec![5i16; 10]);
    }

    #[test]
    fn test_test_query_embedding_f32() {
        let query = sample_query(5);
        let emb = query.query_embedding_f32();
        assert_eq!(emb.len(), 10);
        // 5 / 32767.0
        assert!((emb[0] - 5.0 / 32767.0).abs() < 1e-6);
    }

    #[test]
    fn test_corpus_new_empty() {
        let corpus = TestCorpus::new();
        assert!(corpus.is_empty());
        assert_eq!(corpus.len(), 0);
    }

    #[test]
    fn test_corpus_with_queries() {
        let corpus = TestCorpus {
            queries: (0..10).map(sample_query).collect(),
        };
        assert!(!corpus.is_empty());
        assert_eq!(corpus.len(), 10);
    }

    #[test]
    fn test_corpus_save_load_cbor_roundtrip() {
        let corpus = TestCorpus {
            queries: (0..5).map(sample_query).collect(),
        };

        let tmp = tempfile::NamedTempFile::new().unwrap();
        let path = tmp.path().to_path_buf();

        corpus.save_to_file(&path).unwrap();
        let loaded = TestCorpus::load_from_file(&path).unwrap();

        assert_eq!(loaded.len(), corpus.len());
        for (a, b) in loaded.queries.iter().zip(corpus.queries.iter()) {
            assert_eq!(a.query_text, b.query_text);
            assert_eq!(a.query_embedding, b.query_embedding);
            assert_eq!(a.relevant_chunk_ids, b.relevant_chunk_ids);
            assert_eq!(a.relevance_grades, b.relevance_grades);
        }
    }

    #[test]
    fn test_corpus_load_json() {
        let corpus = TestCorpus {
            queries: vec![sample_query(0)],
        };

        let json = serde_json::to_vec(&corpus).unwrap();
        let tmp = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(tmp.path(), &json).unwrap();

        let loaded = TestCorpus::load_from_file(tmp.path()).unwrap();
        assert_eq!(loaded.len(), 1);
        assert_eq!(loaded.queries[0].query_text, "What is topic 0?");
    }

    #[test]
    fn test_corpus_load_nonexistent() {
        let result = TestCorpus::load_from_file(Path::new("/nonexistent/corpus.cbor"));
        assert!(result.is_err());
    }

    #[test]
    fn test_corpus_cbor_serialization() {
        let corpus = TestCorpus {
            queries: (0..3).map(sample_query).collect(),
        };

        let mut buf = Vec::new();
        ciborium::into_writer(&corpus, &mut buf).unwrap();
        let decoded: TestCorpus = ciborium::from_reader(buf.as_slice()).unwrap();

        assert_eq!(decoded.len(), 3);
    }

    #[test]
    fn test_test_query_serialization() {
        let query = sample_query(42);
        let json = serde_json::to_string(&query).unwrap();
        let decoded: TestQuery = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.query_text, query.query_text);
        assert_eq!(decoded.relevant_chunk_ids, query.relevant_chunk_ids);
    }
}