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.
//!
//! # Usage
//!
//! ```rust,ignore
//! use zclaw_runtime::growth::GrowthIntegration;
//! use zclaw_growth::{VikingAdapter, MemoryExtractor, MemoryRetriever, PromptInjector};
//!
//! // Create growth integration
//! let viking = Arc::new(VikingAdapter::in_memory());
//! let growth = GrowthIntegration::new(viking);
//!
//! // Before conversation: enhance system prompt
//! let enhanced_prompt = growth.enhance_prompt(&agent_id, &base_prompt, &user_input).await?;
//!
//! // After conversation: extract and store memories
//! growth.process_conversation(&agent_id, &messages, session_id).await?;
//! ```

use std::sync::Arc;
use zclaw_growth::{
    GrowthTracker, InjectionFormat, LlmDriverForExtraction,
    MemoryExtractor, MemoryRetriever, PromptInjector, RetrievalResult,
    VikingAdapter,
};
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)
    }

    /// 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);
    }
}