hirag_integration_test.rs

//! Integration test for the complete HiRAG pipeline.
//!
//! This test verifies that the `LayerConsolidator` and `HiRag` orchestrator work
//! together correctly with a real `SurrealDbAdapter` backend.
//!
//! Test Flow:
//! 1. Setup: Initialize `SurrealDbAdapter` (in-memory) and mock LLM/Embedding ports.
//! 2. Ingest: Store 5 base-layer (L0) `MemoryNode`s.
//! 3. Consolidate: Run `LayerConsolidator` to generate a summary (L1) node.
//! 4. Verify Consolidation: Check that the L1 node was created and linked.
//! 5. Query: Execute a query using the `HiRag` orchestrator.
//! 6. Verify HiRAG: Assert that the pipeline retrieves the summary and synthesizes the correct answer.

use std::sync::Arc;
use async_trait::async_trait;
use synapse_core::{
    error::Result,
    MemoryNode,
    MemoryPort,
    LlmPort,
    EmbeddingPort,
    LayerConsolidator,
    HiRag,
};
use synapse_infra::adapters::{SurrealDbAdapter, MockEmbeddingAdapter};
use synapse_core::HolographicPort;

// === Test Setup ===

struct MockHolographic;

#[async_trait]
impl HolographicPort for MockHolographic {
    async fn compress_context(&self, _text: &str) -> Result<Vec<u8>> {
        Ok(Vec::new())
    }

    async fn decompress_context(&self, _neurogram: &[u8]) -> Result<String> {
        Ok("".to_string())
    }
}

/// A mock LLM that provides predictable responses for query decomposition and synthesis.
struct ConfigurableMockLlm {
    decomposition_response: String,
    synthesis_response: String,
}

impl Default for ConfigurableMockLlm {
    fn default() -> Self {
        Self {
            decomposition_response: "What is the summary of the test data?".to_string(),
            synthesis_response: "The final answer is based on the summary of the test data.".to_string(),
        }
    }
}

#[async_trait]
impl LlmPort for ConfigurableMockLlm {
    async fn generate(&self, prompt: &str, _max_tokens: usize) -> Result<String> {
        if prompt.contains("Decompose") {
            Ok(self.decomposition_response.clone())
        } else if prompt.contains("Summarize") {
            Ok("This is a consolidated summary of the test data.".to_string())
        } else {
            Ok(self.synthesis_response.clone())
        }
    }

     async fn generate_with_params(&self, _prompt: &str, _max_tokens: usize, _temp: f32, _top_p: f32) -> Result<String> {
        unimplemented!()
    }

    async fn generate_stream(
        &self,
        _prompt: &str,
        _max_tokens: usize,
    ) -> Result<std::pin::Pin<Box<dyn futures::Stream<Item = Result<String>> + Send>>> {
        let stream = futures::stream::once(async { Ok("Mock summary".to_string()) });
        Ok(Box::pin(stream))
    }
}

#[tokio::test]
async fn test_full_hirag_pipeline() -> Result<()> {
    // 1. Setup
    let memory_adapter: Arc<dyn MemoryPort> = Arc::new(SurrealDbAdapter::new_memory().await?);
    let llm_adapter = Arc::new(ConfigurableMockLlm::default());
    let embedding_adapter = Arc::new(MockEmbeddingAdapter::new());

    // 2. Ingest Layer 0 nodes
    let facts = vec![
        "The sky is blue.",
        "Grass is green.",
        "The sun is bright.",
        "Water is wet.",
        "Fire is hot.",
    ];

    for fact in facts {
        let embedding = embedding_adapter.embed(fact).await?;
        let node = MemoryNode::new(fact.to_string()).with_embedding(embedding);
        memory_adapter.store(node).await?;
    }

    assert_eq!(memory_adapter.count_by_layer(0).await?, 5);

    // 3. Consolidate
    let consolidator = LayerConsolidator::new(
        memory_adapter.clone(),
        llm_adapter.clone(),
        embedding_adapter.clone(),
    ).with_threshold(5); // Consolidate when 5 nodes are present

    let summary_id = consolidator.consolidate_layer(0).await?.unwrap();
    assert!(!summary_id.is_empty());

    // 4. Verify Consolidation
    assert_eq!(memory_adapter.count_by_layer(1).await?, 1, "A layer 1 summary node should be created");
    let summary_node = memory_adapter.get_by_id(&summary_id).await?.unwrap();
    assert_eq!(summary_node.content, "This is a consolidated summary of the test data.");

    // TODO: Verify relationships once SurrealDbAdapter supports graph queries.
    // For now, we trust the implementation based on its unit tests.

    // 5. Query using HiRAG
    let holographic_adapter = Arc::new(MockHolographic);
    let hirag = HiRag::new(
        memory_adapter.clone(),
        llm_adapter.clone(),
        embedding_adapter.clone(),
        holographic_adapter,
    );

    let query = "Give me a high-level overview of the test data.";
    let result = hirag.execute_query(query).await?;

    // 6. Verify HiRAG
    assert_eq!(result, "The final answer is based on the summary of the test data.");

    Ok(())
}