enya-agent 0.1.0

package/crates/enact-context/src/step_context.rs

@@ -0,0 +1,645 @@
+//! Step Context Builder
+//!
+//! Extracts learnings and builds context when steps are discovered.
+//! Used during agentic execution when StepSource::Agent spawns new steps.
+//!
+//! @see packages/enact-schemas/src/context.schemas.ts
+
+use crate::segment::{ContextPriority, ContextSegment, ContextSegmentType};
+use crate::token_counter::TokenCounter;
+use chrono::{DateTime, Utc};
+use enact_core::kernel::{ExecutionId, StepId};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+/// Global sequence counter for segments
+static STEP_SEQUENCE: AtomicU64 = AtomicU64::new(2000);
+
+fn next_sequence() -> u64 {
+    STEP_SEQUENCE.fetch_add(1, Ordering::SeqCst)
+}
+
+/// Context extraction configuration
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct StepContextConfig {
+    /// Maximum tokens for step context
+    pub max_tokens: usize,
+
+    /// Include tool call results
+    pub include_tool_results: bool,
+
+    /// Include reasoning/chain-of-thought
+    pub include_reasoning: bool,
+
+    /// Include error context
+    pub include_errors: bool,
+
+    /// Maximum tool results to include
+    pub max_tool_results: usize,
+
+    /// Truncate long content
+    pub truncate_long_content: bool,
+
+    /// Maximum content length before truncation
+    pub max_content_length: usize,
+}
+
+impl Default for StepContextConfig {
+    fn default() -> Self {
+        Self {
+            max_tokens: 2000,
+            include_tool_results: true,
+            include_reasoning: true,
+            include_errors: true,
+            max_tool_results: 5,
+            truncate_long_content: true,
+            max_content_length: 1000,
+        }
+    }
+}
+
+/// Extracted learning from a step
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct StepLearning {
+    /// Unique ID for this learning
+    pub id: String,
+
+    /// Step that produced this learning
+    pub step_id: StepId,
+
+    /// Execution containing the step
+    pub execution_id: ExecutionId,
+
+    /// Type of learning
+    pub learning_type: LearningType,
+
+    /// The learning content
+    pub content: String,
+
+    /// Confidence in this learning (0.0 - 1.0)
+    pub confidence: f64,
+
+    /// Relevance to future steps (0.0 - 1.0)
+    pub relevance: f64,
+
+    /// Tags for categorization
+    pub tags: Vec<String>,
+
+    /// Timestamp
+    pub created_at: DateTime<Utc>,
+}
+
+/// Types of learnings that can be extracted
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum LearningType {
+    /// Successful action pattern
+    SuccessPattern,
+    /// Error and recovery
+    ErrorRecovery,
+    /// Tool usage insight
+    ToolInsight,
+    /// Decision rationale
+    DecisionRationale,
+    /// Domain knowledge
+    DomainKnowledge,
+    /// Constraint discovered
+    ConstraintDiscovered,
+    /// User preference
+    UserPreference,
+}
+
+/// Result of step context extraction
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct StepContextResult {
+    /// Execution ID
+    pub execution_id: ExecutionId,
+
+    /// Step ID that was processed
+    pub step_id: StepId,
+
+    /// Extracted context segments
+    pub segments: Vec<ContextSegment>,
+
+    /// Extracted learnings
+    pub learnings: Vec<StepLearning>,
+
+    /// Total tokens in extracted context
+    pub total_tokens: usize,
+
+    /// Processing timestamp
+    pub processed_at: DateTime<Utc>,
+}
+
+/// Step Context Builder - extracts learnings when steps are discovered
+pub struct StepContextBuilder {
+    token_counter: TokenCounter,
+    config: StepContextConfig,
+}
+
+impl StepContextBuilder {
+    /// Create a new builder with default config
+    pub fn new() -> Self {
+        Self {
+            token_counter: TokenCounter::default(),
+            config: StepContextConfig::default(),
+        }
+    }
+
+    /// Create with custom config
+    pub fn with_config(config: StepContextConfig) -> Self {
+        Self {
+            token_counter: TokenCounter::default(),
+            config,
+        }
+    }
+
+    /// Build context from step data
+    pub fn build_context(
+        &self,
+        execution_id: ExecutionId,
+        step_id: StepId,
+        step_type: &str,
+        input: &str,
+        output: Option<&str>,
+        tool_calls: &[ToolCallInfo],
+        error: Option<&str>,
+        metadata: &HashMap<String, String>,
+    ) -> StepContextResult {
+        let mut segments = Vec::new();
+        let mut learnings = Vec::new();
+        let mut total_tokens = 0;
+
+        // Extract main step context
+        let step_summary = self.build_step_summary(step_type, input, output);
+        let summary_tokens = self.token_counter.count(&step_summary);
+
+        if total_tokens + summary_tokens <= self.config.max_tokens {
+            segments.push(ContextSegment::history(
+                step_summary.clone(),
+                summary_tokens,
+                next_sequence(),
+            ));
+            total_tokens += summary_tokens;
+        }
+
+        // Extract tool results
+        if self.config.include_tool_results {
+            let tool_context = self.extract_tool_context(tool_calls, step_id.clone());
+            for segment in tool_context {
+                let tokens = segment.token_count;
+                if total_tokens + tokens <= self.config.max_tokens {
+                    total_tokens += tokens;
+                    segments.push(segment);
+                }
+            }
+        }
+
+        // Extract error context
+        if self.config.include_errors {
+            if let Some(err) = error {
+                let error_learning = self.extract_error_learning(
+                    execution_id.clone(),
+                    step_id.clone(),
+                    err,
+                );
+                learnings.push(error_learning);
+
+                let error_content = format!("Error encountered: {}", self.truncate_content(err));
+                let error_tokens = self.token_counter.count(&error_content);
+                let error_segment = ContextSegment::tool_results(
+                    error_content,
+                    error_tokens,
+                    next_sequence(),
+                    step_id.clone(),
+                ).with_priority(ContextPriority::High);
+
+                if total_tokens + error_tokens <= self.config.max_tokens {
+                    total_tokens += error_tokens;
+                    segments.push(error_segment);
+                }
+            }
+        }
+
+        // Extract learnings from successful execution
+        if error.is_none() && output.is_some() {
+            let success_learnings = self.extract_success_learnings(
+                execution_id.clone(),
+                step_id.clone(),
+                step_type,
+                tool_calls,
+                metadata,
+            );
+            learnings.extend(success_learnings);
+        }
+
+        StepContextResult {
+            execution_id,
+            step_id,
+            segments,
+            learnings,
+            total_tokens,
+            processed_at: Utc::now(),
+        }
+    }
+
+    /// Build a summary of the step
+    fn build_step_summary(&self, step_type: &str, input: &str, output: Option<&str>) -> String {
+        let truncated_input = self.truncate_content(input);
+        let truncated_output = output
+            .map(|o| self.truncate_content(o))
+            .unwrap_or_else(|| "(pending)".to_string());
+
+        format!(
+            "[Step: {}]\nInput: {}\nOutput: {}",
+            step_type, truncated_input, truncated_output
+        )
+    }
+
+    /// Extract context from tool calls
+    fn extract_tool_context(&self, tool_calls: &[ToolCallInfo], step_id: StepId) -> Vec<ContextSegment> {
+        tool_calls
+            .iter()
+            .take(self.config.max_tool_results)
+            .map(|tc| {
+                let content = format!(
+                    "Tool: {}\nArgs: {}\nResult: {}",
+                    tc.tool_name,
+                    self.truncate_content(&tc.arguments),
+                    tc.result
+                        .as_ref()
+                        .map(|r| self.truncate_content(r))
+                        .unwrap_or_else(|| "(pending)".to_string())
+                );
+                let tokens = self.token_counter.count(&content);
+
+                ContextSegment::tool_results(
+                    content,
+                    tokens,
+                    next_sequence(),
+                    step_id.clone(),
+                ).with_priority(if tc.success {
+                    ContextPriority::Medium
+                } else {
+                    ContextPriority::High
+                })
+            })
+            .collect()
+    }
+
+    /// Extract learning from an error
+    fn extract_error_learning(
+        &self,
+        execution_id: ExecutionId,
+        step_id: StepId,
+        error: &str,
+    ) -> StepLearning {
+        StepLearning {
+            id: format!("learn_{}", uuid::Uuid::new_v4()),
+            step_id,
+            execution_id,
+            learning_type: LearningType::ErrorRecovery,
+            content: format!("Error encountered: {}. Consider alternative approaches.", error),
+            confidence: 0.7,
+            relevance: 0.8,
+            tags: vec!["error".to_string(), "recovery".to_string()],
+            created_at: Utc::now(),
+        }
+    }
+
+    /// Extract learnings from successful execution
+    fn extract_success_learnings(
+        &self,
+        execution_id: ExecutionId,
+        step_id: StepId,
+        step_type: &str,
+        tool_calls: &[ToolCallInfo],
+        metadata: &HashMap<String, String>,
+    ) -> Vec<StepLearning> {
+        let mut learnings = Vec::new();
+
+        // Learn from successful tool usage
+        for tc in tool_calls.iter().filter(|tc| tc.success) {
+            learnings.push(StepLearning {
+                id: format!("learn_{}", uuid::Uuid::new_v4()),
+                step_id: step_id.clone(),
+                execution_id: execution_id.clone(),
+                learning_type: LearningType::ToolInsight,
+                content: format!(
+                    "Tool '{}' succeeded with pattern: {}",
+                    tc.tool_name,
+                    self.truncate_content(&tc.arguments)
+                ),
+                confidence: 0.8,
+                relevance: 0.6,
+                tags: vec!["tool".to_string(), tc.tool_name.clone()],
+                created_at: Utc::now(),
+            });
+        }
+
+        // Learn from metadata hints
+        if let Some(pattern) = metadata.get("success_pattern") {
+            learnings.push(StepLearning {
+                id: format!("learn_{}", uuid::Uuid::new_v4()),
+                step_id: step_id.clone(),
+                execution_id: execution_id.clone(),
+                learning_type: LearningType::SuccessPattern,
+                content: format!("Step '{}' success pattern: {}", step_type, pattern),
+                confidence: 0.9,
+                relevance: 0.7,
+                tags: vec!["pattern".to_string(), step_type.to_string()],
+                created_at: Utc::now(),
+            });
+        }
+
+        learnings
+    }
+
+    /// Truncate content if too long
+    fn truncate_content(&self, content: &str) -> String {
+        if self.config.truncate_long_content && content.len() > self.config.max_content_length {
+            format!(
+                "{}... [truncated, {} chars total]",
+                &content[..self.config.max_content_length],
+                content.len()
+            )
+        } else {
+            content.to_string()
+        }
+    }
+
+    /// Build context for a child step being spawned
+    pub fn build_child_context(
+        &self,
+        parent_execution_id: ExecutionId,
+        parent_step_id: StepId,
+        child_step_id: StepId,
+        task: &str,
+        parent_context: &[ContextSegment],
+    ) -> StepContextResult {
+        let mut segments = Vec::new();
+        let mut total_tokens = 0;
+
+        // Add child task context
+        let task_content = format!(
+            "Sub-task spawned from parent step.\nTask: {}\nParent step: {}",
+            task,
+            parent_step_id.as_str()
+        );
+        let task_tokens = self.token_counter.count(&task_content);
+        let task_segment = ContextSegment::system(task_content, task_tokens);
+        total_tokens += task_tokens;
+        segments.push(task_segment);
+
+        // Include relevant parent context
+        for segment in parent_context {
+            if segment.priority >= ContextPriority::Medium {
+                let tokens = segment.token_count;
+                if total_tokens + tokens <= self.config.max_tokens {
+                    total_tokens += tokens;
+                    segments.push(segment.clone());
+                }
+            }
+        }
+
+        StepContextResult {
+            execution_id: parent_execution_id,
+            step_id: child_step_id,
+            segments,
+            learnings: Vec::new(),
+            total_tokens,
+            processed_at: Utc::now(),
+        }
+    }
+}
+
+impl Default for StepContextBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Information about a tool call
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ToolCallInfo {
+    /// Tool name
+    pub tool_name: String,
+
+    /// Arguments as JSON string
+    pub arguments: String,
+
+    /// Result if completed
+    pub result: Option<String>,
+
+    /// Whether the call succeeded
+    pub success: bool,
+
+    /// Duration in milliseconds
+    pub duration_ms: Option<u64>,
+}
+
+impl ToolCallInfo {
+    /// Create a successful tool call
+    pub fn success(tool_name: impl Into<String>, arguments: impl Into<String>, result: impl Into<String>) -> Self {
+        Self {
+            tool_name: tool_name.into(),
+            arguments: arguments.into(),
+            result: Some(result.into()),
+            success: true,
+            duration_ms: None,
+        }
+    }
+
+    /// Create a failed tool call
+    pub fn failed(tool_name: impl Into<String>, arguments: impl Into<String>, error: impl Into<String>) -> Self {
+        Self {
+            tool_name: tool_name.into(),
+            arguments: arguments.into(),
+            result: Some(error.into()),
+            success: false,
+            duration_ms: None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_execution_id() -> ExecutionId {
+        ExecutionId::new()
+    }
+
+    fn test_step_id() -> StepId {
+        StepId::new()
+    }
+
+    #[test]
+    fn test_step_context_config_defaults() {
+        let config = StepContextConfig::default();
+        assert_eq!(config.max_tokens, 2000);
+        assert!(config.include_tool_results);
+        assert!(config.include_errors);
+    }
+
+    #[test]
+    fn test_build_context_basic() {
+        let builder = StepContextBuilder::new();
+        let result = builder.build_context(
+            test_execution_id(),
+            test_step_id(),
+            "llm_call",
+            "What is 2+2?",
+            Some("4"),
+            &[],
+            None,
+            &HashMap::new(),
+        );
+
+        assert!(!result.segments.is_empty());
+        assert!(result.total_tokens > 0);
+    }
+
+    #[test]
+    fn test_build_context_with_error() {
+        let builder = StepContextBuilder::new();
+        let result = builder.build_context(
+            test_execution_id(),
+            test_step_id(),
+            "tool_call",
+            "fetch data",
+            None,
+            &[],
+            Some("Connection timeout"),
+            &HashMap::new(),
+        );
+
+        assert!(!result.learnings.is_empty());
+        assert_eq!(result.learnings[0].learning_type, LearningType::ErrorRecovery);
+    }
+
+    #[test]
+    fn test_build_context_with_tool_calls() {
+        let builder = StepContextBuilder::new();
+        let tool_calls = vec![
+            ToolCallInfo::success("search", r#"{"query": "test"}"#, "Found 5 results"),
+            ToolCallInfo::failed("fetch", r#"{"url": "..."}"#, "404 Not Found"),
+        ];
+
+        let result = builder.build_context(
+            test_execution_id(),
+            test_step_id(),
+            "multi_tool",
+            "search and fetch",
+            Some("partial results"),
+            &tool_calls,
+            None,
+            &HashMap::new(),
+        );
+
+        // Should have tool result segments
+        assert!(result.segments.len() >= 2);
+        // Should have tool insight learning from successful call
+        assert!(result
+            .learnings
+            .iter()
+            .any(|l| l.learning_type == LearningType::ToolInsight));
+    }
+
+    #[test]
+    fn test_truncate_long_content() {
+        let config = StepContextConfig {
+            max_content_length: 50,
+            ..Default::default()
+        };
+        let builder = StepContextBuilder::with_config(config);
+
+        let long_content = "a".repeat(100);
+        let result = builder.build_context(
+            test_execution_id(),
+            test_step_id(),
+            "test",
+            &long_content,
+            None,
+            &[],
+            None,
+            &HashMap::new(),
+        );
+
+        // Content should be truncated
+        assert!(result.segments[0].content.contains("truncated"));
+    }
+
+    #[test]
+    fn test_build_child_context() {
+        let builder = StepContextBuilder::new();
+        let token_counter = TokenCounter::default();
+
+        let system_content = "Parent system context";
+        let system_tokens = token_counter.count(system_content);
+
+        let history_content = "Some history";
+        let history_tokens = token_counter.count(history_content);
+
+        let parent_context = vec![
+            ContextSegment::system(system_content, system_tokens),
+            ContextSegment::new(
+                ContextSegmentType::History,
+                history_content.to_string(),
+                history_tokens,
+                1,
+            ).with_priority(ContextPriority::Low),
+        ];
+
+        let result = builder.build_child_context(
+            test_execution_id(),
+            test_step_id(),
+            StepId::new(),
+            "Analyze the data",
+            &parent_context,
+        );
+
+        // Should have task segment
+        assert!(result
+            .segments
+            .iter()
+            .any(|s| s.content.contains("Sub-task")));
+        // Should include high-priority parent context but not low
+        assert!(result
+            .segments
+            .iter()
+            .any(|s| s.content.contains("Parent system")));
+    }
+
+    #[test]
+    fn test_learning_types() {
+        let builder = StepContextBuilder::new();
+        let mut metadata = HashMap::new();
+        metadata.insert("success_pattern".to_string(), "retry with backoff".to_string());
+
+        let result = builder.build_context(
+            test_execution_id(),
+            test_step_id(),
+            "api_call",
+            "fetch user",
+            Some("user data"),
+            &[ToolCallInfo::success("http", "{}", "200 OK")],
+            None,
+            &metadata,
+        );
+
+        // Should have both tool insight and success pattern learnings
+        assert!(result
+            .learnings
+            .iter()
+            .any(|l| l.learning_type == LearningType::ToolInsight));
+        assert!(result
+            .learnings
+            .iter()
+            .any(|l| l.learning_type == LearningType::SuccessPattern));
+    }
+}