Session 1 - Module A: Advanced Agent Patterns¶

# From src/session1/reflection_agent.py
from base_agent import BaseAgent
from typing import Dict

class ReflectionAgent(BaseAgent):
    """Agent that reflects on and improves its own data processing outputs"""

    def __init__(self, name: str, llm_client, max_iterations: int = 3):
        super().__init__(name, "Agent with data processing reflection capabilities", llm_client)
        self.max_iterations = max_iterations  # Prevent infinite optimization loops
        self.reflection_history = []          # Track data processing improvements

# From src/session1/reflection_agent.py (continued)
async def _generate_response(self, message: str, context: Dict = None) -> str:
    """Generate data processing response with reflection and optimization"""
    current_response = await self._initial_data_processing(message, context)

    for iteration in range(self.max_iterations):
        # Step 1: Reflect on current data processing approach
        critique = await self._reflect_on_processing(message, current_response)

        # Step 2: If processing approach is optimal, return it
        if self._is_processing_satisfactory(critique):
            self.logger.info(f"Data processing optimal after {iteration + 1} iterations")
            break

        # Step 3: Optimize processing based on critique
        improved_response = await self._improve_processing(
            message, current_response, critique
        )

        # Step 4: Track the optimization process for pipeline metrics
        self._track_reflection(iteration, current_response, critique, improved_response)

        current_response = improved_response

    return current_response

async def _initial_data_processing(self, message: str, context: Dict = None) -> str:
    """Generate initial data processing approach without reflection"""
    system_prompt = f"""
    You are {self.name}, {self.description}.
    Design an optimal data processing approach for the given requirements.
    Focus on throughput, data quality, and cost efficiency for large-scale systems.
    Consider partitioning strategies, batch sizing, and resource utilization.
    """

    prompt = f"{system_prompt}\n\nProcessing requirements: {message}"
    response = await self._call_llm(prompt)
    return response

async def _reflect_on_processing(self, original_message: str, response: str) -> str:
    """Generate critique of the current data processing approach"""
    critique_prompt = f"""
    Analyze this data processing approach and provide constructive criticism.

    Processing requirements: {original_message}
    Proposed approach: {response}

    Evaluate from a data engineering perspective:
    1. Scalability: Can this handle petabyte-scale datasets?
    2. Performance: Are batch sizes and parallelization optimal?
    3. Cost efficiency: Does this minimize cloud compute costs?
    4. Data quality: Are validation and error handling comprehensive?
    5. Reliability: Will this maintain SLAs under high load?
    6. Monitoring: Can pipeline health be effectively tracked?

    Provide specific suggestions for improvement based on production data engineering best practices.
    If the approach is already optimal for enterprise scale, say "SATISFACTORY".
    """

    critique = await self._call_llm(critique_prompt)
    return critique

async def _improve_processing(self, message: str, current_response: str, critique: str) -> str:
    """Improve data processing approach based on critique"""
    improvement_prompt = f"""
    Optimize the following data processing approach based on the critique provided.

    Original requirements: {message}
    Current processing approach: {current_response}
    Critique focusing on data engineering concerns: {critique}

    Generate an improved approach that addresses the critique while maintaining:
    - High throughput for large datasets
    - Cost optimization for cloud deployment
    - Robust error handling and data quality validation
    - Scalable architecture for petabyte-scale processing
    """

    improved = await self._call_llm(improvement_prompt)
    return improved

from typing import List, Dict, Optional
from dataclasses import dataclass
from enum import Enum

