neuromance-common 0.0.5

//! Agent state management and context types.
//!
//! This module provides types for managing agent state, memory, and context in autonomous
//! LLM-powered agents. These types enable agents to maintain conversation history, track
//! execution statistics, store short/long-term memory, and manage task context.
//!
//! ## Core Types
//!
//! - [`AgentState`]: Complete agent state including messages, memory, context, and stats
//! - [`AgentMemory`]: Memory storage with short-term, long-term, and working memory
//! - [`AgentContext`]: Task context including goals, constraints, and environment
//! - [`AgentMessage`]: Enum for different message types agents can process
//! - [`AgentResponse`]: Response from agent execution including content and tool results
//! - [`AgentStats`]: Execution statistics for monitoring and debugging
//!
//! ## Example: Using Agent State
//!
//! ```
//! use neuromance_common::agents::{AgentState, AgentContext, AgentMemory, AgentStats};
//!
//! // Create a new agent state with default values
//! let mut state = AgentState::default();
//!
//! // Configure the agent's task context
//! state.context.task = Some("Research the Rust programming language".to_string());
//! state.context.goals = vec![
//!     "Find creation date".to_string(),
//!     "Identify key features".to_string(),
//! ];
//! state.context.constraints = vec!["Use only reliable sources".to_string()];
//!
//! // Add to short-term memory
//! state.memory.short_term.push("User prefers concise answers".to_string());
//!
//! // Track statistics
//! state.stats.total_messages += 1;
//! state.stats.successful_tool_calls += 2;
//! ```
//!
//! ## Memory Management
//!
//! [`AgentMemory`] provides three types of memory storage:
//!
//! - **Short-term**: Recent context and observations (Vec of strings)
//! - **Long-term**: Persistent knowledge (key-value string storage)
//! - **Working memory**: Arbitrary JSON data for active processing
//!
//! ## Context Updates
//!
//! The [`ContextUpdate`] enum allows dynamic modification of agent context:
//!
//! ```
//! use neuromance_common::agents::ContextUpdate;
//!
//! // Examples of context updates
//! let updates = vec![
//!     ContextUpdate::SetTask("New task description".to_string()),
//!     ContextUpdate::AddGoal("Complete research".to_string()),
//!     ContextUpdate::AddConstraint("Time limit: 5 minutes".to_string()),
//!     ContextUpdate::SetEnvironmentVariable("debug".to_string(), "true".to_string()),
//! ];
//! ```

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::chat::Message;

/// Updates that can be applied to an agent's context.
///
/// Allows dynamic modification of task, goals, constraints, and environment during execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ContextUpdate {
    /// Set or replace the current task description.
    SetTask(String),
    /// Add a new goal to the agent's objectives.
    AddGoal(String),
    /// Remove a goal from the agent's objectives.
    RemoveGoal(String),
    /// Add a constraint that limits agent behavior.
    AddConstraint(String),
    /// Remove a previously added constraint.
    RemoveConstraint(String),
    /// Set an environment variable (key-value pair).
    SetEnvironmentVariable(String, String),
    /// Clear all memory (short-term, long-term, and working).
    ClearMemory,
}

/// Messages that can be sent to or from an agent.
///
/// Represents different types of communication in an agent system.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AgentMessage {
    /// Input from a user.
    UserInput(String),
    /// System-level prompt or instruction.
    SystemPrompt(String),
    /// Result from tool execution.
    ToolResult(String),
    /// Response generated by the agent.
    AgentResponse(Box<AgentResponse>),
    /// Request to update agent context.
    ContextUpdate(ContextUpdate),
}

/// Response from agent execution.
///
/// Contains the main content, optional reasoning trace, and any tool execution results.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentResponse {
    /// The primary response message from the agent.
    pub content: Message,
    /// Optional chain-of-thought or reasoning explanation.
    pub reasoning: Option<String>,
    /// Messages from tool executions that contributed to this response.
    pub tool_responses: Vec<Message>,
}

/// Complete state of an agent.
///
/// Maintains all messages, conversation history, memory, context, and execution statistics.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentState {
    /// All messages in the current conversation.
    pub messages: Vec<Message>,
    /// History of agent interactions (message-response pairs).
    pub conversation_history: Vec<(AgentMessage, AgentResponse)>,
    /// Agent's memory systems.
    pub memory: AgentMemory,
    /// Current task context and objectives.
    pub context: AgentContext,
    /// Execution statistics and metrics.
    pub stats: AgentStats,
}

/// Execution statistics for monitoring agent performance.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentStats {
    /// Total number of messages processed.
    pub total_messages: usize,
    /// Total tokens consumed across all interactions.
    pub tokens_used: usize,
    /// Number of successful tool executions.
    pub successful_tool_calls: usize,
    /// Number of failed tool executions.
    pub failed_tool_calls: usize,
}

/// Context describing the agent's current task and environment.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentContext {
    /// The current task or objective.
    pub task: Option<String>,
    /// List of goals the agent should achieve.
    pub goals: Vec<String>,
    /// Constraints that limit agent behavior.
    pub constraints: Vec<String>,
    /// Environment variables available to the agent.
    pub environment: HashMap<String, String>,
}

/// Memory storage for agents.
///
/// Provides three types of memory for different use cases:
/// - Short-term: Recent observations and context
/// - Long-term: Persistent facts and knowledge
/// - Working: Arbitrary data for active processing
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentMemory {
    /// Recent observations and temporary context.
    pub short_term: Vec<String>,
    /// Persistent knowledge stored as key-value pairs.
    pub long_term: HashMap<String, String>,
    /// Arbitrary JSON data for active processing.
    pub working_memory: HashMap<String, serde_json::Value>,
}