zclaw_openfang/crates/zclaw-runtime/src/growth.rs

//! Growth System Integration for ZCLAW Runtime
//!
//! This module provides integration between the AgentLoop and the Growth System,
//! enabling automatic memory retrieval before conversations and memory extraction
//! after conversations.
//!
//! **Note (2026-03-30)**: GrowthIntegration IS wired into the Kernel's middleware
//! chain (MemoryMiddleware + CompactionMiddleware). In the Tauri desktop deployment,
//! `kernel_commands::kernel_init()` bridges the persistent SqliteStorage to the Kernel
//! via `set_viking()` + `set_extraction_driver()`, so the middleware chain and the
//! Tauri intelligence_hooks share the same persistent storage backend.

use std::sync::Arc;
use zclaw_growth::{
    GrowthTracker, InjectionFormat, LlmDriverForExtraction,
    MemoryExtractor, MemoryRetriever, PromptInjector, RetrievalResult,
    VikingAdapter,
};
use zclaw_memory::{ExtractedFactBatch, Fact, FactCategory};
use zclaw_types::{AgentId, Message, Result, SessionId};

/// Growth system integration for AgentLoop
///
/// This struct wraps the growth system components and provides
/// a simplified interface for integration with the agent loop.
pub struct GrowthIntegration {
    /// Memory retriever for fetching relevant memories
    retriever: MemoryRetriever,
    /// Memory extractor for extracting memories from conversations
    extractor: MemoryExtractor,
    /// Prompt injector for injecting memories into prompts
    injector: PromptInjector,
    /// Growth tracker for tracking growth metrics
    tracker: GrowthTracker,
    /// Configuration
    config: GrowthConfigInner,
}

/// Internal configuration for growth integration
#[derive(Debug, Clone)]
struct GrowthConfigInner {
    /// Enable/disable growth system
    pub enabled: bool,
    /// Auto-extract after each conversation
    pub auto_extract: bool,
}

impl Default for GrowthConfigInner {
    fn default() -> Self {
        Self {
            enabled: true,
            auto_extract: true,
        }
    }
}

impl GrowthIntegration {
    /// Create a new growth integration with in-memory storage
    pub fn in_memory() -> Self {
        let viking = Arc::new(VikingAdapter::in_memory());
        Self::new(viking)
    }

    /// Create a new growth integration with the given Viking adapter
    pub fn new(viking: Arc<VikingAdapter>) -> Self {
        // Create extractor without LLM driver - can be set later
        let extractor = MemoryExtractor::new_without_driver()
            .with_viking(viking.clone());

        let retriever = MemoryRetriever::new(viking.clone());
        let injector = PromptInjector::new();
        let tracker = GrowthTracker::new(viking);

        Self {
            retriever,
            extractor,
            injector,
            tracker,
            config: GrowthConfigInner::default(),
        }
    }

    /// Set the injection format
    pub fn with_format(mut self, format: InjectionFormat) -> Self {
        self.injector = self.injector.with_format(format);
        self
    }

    /// Set the LLM driver for memory extraction
    pub fn with_llm_driver(mut self, driver: Arc<dyn LlmDriverForExtraction>) -> Self {
        self.extractor = self.extractor.with_llm_driver(driver);
        self
    }

    /// Enable or disable growth system
    pub fn set_enabled(&mut self, enabled: bool) {
        self.config.enabled = enabled;
    }

    /// Check if growth system is enabled
    pub fn is_enabled(&self) -> bool {
        self.config.enabled
    }

    /// Enable or disable auto extraction
    pub fn set_auto_extract(&mut self, auto_extract: bool) {
        self.config.auto_extract = auto_extract;
    }

    /// Enhance system prompt with retrieved memories
    ///
    /// This method:
    /// 1. Retrieves relevant memories based on user input
    /// 2. Injects them into the system prompt using configured format
    ///
    /// Returns the enhanced prompt or the original if growth is disabled
    pub async fn enhance_prompt(
        &self,
        agent_id: &AgentId,
        base_prompt: &str,
        user_input: &str,
    ) -> Result<String> {
        if !self.config.enabled {
            return Ok(base_prompt.to_string());
        }

        tracing::debug!(
            "[GrowthIntegration] Enhancing prompt for agent: {}",
            agent_id
        );

        // Retrieve relevant memories
        let memories = self
            .retriever
            .retrieve(agent_id, user_input)
            .await
            .unwrap_or_else(|e| {
                tracing::warn!("[GrowthIntegration] Retrieval failed: {}", e);
                RetrievalResult::default()
            });

        if memories.is_empty() {
            tracing::debug!("[GrowthIntegration] No memories retrieved");
            return Ok(base_prompt.to_string());
        }

        tracing::info!(
            "[GrowthIntegration] Injecting {} memories ({} tokens)",
            memories.total_count(),
            memories.total_tokens
        );

        // Inject memories into prompt
        let enhanced = self.injector.inject_with_format(base_prompt, &memories);

        Ok(enhanced)
    }

