use crate::action::{Action, ActionResult};
use crate::config::AgentConfig;
use crate::goal::{Goal, GoalManager};
use crate::learning::{ActionId, LearningConfig, LearningEngine, StateId};
use crate::observation::{Observation, ObservationBuffer};
use crate::policy::{Policy, PolicyEngine, Rule};
use crate::types::Timestamp;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct AgentId(pub String);
impl AgentId {
pub fn new(name: &str) -> Self {
Self(format!("{}_{}", name, Timestamp::now().0))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum AgentState {
#[default]
Initializing,
Idle,
Processing,
Executing,
Learning,
Paused,
Stopped,
Error,
}
pub trait Agent {
fn id(&self) -> &AgentId;
fn name(&self) -> &str;
fn state(&self) -> AgentState;
fn observe(&mut self, observation: Observation);
fn decide(&self) -> Action;
fn execute(&mut self, action: Action) -> ActionResult;
fn learn(&mut self, observation: &Observation, action: &Action, result: &ActionResult);
fn step(&mut self) -> Option<ActionResult> {
let action = self.decide();
if action.is_noop() {
return None;
}
Some(self.execute(action))
}
fn config(&self) -> &AgentConfig;
}
pub struct SimpleAgent {
id: AgentId,
config: AgentConfig,
state: AgentState,
policy_engine: PolicyEngine,
goals: GoalManager,
observations: ObservationBuffer,
last_observation: Option<Observation>,
action_history: Vec<(Action, ActionResult)>,
stats: AgentStats,
learning_engine: Option<LearningEngine>,
last_state_action: Option<(StateId, ActionId)>,
}
impl SimpleAgent {
pub fn new(name: &str) -> Self {
let config = AgentConfig::new(name);
Self::with_config(name, config)
}
pub fn with_config(name: &str, config: AgentConfig) -> Self {
let max_goals = config.max_goals;
let learning_enabled = config.learning_enabled;
let mut config = config;
config.name = name.to_string();
let learning_engine = if learning_enabled {
Some(LearningEngine::default_config())
} else {
None
};
Self {
id: AgentId::new(name),
config,
state: AgentState::Idle,
policy_engine: PolicyEngine::new(),
goals: GoalManager::new(max_goals),
observations: ObservationBuffer::new(100),
last_observation: None,
action_history: Vec::new(),
stats: AgentStats::default(),
learning_engine,
last_state_action: None,
}
}
pub fn add_policy(&mut self, policy: Policy) {
self.policy_engine.add_policy(policy);
}
pub fn add_rule(&mut self, rule: Rule) {
let mut policy = Policy::new("default");
policy.add_rule(rule);
self.policy_engine.add_policy(policy);
}
pub fn set_exploration_rate(&mut self, rate: f32) {
self.policy_engine.set_exploration_rate(rate);
}
pub fn add_goal(&mut self, goal: Goal) -> Option<String> {
self.goals.add(goal)
}
pub fn set_goal(&mut self, goal: Goal) {
self.goals.add(goal);
}
pub fn active_goals(&self) -> Vec<&Goal> {
self.goals.active_goals()
}
pub fn recent_observations(&self, count: usize) -> Vec<&Observation> {
self.observations.get_recent(count)
}
pub fn stats(&self) -> &AgentStats {
&self.stats
}
pub fn pause(&mut self) {
self.state = AgentState::Paused;
}
pub fn resume(&mut self) {
if self.state == AgentState::Paused {
self.state = AgentState::Idle;
}
}
pub fn stop(&mut self) {
self.state = AgentState::Stopped;
}
pub fn is_running(&self) -> bool {
!matches!(self.state, AgentState::Stopped | AgentState::Error)
}
pub fn enable_learning(&mut self, config: LearningConfig) {
self.config.learning_enabled = true;
self.learning_engine = Some(LearningEngine::new(config));
}
pub fn enable_learning_default(&mut self) {
self.config.learning_enabled = true;
self.learning_engine = Some(LearningEngine::default_config());
}
pub fn disable_learning(&mut self) {
self.config.learning_enabled = false;
self.learning_engine = None;
}
pub fn learning_engine(&self) -> Option<&LearningEngine> {
self.learning_engine.as_ref()
}
pub fn learning_engine_mut(&mut self) -> Option<&mut LearningEngine> {
self.learning_engine.as_mut()
}
fn get_available_actions(&self, obs: &Observation) -> Vec<ActionId> {
let policy_action = self.policy_engine.decide(obs);
vec![
ActionId::from_action(&policy_action),
ActionId::from_action(&Action::noop()),
ActionId::from_action(&Action::wait()),
]
}
}
impl Agent for SimpleAgent {
fn id(&self) -> &AgentId {
&self.id
}
fn name(&self) -> &str {
&self.config.name
}
fn state(&self) -> AgentState {
self.state
}
fn observe(&mut self, observation: Observation) {
self.state = AgentState::Processing;
self.observations.push(observation.clone());
self.last_observation = Some(observation);
self.stats.observations_received += 1;
}
fn decide(&self) -> Action {
if let Some(ref obs) = self.last_observation {
self.policy_engine.decide(obs)
} else {
Action::noop()
}
}
fn execute(&mut self, action: Action) -> ActionResult {
self.state = AgentState::Executing;
self.stats.actions_executed += 1;
if self.config.learning_enabled && self.last_observation.is_some() {
let state_id = StateId::from_observation(self.last_observation.as_ref().unwrap());
let action_id = ActionId::from_action(&action);
self.last_state_action = Some((state_id, action_id));
}
log::debug!("Executing action: {:?}", action.action_type);
let result = ActionResult::success(&action.id);
if self.action_history.len() >= 100 {
self.action_history.remove(0);
}
self.action_history.push((action, result.clone()));
self.state = AgentState::Idle;
result
}
fn learn(&mut self, observation: &Observation, action: &Action, result: &ActionResult) {
if !self.config.learning_enabled {
return;
}
self.state = AgentState::Learning;
if result.success {
self.stats.successful_actions += 1;
}
let next_state = StateId::from_observation(observation);
let available_actions = self.get_available_actions(observation);
let reward = if result.success { 1.0 } else { -1.0 };
let next_action = ActionId::from_action(action);
if let Some(ref mut engine) = self.learning_engine {
if let Some((prev_state, prev_action)) = &self.last_state_action {
engine.update(
prev_state,
prev_action,
reward,
&next_state,
Some(&next_action),
&available_actions,
);
self.stats.learning_updates += 1;
}
}
self.state = AgentState::Idle;
}
fn config(&self) -> &AgentConfig {
&self.config
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentStats {
pub observations_received: u64,
pub actions_executed: u64,
pub successful_actions: u64,
pub goals_achieved: u64,
pub goals_failed: u64,
pub learning_updates: u64,
}
impl AgentStats {
pub fn success_rate(&self) -> f64 {
if self.actions_executed == 0 {
0.0
} else {
self.successful_actions as f64 / self.actions_executed as f64
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::policy::Condition;
#[test]
fn test_agent_creation() {
let agent = SimpleAgent::new("test_agent");
assert_eq!(agent.name(), "test_agent");
assert_eq!(agent.state(), AgentState::Idle);
}
#[test]
fn test_agent_observe() {
let mut agent = SimpleAgent::new("test");
let obs = Observation::sensor("temp", 25.0);
agent.observe(obs);
assert_eq!(agent.stats().observations_received, 1);
}
#[test]
fn test_agent_with_rule() {
let mut agent = SimpleAgent::new("temp_monitor");
agent.set_exploration_rate(0.0);
let rule = Rule::new(
"high_temp",
Condition::above("temperature", 30.0),
Action::alert("Temperature too high!"),
);
agent.add_rule(rule);
let obs = Observation::sensor("temperature", 35.0);
agent.observe(obs);
let action = agent.decide();
assert!(matches!(
action.action_type,
crate::action::ActionType::Alert(_)
));
}
#[test]
fn test_agent_execute() {
let mut agent = SimpleAgent::new("test");
let action = Action::store("key", "value");
let result = agent.execute(action);
assert!(result.success);
assert_eq!(agent.stats().actions_executed, 1);
}
#[test]
fn test_agent_lifecycle() {
let mut agent = SimpleAgent::new("test");
assert!(agent.is_running());
agent.pause();
assert_eq!(agent.state(), AgentState::Paused);
agent.resume();
assert_eq!(agent.state(), AgentState::Idle);
agent.stop();
assert!(!agent.is_running());
}
#[test]
fn test_agent_learning() {
use crate::learning::{LearningAlgorithm, LearningConfig};
let mut agent = SimpleAgent::new("learning_agent");
let config = LearningConfig {
learning_rate: 0.1,
discount_factor: 0.9,
algorithm: LearningAlgorithm::QLearning,
..Default::default()
};
agent.enable_learning(config);
assert!(agent.learning_engine().is_some());
let obs1 = Observation::sensor("temperature", 25.0);
agent.observe(obs1.clone());
let action1 = agent.decide();
let result1 = agent.execute(action1.clone());
agent.learn(&obs1, &action1, &result1);
assert_eq!(agent.stats().learning_updates, 1);
let engine = agent.learning_engine().unwrap();
assert!(engine.total_updates() > 0);
}
#[test]
fn test_agent_learning_toggle() {
let mut agent = SimpleAgent::with_config("test", crate::config::AgentConfig::iot_mode());
assert!(agent.learning_engine().is_none());
agent.enable_learning_default();
assert!(agent.learning_engine().is_some());
agent.disable_learning();
assert!(agent.learning_engine().is_none());
}
#[test]
fn test_agent_multi_step_learning() {
let mut agent = SimpleAgent::new("multi_step_agent");
for i in 0..10 {
let obs = Observation::sensor("step", i as f64);
agent.observe(obs.clone());
let action = agent.decide();
let result = agent.execute(action.clone());
agent.learn(&obs, &action, &result);
}
assert_eq!(agent.stats().learning_updates, 10);
assert_eq!(agent.stats().actions_executed, 10);
let engine = agent.learning_engine().unwrap();
assert!(engine.total_updates() >= 10);
assert!(engine.state_action_count() > 0);
}
}