class DataProcessingStatus(Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"

@dataclass
class DataPipelineStep:
    id: str
    description: str
    dependencies: List[str]
    estimated_time: int
    resource_requirements: Dict[str, any]  # CPU, memory, storage needs
    status: DataProcessingStatus = DataProcessingStatus.PENDING
    result: Optional[str] = None

class DataPipelinePlanningAgent(BaseAgent):
    """Agent that breaks down complex data processing tasks into executable pipeline steps"""

    def __init__(self, name: str, llm_client):
        super().__init__(name, "Multi-step data pipeline planning agent", llm_client)
        self.current_plan: Optional[List[DataPipelineStep]] = None
        self.execution_history: List[Dict] = []
        self.resource_tracker = DataResourceTracker()  # Track cloud resource usage

async def create_data_pipeline_plan(self, complex_data_task: str) -> List[DataPipelineStep]:
    """Break down complex data processing task into executable pipeline steps"""

    planning_prompt = f"""
    Break down this complex data processing task into specific, executable pipeline steps:
    Task: {complex_data_task}

    For each step, provide:
    1. Clear description of data processing operation
    2. Dependencies (which steps must complete first for data availability)
    3. Estimated time in minutes for large-scale processing
    4. Resource requirements (CPU, memory, storage for cloud deployment)
    5. Data quality checkpoints and validation requirements

    Consider:
    - Optimal batch sizing for throughput
    - Parallelization opportunities
    - Data partitioning strategies
    - Error handling and recovery mechanisms
    - Cost optimization through efficient resource usage

    Format as a numbered list with dependencies and resource requirements noted.
    Be specific about data volumes, formats, and processing requirements.
    """

    plan_text = await self._call_llm(planning_prompt)

    # Parse the plan into DataPipelineStep objects
    steps = self._parse_data_pipeline_plan(plan_text)
    self.current_plan = steps

    return steps

def _parse_data_pipeline_plan(self, plan_text: str) -> List[DataPipelineStep]:
    """Parse LLM-generated data pipeline plan into structured steps"""
    steps = []
    lines = plan_text.strip().split('\n')

    for i, line in enumerate(lines):
        if line.strip() and any(line.startswith(str(j)) for j in range(1, 20)):
            # Extract data processing step information
            step_id = f"data_step_{i+1}"
            description = line.split('.', 1)[1].strip() if '.' in line else line.strip()

            # Extract data processing dependencies and resource requirements
            dependencies = self._extract_data_dependencies(description)
            resource_reqs = self._extract_resource_requirements(description)

            steps.append(DataPipelineStep(
                id=step_id,
                description=description,
                dependencies=dependencies,
                estimated_time=15,  # Default for data processing operations
                resource_requirements=resource_reqs
            ))

    return steps

async def execute_data_pipeline(self) -> Dict[str, any]:
    """Execute the current data pipeline with dependency and resource management"""
    if not self.current_plan:
        return {"error": "No data pipeline to execute"}

    execution_log = []
    completed_steps = set()
    resource_usage = {"total_cost": 0, "peak_memory": 0, "total_compute_hours": 0}

    while len(completed_steps) < len(self.current_plan):
        # Find next executable step (dependencies satisfied, resources available)
        next_step = self._find_next_executable_data_step(completed_steps)

        if not next_step:
            return {"error": "No executable steps found - check dependencies or resource constraints"}

        # Allocate resources and execute the data processing step
        self.logger.info(f"Executing data processing step: {next_step.description}")
        next_step.status = DataProcessingStatus.IN_PROGRESS

        # Track resource allocation for cost monitoring
        allocated_resources = await self.resource_tracker.allocate_resources(
            next_step.resource_requirements
        )

        try:
            result = await self._execute_data_step(next_step, allocated_resources)
            next_step.result = result
            next_step.status = DataProcessingStatus.COMPLETED
            completed_steps.add(next_step.id)

            # Update resource usage tracking
            resource_usage = self.resource_tracker.update_usage(
                allocated_resources, result.get('metrics', {})
            )

            execution_log.append({
                "step": next_step.description,
                "result": result,
                "status": "completed",
                "resources_used": allocated_resources,
                "processing_metrics": result.get('metrics', {})
            })

        except Exception as e:
            next_step.status = DataProcessingStatus.FAILED
            execution_log.append({
                "step": next_step.description,
                "error": str(e),
                "status": "failed",
                "resources_allocated": allocated_resources
            })
            # Release allocated resources on failure
            await self.resource_tracker.release_resources(allocated_resources)
            break
        finally:
            # Always release resources after step completion
            await self.resource_tracker.release_resources(allocated_resources)

    return {
        "execution_log": execution_log,
        "completed_steps": len(completed_steps),
        "total_steps": len(self.current_plan),
        "success": len(completed_steps) == len(self.current_plan),
        "resource_usage": resource_usage,
        "cost_analysis": self.resource_tracker.generate_cost_report()
    }

async def _execute_data_step(self, step: DataPipelineStep, resources: Dict) -> str:
    """Execute an individual data pipeline step with allocated resources"""
    execution_prompt = f"""
    Execute this specific data processing task:
    {step.description}

    Allocated resources: {resources}

    Provide a detailed result including:
    1. What data processing operations were completed
    2. Data volume processed and throughput achieved
    3. Quality metrics and validation results
    4. Performance metrics (processing time, resource utilization)
    5. Any data quality issues or anomalies detected

    Be specific about data formats, processing efficiency, and outcomes.
    """

    result = await self._call_llm(execution_prompt)
    return result

from typing import Dict, List, Any
from base_agent import BaseAgent, AgentMessage
from datetime import datetime

class DataProcessingCoordinator:
    """Coordinates multiple specialized data processing agents in a distributed system"""

    def __init__(self):
        self.agents: Dict[str, BaseAgent] = {}  # Specialized data processing agents
        self.message_history = []               # All inter-agent data processing messages
        self.communication_patterns = {}       # Track data flow patterns
        self.pipeline_metrics = {}            # Track processing performance metrics

    def register_data_agent(self, agent: BaseAgent):
        """Register a specialized data processing agent with the coordinator"""
        self.agents[agent.name] = agent
        self.communication_patterns[agent.name] = []
        self.pipeline_metrics[agent.name] = {"processed_data_gb": 0, "processing_time": 0}
        print(f"✓ Registered data processing agent: {agent.name} ({agent.description})")

async def route_data_message(self, message: str, to_agent: str, from_agent: str = "pipeline_controller") -> str:
    """Route data processing message to specific agent and track communication"""
    if to_agent not in self.agents:
        return f"Error: Data processing agent '{to_agent}' not found"

    # Log data processing communication pattern
    self.communication_patterns[to_agent].append({
        "from": from_agent,
        "message_type": self._classify_data_message(message),
        "message": message[:50] + "..." if len(message) > 50 else message,
        "timestamp": datetime.now()
    })

    # Process message with target data processing agent
    agent = self.agents[to_agent]
    response = await agent.process_message(message)

    # Store complete data processing conversation
    self.message_history.append({
        "from": from_agent,
        "to": to_agent,
        "message": message,
        "response": response,
        "processing_metrics": self._extract_processing_metrics(response),
        "timestamp": datetime.now()
    )

    return response

async def delegate_data_processing_task(self, task: str) -> Dict[str, Any]:
    """Delegate complex data processing task to appropriate specialized agents"""

    # Step 1: Analyze data processing requirements to determine which agents are needed
    task_analysis = await self._analyze_data_task_requirements(task)

    # Step 2: Create delegation plan based on data processing capabilities
    delegation_plan = self._create_data_delegation_plan(task_analysis)

    # Step 3: Execute delegation plan with resource coordination

    results = {}
    processing_metrics = {}
    for agent_name, subtask in delegation_plan.items():
        if agent_name in self.agents:
            start_time = datetime.now()
            result = await self.route_data_message(subtask, agent_name, "coordinator")
            end_time = datetime.now()

            results[agent_name] = result
            processing_metrics[agent_name] = {
                "processing_time": (end_time - start_time).total_seconds(),
                "data_processed": self._extract_data_volume(result),
                "resource_efficiency": self._calculate_efficiency(result)
            }

    # Step 4: Integrate data processing results
    final_result = await self._integrate_data_results(task, results)

    return {
        "original_task": task,
        "delegation_plan": delegation_plan,
        "agent_results": results,
        "processing_metrics": processing_metrics,
        "integrated_result": final_result,
        "total_data_processed": sum(m.get('data_processed', 0) for m in processing_metrics.values()),
        "pipeline_efficiency": self._calculate_pipeline_efficiency(processing_metrics)
    }

async def _analyze_data_task_requirements(self, task: str) -> Dict[str, Any]:
    """Analyze what types of data processing agents are needed for this task"""
    analysis_prompt = f"""
    Analyze this data processing task and determine what types of specialized agents would be needed:
    Task: {task}

    Available specialized data processing agents: {list(self.agents.keys())}

    Consider the data processing pipeline requirements:
    1. Data ingestion and validation needs
    2. Transformation and enrichment requirements
    3. Quality assurance and anomaly detection needs
    4. Storage optimization and partitioning requirements
    5. Analytics and reporting capabilities needed

    For each needed agent type, specify:
    1. What specific data processing subtask they should handle
    2. Why this specialized agent is appropriate for this data operation
    3. Expected data volume and processing requirements
    4. Dependencies on other agents' outputs

    Return analysis focusing on data engineering optimization and scalability.
    """

    # This would use an LLM to analyze - simplified for example
    return {"required_agents": list(self.agents.keys())[:2]}  # Simplified