    /// Process conversation after completion
    ///
    /// This method:
    /// 1. Extracts memories from the conversation using LLM (if driver available)
    /// 2. Stores the extracted memories
    /// 3. Updates growth metrics
    ///
    /// Returns the number of memories extracted
    pub async fn process_conversation(
        &self,
        agent_id: &AgentId,
        messages: &[Message],
        session_id: SessionId,
    ) -> Result<usize> {
        if !self.config.enabled || !self.config.auto_extract {
            return Ok(0);
        }

        tracing::debug!(
            "[GrowthIntegration] Processing conversation for agent: {}",
            agent_id
        );

        // Extract memories from conversation
        let extracted = self
            .extractor
            .extract(messages, session_id.clone())
            .await
            .unwrap_or_else(|e| {
                tracing::warn!("[GrowthIntegration] Extraction failed: {}", e);
                Vec::new()
            });

        if extracted.is_empty() {
            tracing::debug!("[GrowthIntegration] No memories extracted");
            return Ok(0);
        }

        tracing::info!(
            "[GrowthIntegration] Extracted {} memories",
            extracted.len()
        );

        // Store extracted memories
        let count = extracted.len();
        self.extractor
            .store_memories(&agent_id.to_string(), &extracted)
            .await?;

        // Track learning event
        self.tracker
            .record_learning(agent_id, &session_id.to_string(), count)
            .await?;

        Ok(count)
    }

    /// Combined extraction: single LLM call that produces both stored memories
    /// and structured facts, avoiding double extraction overhead.
    ///
    /// Returns `(memory_count, Option<ExtractedFactBatch>)` on success.
    pub async fn extract_combined(
        &self,
        agent_id: &AgentId,
        messages: &[Message],
        session_id: &SessionId,
    ) -> Result<Option<(usize, ExtractedFactBatch)>> {
        if !self.config.enabled || !self.config.auto_extract {
            return Ok(None);
        }

        // Single LLM extraction call
        let extracted = self
            .extractor
            .extract(messages, session_id.clone())
            .await
            .unwrap_or_else(|e| {
                tracing::warn!("[GrowthIntegration] Combined extraction failed: {}", e);
                Vec::new()
            });

        if extracted.is_empty() {
            return Ok(None);
        }

        let mem_count = extracted.len();

        // Store raw memories
        self.extractor
            .store_memories(&agent_id.to_string(), &extracted)
            .await?;

        // Track learning event
        self.tracker
            .record_learning(agent_id, &session_id.to_string(), mem_count)
            .await?;

        // Convert same extracted memories to structured facts (no extra LLM call)
        let facts: Vec<Fact> = extracted
            .into_iter()
            .map(|m| {
                let category = match m.memory_type {
                    zclaw_growth::types::MemoryType::Preference => FactCategory::Preference,
                    zclaw_growth::types::MemoryType::Knowledge => FactCategory::Knowledge,
                    zclaw_growth::types::MemoryType::Experience => FactCategory::Behavior,
                    _ => FactCategory::General,
                };
                Fact::new(m.content, category, f64::from(m.confidence))
                    .with_source(session_id.to_string())
            })
            .collect();

        let batch = ExtractedFactBatch {
            facts,
            agent_id: agent_id.to_string(),
            session_id: session_id.to_string(),
        }
        .deduplicate()
        .filter_by_confidence(0.7);

        if batch.is_empty() {
            return Ok(Some((mem_count, ExtractedFactBatch {
                facts: vec![],
                agent_id: agent_id.to_string(),
                session_id: session_id.to_string(),
            })));
        }

        Ok(Some((mem_count, batch)))
    }

    /// Retrieve memories for a query without injection
    pub async fn retrieve_memories(
        &self,
        agent_id: &AgentId,
        query: &str,
    ) -> Result<RetrievalResult> {
        self.retriever.retrieve(agent_id, query).await
    }

    /// Get growth statistics for an agent
    pub async fn get_stats(&self, agent_id: &AgentId) -> Result<zclaw_growth::GrowthStats> {
        self.tracker.get_stats(agent_id).await
    }

    /// Warm up cache with hot memories
    pub async fn warmup_cache(&self, agent_id: &AgentId) -> Result<usize> {
        self.retriever.warmup_cache(agent_id).await
    }

    /// Clear the semantic index
    pub async fn clear_index(&self) {
        self.retriever.clear_index().await;
    }
}

impl Default for GrowthIntegration {
    fn default() -> Self {
        Self::in_memory()
    }
}

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

    #[tokio::test]
    async fn test_growth_integration_creation() {
        let growth = GrowthIntegration::in_memory();
        assert!(growth.is_enabled());
    }

    #[tokio::test]
    async fn test_enhance_prompt_empty() {
        let growth = GrowthIntegration::in_memory();
        let agent_id = AgentId::new();
        let base = "You are helpful.";
        let user_input = "Hello";

        let enhanced = growth
            .enhance_prompt(&agent_id, base, user_input)
            .await
            .unwrap();

        // Without any stored memories, should return base prompt
        assert_eq!(enhanced, base);
    }

    #[tokio::test]
    async fn test_disabled_growth() {
        let mut growth = GrowthIntegration::in_memory();
        growth.set_enabled(false);

        let agent_id = AgentId::new();
        let base = "You are helpful.";

        let enhanced = growth
            .enhance_prompt(&agent_id, base, "test")
            .await
            .unwrap();

        assert_eq!(enhanced, base);
    }

    #[tokio::test]
    async fn test_process_conversation_disabled() {
        let mut growth = GrowthIntegration::in_memory();
        growth.set_auto_extract(false);

        let agent_id = AgentId::new();
        let messages = vec![Message::user("Hello")];
        let session_id = SessionId::new();

        let count = growth
            .process_conversation(&agent_id, &messages, session_id)
            .await
            .unwrap();

        assert_eq!(count, 0);
    }
}