anya_core/ml/agents/

mod.rs

// Machine Learning Agents Module
//
// This module provides a modular system of ML agents that enhance decision-making
// across Anya's core functions, including Stacks blockchain operations, DAO governance,
// and Web5 capabilities. All agents adhere to the core principle of "read first always"
// to ensure informed decision-making, Bitcoin principles of decentralization, and ethical AI.

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::sync::{Arc, RwLock};

// System maps for global state tracking and indexing
pub mod system_map;
pub use system_map::*;

// Re-export agents
pub mod federated_agent;
pub use federated_agent::FederatedAgent;

pub mod dao_agent;
pub use dao_agent::DaoAgent;

// pub mod web5_agent;
// pub use web5_agent::Web5Agent;

// pub mod stacks_agent;
// pub use stacks_agent::StacksAgent;

/// Unique identifier for an agent
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct AgentId(pub String);

impl fmt::Display for AgentId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Observation provided to an agent
#[derive(Debug, Clone)]
pub enum Observation {
    /// Text-based observation
    Text(String),

    /// Numeric observation with a specific metric name
    Numeric(String, f64),

    /// JSON-structured observation
    Json(serde_json::Value),

    /// Custom observation with binary data
    Custom(String, Vec<u8>),

    /// System state observation containing a snapshot of the system's state
    SystemState(SystemStateRef),
}

/// Reference to system state for cloneable observations
#[derive(Debug, Clone)]
pub struct SystemStateRef {
    /// Timestamp of the observation
    pub timestamp: u64,
    /// Reference identifier for the system state
    pub state_id: String,
}

/// Action taken by an agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Action {
    /// Recommendation for a human operator
    Recommendation(String),

    /// Notification of an event or insight
    Notification(String, String), // (title, message)

    /// JSON-structured response
    Json(serde_json::Value),

    /// Custom action with binary data
    Custom(String, Vec<u8>),

    /// System update action with payload
    SystemUpdate(SystemUpdateType, serde_json::Value),
}

/// Types of system updates
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SystemUpdateType {
    /// Update to the global index
    IndexUpdate,

    /// Update to the system map
    MapUpdate,

    /// Updates to configuration
    ConfigUpdate,

    /// Update to agent states
    AgentStateUpdate,
}

/// System state observation
#[derive(Debug)]
pub struct SystemState {
    /// Current system index snapshot
    pub index: Option<SystemIndex>,
    /// Current system map snapshot  
    pub map: Option<SystemMap>,
    /// Timestamp of the observation
    pub timestamp: u64,
}

/// Feedback provided to an agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Feedback {
    /// Numerical score for the feedback (0.0 to 1.0)
    pub score: f32,

    /// Textual description of the feedback
    pub description: Option<String>,

    /// Source of the feedback (human, another agent, system)
    pub source: FeedbackSource,
}

/// Source of feedback
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FeedbackSource {
    /// Feedback from a human
    Human,

    /// Feedback from another agent
    Agent(AgentId),

    /// Feedback from the system
    System,
}

/// Metrics for agent performance
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct AgentMetrics {
    /// Number of observations processed
    pub observations_processed: u64,

    /// Number of actions taken
    pub actions_taken: u64,

    /// Average feedback score
    pub average_feedback: f32,

    /// Processing time in milliseconds (average)
    pub avg_processing_time_ms: f64,

    /// Custom metrics
    pub custom_metrics: HashMap<String, f64>,
}

/// Error from agent operations
#[derive(Debug, thiserror::Error)]
pub enum AgentError {
    /// Invalid observation format
    #[error("Invalid observation: {0}")]
    InvalidObservation(String),

    /// Error during processing
    #[error("Processing error: {0}")]
    ProcessingError(String),

