tkb_timeshift/claude_n8n/tools/workflow_improver.py

#!/usr/bin/env python3
"""
Workflow Improver - Iterative improvement and testing framework for N8N workflows
Implements automated testing, optimization, and iterative refinement capabilities
"""

import json
import copy
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass
from datetime import datetime
import logging


@dataclass
class TestCase:
    """Represents a test case for workflow validation"""
    name: str
    input_data: Dict
    expected_output: Optional[Dict] = None
    expected_status: str = "success"
    description: str = ""


@dataclass
class ImprovementResult:
    """Result of workflow improvement iteration"""
    iteration: int
    original_workflow: Dict
    improved_workflow: Dict
    test_results: List[Dict]
    improvements_made: List[str]
    performance_metrics: Dict
    success: bool
    error_message: Optional[str] = None


class WorkflowImprover:
    """Implements iterative workflow improvement and testing"""
    
    def __init__(self, n8n_client, analyzer, monitor):
        """Initialize workflow improver"""
        self.client = n8n_client
        self.analyzer = analyzer
        self.monitor = monitor
        self.logger = self._setup_logger()
        
        # Improvement strategies
        self.improvement_strategies = {
            'add_error_handling': self._add_error_handling,
            'optimize_timeouts': self._optimize_timeouts,
            'add_retry_logic': self._add_retry_logic,
            'improve_validation': self._improve_validation,
            'optimize_performance': self._optimize_performance,
            'fix_connections': self._fix_connections
        }
    
    def _setup_logger(self) -> logging.Logger:
        """Setup logging for the improver"""
        logger = logging.getLogger('N8NWorkflowImprover')
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def create_test_suite(self, workflow: Dict, sample_data: List[Dict] = None) -> List[TestCase]:
        """Create comprehensive test suite for a workflow"""
        test_cases = []
        
        # Basic functionality test
        test_cases.append(TestCase(
            name="basic_functionality",
            input_data=sample_data[0] if sample_data else {},
            expected_status="success",
            description="Test basic workflow functionality"
        ))
        
        # Error handling tests
        test_cases.append(TestCase(
            name="invalid_input",
            input_data={"invalid": "data"},
            expected_status="error",
            description="Test error handling with invalid input"
        ))
        
        # Empty data test
        test_cases.append(TestCase(
            name="empty_input",
            input_data={},
            expected_status="success",
            description="Test workflow with empty input data"
        ))
        
        # Large data test (if applicable)
        if sample_data and len(sample_data) > 1:
            test_cases.append(TestCase(
                name="large_dataset",
                input_data={"batch": sample_data},
                expected_status="success",
                description="Test workflow with larger dataset"
            ))
        
        return test_cases
    
    def run_test_suite(self, workflow_id: str, test_cases: List[TestCase]) -> List[Dict]:
        """Run complete test suite against a workflow"""
        results = []
        
        for test_case in test_cases:
            self.logger.info(f"Running test case: {test_case.name}")
            
            try:
                # Execute workflow with test data
                execution_event = self.monitor.execute_and_monitor(
                    workflow_id, 
                    test_case.input_data,
                    timeout=120
                )
                
                # Analyze results
                test_result = {
                    'test_name': test_case.name,
                    'description': test_case.description,
                    'input_data': test_case.input_data,
                    'expected_status': test_case.expected_status,
                    'actual_status': execution_event.status.value,
                    'execution_time': execution_event.duration,
                    'passed': execution_event.status.value == test_case.expected_status,
                    'execution_id': execution_event.execution_id,
                    'error_message': execution_event.error_message,
                    'timestamp': datetime.now().isoformat()
                }
                
                # Validate output if expected output is provided
                if test_case.expected_output and execution_event.node_data:
                    output_match = self._validate_output(
                        execution_event.node_data, 
                        test_case.expected_output
                    )
                    test_result['output_validation'] = output_match
                    test_result['passed'] = test_result['passed'] and output_match
                
                results.append(test_result)
                
            except Exception as e:
                self.logger.error(f"Test case {test_case.name} failed with exception: {e}")
                results.append({
                    'test_name': test_case.name,
                    'description': test_case.description,
                    'passed': False,
                    'error_message': str(e),
                    'timestamp': datetime.now().isoformat()
                })
        
        return results
    
    def iterative_improvement(self, workflow_id: str, test_cases: List[TestCase], 
                            max_iterations: int = 5) -> List[ImprovementResult]:
        """Perform iterative improvement on a workflow"""
        results = []
        current_workflow = self.client.get_workflow(workflow_id)
        
        for iteration in range(max_iterations):
            self.logger.info(f"Starting improvement iteration {iteration + 1}")
            
            try:
                # Run tests on current workflow
                test_results = self.run_test_suite(workflow_id, test_cases)
                
                # Analyze workflow for issues
                analysis = self.analyzer.analyze_workflow_structure(current_workflow)
                
                # Check if workflow is already performing well
                passed_tests = len([r for r in test_results if r.get('passed', False)])
                test_success_rate = passed_tests / len(test_results) if test_results else 0
                
                if test_success_rate >= 0.9 and len(analysis['issues']) == 0:
                    self.logger.info("Workflow is already performing well, no improvements needed")
                    break
                
                # Generate improvements
                improved_workflow, improvements_made = self._generate_improvements(
                    current_workflow, analysis, test_results
                )
                
                if not improvements_made:
                    self.logger.info("No more improvements can be made")
                    break
                
                # Apply improvements
                self.client.update_workflow(workflow_id, improved_workflow)
                
                # Run tests again to validate improvements
                new_test_results = self.run_test_suite(workflow_id, test_cases)
                
                # Calculate performance metrics
                performance_metrics = self._calculate_performance_improvement(
                    test_results, new_test_results
                )
                
                result = ImprovementResult(
                    iteration=iteration + 1,
                    original_workflow=current_workflow,
                    improved_workflow=improved_workflow,
                    test_results=new_test_results,
                    improvements_made=improvements_made,
                    performance_metrics=performance_metrics,
                    success=True
                )
                
                results.append(result)
                current_workflow = improved_workflow
                
                self.logger.info(f"Iteration {iteration + 1} completed with {len(improvements_made)} improvements")
                
            except Exception as e:
                self.logger.error(f"Error in iteration {iteration + 1}: {e}")
                result = ImprovementResult(
                    iteration=iteration + 1,
                    original_workflow=current_workflow,
                    improved_workflow=current_workflow,
                    test_results=[],
                    improvements_made=[],
                    performance_metrics={},
                    success=False,
                    error_message=str(e)
                )
                results.append(result)
                break
        
        return results
    
    def _generate_improvements(self, workflow: Dict, analysis: Dict, 
                             test_results: List[Dict]) -> Tuple[Dict, List[str]]:
        """Generate workflow improvements based on analysis and test results"""
        improved_workflow = copy.deepcopy(workflow)
        improvements_made = []
        
        # Apply improvements based on structural issues
        for issue in analysis.get('issues', []):
            issue_type = issue.get('type')
            
            if issue_type in self.improvement_strategies:
                strategy_func = self.improvement_strategies[issue_type]
                workflow_modified, improvement_desc = strategy_func(
                    improved_workflow, issue
                )
                
                if workflow_modified:
                    improvements_made.append(improvement_desc)
        
        # Apply improvements based on test failures
        failed_tests = [r for r in test_results if not r.get('passed', False)]
        for test_result in failed_tests:
            improvement = self._improve_based_on_test_failure(
                improved_workflow, test_result
            )
            if improvement:
                improvements_made.append(improvement)
        
        return improved_workflow, improvements_made
    
    def _add_error_handling(self, workflow: Dict, issue: Dict) -> Tuple[bool, str]:
        """Add error handling to nodes"""
        node_name = issue.get('node')
        if not node_name:
            return False, ""
        
        # Find the node and add error handling
        for node in workflow.get('nodes', []):
            if node.get('name') == node_name:
                parameters = node.get('parameters', {})
                parameters['continueOnFail'] = True
                
                # Add error handling parameters based on node type
                node_type = node.get('type', '')
                if 'httpRequest' in node_type:
                    parameters['retry'] = {
                        'retries': 3,
                        'waitBetween': 1000
                    }
                
                return True, f"Added error handling to node '{node_name}'"
        
        return False, ""
    
    def _optimize_timeouts(self, workflow: Dict, issue: Dict) -> Tuple[bool, str]:
        """Optimize timeout settings"""
        node_name = issue.get('node')
        if not node_name:
            return False, ""
        
        for node in workflow.get('nodes', []):
            if node.get('name') == node_name:
                parameters = node.get('parameters', {})
                current_timeout = parameters.get('timeout', 300)
                
                # Increase timeout if it's too aggressive
                if current_timeout < 60:
                    parameters['timeout'] = 60
                    return True, f"Increased timeout for node '{node_name}' to 60 seconds"
        
        return False, ""
    
    def _add_retry_logic(self, workflow: Dict, issue: Dict) -> Tuple[bool, str]:
        """Add retry logic to nodes"""
        # This would add retry nodes or modify existing nodes with retry parameters
        return False, "Retry logic addition not implemented"
    
    def _improve_validation(self, workflow: Dict, issue: Dict) -> Tuple[bool, str]:
        """Improve input validation"""
        # This would add validation nodes or improve existing validation
        return False, "Validation improvement not implemented"
    
    def _optimize_performance(self, workflow: Dict, issue: Dict) -> Tuple[bool, str]:
        """Optimize workflow performance"""
        # This could involve optimizing loops, reducing unnecessary operations, etc.
        return False, "Performance optimization not implemented"
    
    def _fix_connections(self, workflow: Dict, issue: Dict) -> Tuple[bool, str]:
        """Fix disconnected nodes"""
        description = issue.get('description', '')
        
        # Extract disconnected node names from description
        if "Disconnected nodes found:" in description:
            disconnected_nodes = description.split(": ")[1].split(", ")
            
            # Remove disconnected nodes
            original_count = len(workflow.get('nodes', []))
            workflow['nodes'] = [
                node for node in workflow.get('nodes', [])
                if node.get('name') not in disconnected_nodes
            ]
            
            removed_count = original_count - len(workflow['nodes'])
            if removed_count > 0:
                return True, f"Removed {removed_count} disconnected nodes"
        
        return False, ""
    
    def _improve_based_on_test_failure(self, workflow: Dict, test_result: Dict) -> Optional[str]:
        """Improve workflow based on specific test failure"""
        test_name = test_result.get('test_name')
        error_message = test_result.get('error_message', '')
        
        if test_name == "invalid_input" and "validation" in error_message.lower():
            # Add input validation
            return "Added input validation based on test failure"
        
        elif "timeout" in error_message.lower():
            # Increase timeouts
            return "Increased timeouts based on test failure"
        
        return None
    
    def _validate_output(self, actual_output: Dict, expected_output: Dict) -> bool:
        """Validate workflow output against expected results"""
        try:
            # Simple validation - check if expected keys exist and values match
            for key, expected_value in expected_output.items():
                if key not in actual_output:
                    return False
                
                if isinstance(expected_value, dict):
                    if not self._validate_output(actual_output[key], expected_value):
                        return False
                elif actual_output[key] != expected_value:
                    return False
            
            return True
        except Exception:
            return False
    
    def _calculate_performance_improvement(self, old_results: List[Dict], 
                                         new_results: List[Dict]) -> Dict:
        """Calculate performance improvement metrics"""
        old_success_rate = len([r for r in old_results if r.get('passed', False)]) / len(old_results) if old_results else 0
        new_success_rate = len([r for r in new_results if r.get('passed', False)]) / len(new_results) if new_results else 0
        
        old_avg_time = sum([r.get('execution_time', 0) for r in old_results if r.get('execution_time')]) / len(old_results) if old_results else 0
        new_avg_time = sum([r.get('execution_time', 0) for r in new_results if r.get('execution_time')]) / len(new_results) if new_results else 0
        
        return {
            'success_rate_improvement': new_success_rate - old_success_rate,
            'performance_improvement_percent': ((old_avg_time - new_avg_time) / old_avg_time * 100) if old_avg_time > 0 else 0,
            'old_success_rate': old_success_rate,
            'new_success_rate': new_success_rate,
            'old_avg_execution_time': old_avg_time,
            'new_avg_execution_time': new_avg_time
        }
    
    def create_test_data_from_execution(self, execution_id: str) -> Dict:
        """Create test data from a successful execution"""
        try:
            execution = self.client.get_execution(execution_id)
            
            if execution.get('status') != 'success':
                raise ValueError("Can only create test data from successful executions")
            
            # Extract input data from the execution
            data = execution.get('data', {})
            if 'resultData' in data and 'runData' in data['resultData']:
                run_data = data['resultData']['runData']
                
                # Find the trigger or start node data
                for node_name, node_runs in run_data.items():
                    if node_runs and 'data' in node_runs[0]:
                        node_data = node_runs[0]['data']
                        if 'main' in node_data and node_data['main']:
                            return node_data['main'][0][0]  # First item of first output
            
            return {}
        
        except Exception as e:
            self.logger.error(f"Error creating test data from execution: {e}")
            return {}
    
    def benchmark_workflow(self, workflow_id: str, iterations: int = 10) -> Dict:
        """Benchmark workflow performance"""
        results = []
        
        for i in range(iterations):
            try:
                execution_event = self.monitor.execute_and_monitor(workflow_id, {})
                results.append({
                    'iteration': i + 1,
                    'status': execution_event.status.value,
                    'duration': execution_event.duration,
                    'success': execution_event.status.value == 'success'
                })
            except Exception as e:
                results.append({
                    'iteration': i + 1,
                    'status': 'error',
                    'duration': None,
                    'success': False,
                    'error': str(e)
                })
        
        successful_runs = [r for r in results if r['success']]
        durations = [r['duration'] for r in successful_runs if r['duration']]
        
        return {
            'total_iterations': iterations,
            'successful_runs': len(successful_runs),
            'success_rate': len(successful_runs) / iterations * 100,
            'average_duration': sum(durations) / len(durations) if durations else 0,
            'min_duration': min(durations) if durations else 0,
            'max_duration': max(durations) if durations else 0,
            'detailed_results': results
        }


if __name__ == "__main__":
    print("Workflow Improver initialized successfully.")