    /// Input/output error
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),

    /// Serialization/deserialization error
    #[error("Serialization error: {0}")]
    SerializationError(#[from] serde_json::Error),

    /// Internal error
    #[error("Internal error: {0}")]
    InternalError(String),

    /// Ethical compliance error
    #[error("Ethical compliance error: {0}")]
    EthicalComplianceError(String),

    /// System time error
    #[error("System time error: {0}")]
    SystemTimeError(#[from] std::time::SystemTimeError),

    /// Version parsing error
    #[error("Version error: {0}")]
    VersionError(#[from] semver::Error),
}

/// The core agent trait that all ML agents must implement
#[async_trait]
pub trait Agent: Send + Sync {
    /// Get the agent's unique identifier
    fn id(&self) -> &AgentId;

    /// Get the agent's type
    fn agent_type(&self) -> &str;

    /// Process an observation and optionally return an action
    async fn process(&self, observation: Observation) -> Result<Option<Action>, AgentError>;

    /// Receive feedback on a previous action
    async fn receive_feedback(&mut self, feedback: Feedback) -> Result<(), AgentError>;

    /// Get the agent's performance metrics
    fn metrics(&self) -> AgentMetrics;

    /// Get the agent's ethical compliance score (0.0 to 1.0)
    fn ethical_compliance(&self) -> f32 {
        0.8 // Default reasonable compliance score
    }

    /// Read system state before processing (implements "read first always")
    async fn read_system_state(&self) -> Result<SystemState, AgentError> {
        use crate::ml::agents::system_map::{system_index, system_map};

        // Check that we can read the index and map
        system_index().read_index().await?;
        system_map().read_map().await?;

        // Return a simple state indicating success
        Ok(SystemState {
            index: None, // Can't return actual index due to atomic types
            map: None,   // Can't return actual map due to non-cloneable types
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
        })
    }

    /// Update system state after processing
    async fn update_system_state(&self, updates: &[SystemUpdateType]) -> Result<(), AgentError> {
        use crate::ml::agents::system_map::{system_index, system_map};

        for update_type in updates {
            match update_type {
                SystemUpdateType::IndexUpdate => {
                    system_index().increment_version().await?;
                }
                SystemUpdateType::MapUpdate => {
                    system_map().update_map().await?;
                }
                _ => {} // Other updates handled elsewhere
            }
        }
        Ok(())
    }
}

/// Orchestrates multiple agents working together
pub struct AgentSystem {
    /// Registered agents by their ID
    agents: RwLock<HashMap<AgentId, Arc<dyn Agent>>>,

    /// System configuration
    config: RwLock<AgentSystemConfig>,

    /// System metrics
    metrics: RwLock<AgentSystemMetrics>,
}

/// Configuration for the agent system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentSystemConfig {
    /// Whether to enforce the "read first always" principle
    pub enforce_read_first: bool,

    /// Minimum ethical compliance score required for agents
    pub min_ethical_compliance: f32,

    /// Maximum number of agents that can be registered
    pub max_agents: usize,

    /// Default timeout for agent processing in milliseconds
    pub default_timeout_ms: u64,
}

impl Default for AgentSystemConfig {
    fn default() -> Self {
        Self {
            enforce_read_first: true, // Always enforce read-first by default
            min_ethical_compliance: 0.7,
            max_agents: 100,
            default_timeout_ms: 5000,
        }
    }
}

/// Metrics for the agent system
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct AgentSystemMetrics {
    /// Total number of observations processed
    pub total_observations: u64,

    /// Total number of actions taken
    pub total_actions: u64,

    /// Total number of errors
    pub total_errors: u64,

    /// Average processing time in milliseconds
    pub avg_processing_time_ms: f64,

    /// System uptime in seconds
    pub uptime_seconds: u64,
}

impl AgentSystem {
    /// Create a new agent system with default configuration
    pub fn new() -> Self {
        Self {
            agents: RwLock::new(HashMap::new()),
            config: RwLock::new(AgentSystemConfig::default()),
            metrics: RwLock::new(AgentSystemMetrics::default()),
        }
    }

    /// Create a new agent system with custom configuration
    pub fn with_config(config: AgentSystemConfig) -> Self {
        Self {
            agents: RwLock::new(HashMap::new()),
            config: RwLock::new(config),
            metrics: RwLock::new(AgentSystemMetrics::default()),
        }
    }

    /// Register a new agent with the system
    pub async fn register_agent(&self, agent: Arc<dyn Agent>) -> Result<(), AgentError> {
        let agent_id = agent.id().clone();
        let config = self.config.read().map_err(|_| {
            AgentError::InternalError("Failed to acquire read lock on config".to_string())
        })?;

        // Check ethical compliance
        let compliance = agent.ethical_compliance();
        if compliance < config.min_ethical_compliance {
            return Err(AgentError::EthicalComplianceError(format!(
                "Agent {} has insufficient ethical compliance score: {} < {}",
                agent_id, compliance, config.min_ethical_compliance
            )));
        }

        // Register the agent
        let mut agents = self.agents.write().map_err(|_| {
            AgentError::InternalError("Failed to acquire write lock on agents".to_string())
        })?;

        if agents.len() >= config.max_agents {
            return Err(AgentError::ProcessingError(format!(
                "Maximum number of agents ({}) reached",
                config.max_agents
            )));
        }

        agents.insert(agent_id, agent);

        Ok(())
    }

    /// Unregister an agent from the system
    pub async fn unregister_agent(&self, agent_id: &AgentId) -> Result<(), AgentError> {
        let mut agents = self.agents.write().map_err(|_| {
            AgentError::InternalError("Failed to acquire write lock on agents".to_string())
        })?;

        if agents.remove(agent_id).is_none() {
            return Err(AgentError::ProcessingError(format!(
                "Agent {agent_id} not found"
            )));
        }

        Ok(())
    }

    /// Process an observation with a specific agent
    pub async fn process_with_agent(
        &self,
        agent_id: &AgentId,
        observation: Observation,
    ) -> Result<Option<Action>, AgentError> {
        // Get agent and config values without holding the locks through await points
        let (agent, enforce_read_first) = {
            // Get agent reference
            let agents = self.agents.read().map_err(|_| {
                AgentError::InternalError("Failed to acquire read lock on agents".to_string())
            })?;

            let agent = agents
                .get(agent_id)
                .ok_or_else(|| AgentError::ProcessingError(format!("Agent {agent_id} not found")))?
                .clone();

            // Get config value
            let config = self.config.read().map_err(|_| {
                AgentError::InternalError("Failed to acquire read lock on config".to_string())
            })?;

            (agent, config.enforce_read_first)
        };

        // Enforce read-first principle if configured
        if enforce_read_first {
            // First read the current system state
            let _system_state = agent.read_system_state().await?;

            // Then process with the original observation plus system state
            let combined_observation = match observation {
                Observation::SystemState(_) => observation,
                _ => Observation::SystemState(SystemStateRef {
                    timestamp: chrono::Utc::now().timestamp() as u64,
                    state_id: format!("state_{}", chrono::Utc::now().timestamp()),
                }),
            };

            let start_time = std::time::Instant::now();
            let result = agent.process(combined_observation).await;
            let processing_time = start_time.elapsed();

            // Update metrics
            {
                let mut metrics = self.metrics.write().map_err(|_| {
                    AgentError::InternalError("Failed to acquire write lock on metrics".to_string())
                })?;

                metrics.total_observations += 1;
                if result.is_ok() && result.as_ref().unwrap().is_some() {
                    metrics.total_actions += 1;
                }
                if result.is_err() {
                    metrics.total_errors += 1;
                }

                // Update average processing time
                let current_avg = metrics.avg_processing_time_ms;
                let current_count = metrics.total_observations;
                metrics.avg_processing_time_ms = (current_avg * (current_count - 1) as f64
                    + processing_time.as_millis() as f64)
                    / current_count as f64;
            }

            // Update system state after processing if there was an action
            if let Ok(Some(_)) = &result {
                agent
                    .update_system_state(&[
                        SystemUpdateType::IndexUpdate,
                        SystemUpdateType::MapUpdate,
                    ])
                    .await?;
            }

            result
        } else {
            // Standard processing without enforcing read-first
            agent.process(observation).await
        }
    }

    /// Broadcast an observation to all agents and collect their actions
    pub async fn broadcast(
        &self,
        observation: Observation,
    ) -> HashMap<AgentId, Result<Option<Action>, AgentError>> {
        // Get agent IDs without holding the lock through await points
        let agent_ids = {
            let agents = match self.agents.read() {
                Ok(agents) => agents,
                Err(_) => return HashMap::new(),
            };

            // Collect agent IDs to process after releasing the lock
            agents.keys().cloned().collect::<Vec<_>>()
        };

        let mut results = HashMap::new();

        for agent_id in agent_ids {
            let result = self
                .process_with_agent(&agent_id, observation.clone())
                .await;
            results.insert(agent_id, result);
        }

        results
    }

    /// Get the configuration
    pub fn config(&self) -> Result<AgentSystemConfig, AgentError> {
        self.config.read().map(|c| c.clone()).map_err(|_| {
            AgentError::InternalError("Failed to acquire read lock on config".to_string())
        })
    }

    /// Update the configuration
    pub fn update_config(&self, config: AgentSystemConfig) -> Result<(), AgentError> {
        let mut current_config = self.config.write().map_err(|_| {
            AgentError::InternalError("Failed to acquire write lock on config".to_string())
        })?;

        *current_config = config;

        Ok(())
    }

    /// Get the system metrics
    pub fn metrics(&self) -> Result<AgentSystemMetrics, AgentError> {
        self.metrics.read().map(|m| m.clone()).map_err(|_| {
            AgentError::InternalError("Failed to acquire read lock on metrics".to_string())
        })
    }
}

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

#[allow(dead_code)]
pub struct MLAgentCoordinator {
    agents: Vec<Box<dyn Agent>>,
    // Add basic placeholders for missing types
}

/// Basic resource pool placeholder
#[allow(dead_code)]
pub struct ResourcePool {
    capacity: usize,
}

/// Basic health monitor placeholder
#[allow(dead_code)]
pub struct HealthMonitor {
    is_healthy: bool,
}

#[cfg(test)]
mod tests {
    #[tokio::test]
    async fn test_agent_system_registration() {
        // Test agent registration and unregistration
    }

    #[tokio::test]
    async fn test_read_first_principle() {
        // Test that the read-first principle is enforced
    }

    #[tokio::test]
    async fn test_ethical_compliance() {
        // Test that ethical compliance is properly enforced
    }
}