igny8/backend/igny8_core/ai/tracker.py

"""
Progress and Step Tracking utilities for AI framework
"""
import time
import logging
from typing import List, Dict, Any, Optional, Callable
from igny8_core.ai.types import StepLog, ProgressState

logger = logging.getLogger(__name__)


class StepTracker:
    """Tracks detailed request and response steps for debugging"""
    
    def __init__(self, function_name: str):
        self.function_name = function_name
        self.request_steps: List[Dict] = []
        self.response_steps: List[Dict] = []
        self.step_counter = 0
    
    def add_request_step(
        self,
        step_name: str,
        status: str = 'success',
        message: str = '',
        error: str = None,
        duration: int = None
    ) -> Dict:
        """Add a request step with automatic timing"""
        self.step_counter += 1
        step = {
            'stepNumber': self.step_counter,
            'stepName': step_name,
            'functionName': self.function_name,
            'status': status,
            'message': message,
            'duration': duration
        }
        if error:
            step['error'] = error
        
        self.request_steps.append(step)
        return step
    
    def add_response_step(
        self,
        step_name: str,
        status: str = 'success',
        message: str = '',
        error: str = None,
        duration: int = None
    ) -> Dict:
        """Add a response step with automatic timing"""
        self.step_counter += 1
        step = {
            'stepNumber': self.step_counter,
            'stepName': step_name,
            'functionName': self.function_name,
            'status': status,
            'message': message,
            'duration': duration
        }
        if error:
            step['error'] = error
        
        self.response_steps.append(step)
        return step
    
    def get_meta(self) -> Dict:
        """Get metadata for progress callback"""
        return {
            'request_steps': self.request_steps,
            'response_steps': self.response_steps
        }


class ProgressTracker:
    """Tracks progress updates for AI tasks"""
    
    def __init__(self, celery_task=None):
        self.task = celery_task
        self.current_phase = 'INIT'
        self.current_message = 'Initializing...'
        self.current_percentage = 0
        self.start_time = time.time()
        self.current = 0
        self.total = 0
    
    def update(
        self,
        phase: str,
        percentage: int,
        message: str,
        current: int = None,
        total: int = None,
        current_item: str = None,
        meta: Dict = None
    ):
        """Update progress with consistent format"""
        self.current_phase = phase
        self.current_message = message
        self.current_percentage = percentage
        
        if current is not None:
            self.current = current
        if total is not None:
            self.total = total
        
        progress_meta = {
            'phase': phase,
            'percentage': percentage,
            'message': message,
            'current': self.current,
            'total': self.total,
        }
        
        if current_item:
            progress_meta['current_item'] = current_item
        
        if meta:
            progress_meta.update(meta)
        
        # Update Celery task state if available
        if self.task:
            try:
                self.task.update_state(
                    state='PROGRESS',
                    meta=progress_meta
                )
            except Exception as e:
                logger.warning(f"Failed to update Celery task state: {e}")
        
        logger.info(f"[{phase}] {percentage}%: {message}")
    
    def set_phase(self, phase: str, percentage: int, message: str, meta: Dict = None):
        """Set progress phase"""
        self.update(phase, percentage, message, meta=meta)
    
    def complete(self, message: str = "Task complete!", meta: Dict = None):
        """Mark task as complete"""
        final_meta = {
            'phase': 'DONE',
            'percentage': 100,
            'message': message,
            'status': 'success'
        }
        if meta:
            final_meta.update(meta)
        
        if self.task:
            try:
                self.task.update_state(
                    state='SUCCESS',
                    meta=final_meta
                )
            except Exception as e:
                logger.warning(f"Failed to update Celery task state: {e}")
    
    def error(self, error_message: str, meta: Dict = None):
        """Mark task as failed"""
        error_meta = {
            'phase': 'ERROR',
            'percentage': 0,
            'message': f'Error: {error_message}',
            'status': 'error',
            'error': error_message
        }
        if meta:
            error_meta.update(meta)
        
        if self.task:
            try:
                self.task.update_state(
                    state='FAILURE',
                    meta=error_meta
                )
            except Exception as e:
                logger.warning(f"Failed to update Celery task state: {e}")
    
    def get_duration(self) -> int:
        """Get elapsed time in milliseconds"""
        return int((time.time() - self.start_time) * 1000)
    
    def update_ai_progress(self, state: str, meta: Dict):
        """Callback for AI processor progress updates"""
        if isinstance(meta, dict):
            percentage = meta.get('percentage', self.current_percentage)
            message = meta.get('message', self.current_message)
            phase = meta.get('phase', self.current_phase)
            self.update(phase, percentage, message, meta=meta)


class CostTracker:
    """Tracks API costs and token usage"""
    
    def __init__(self):
        self.total_cost = 0.0
        self.total_tokens = 0
        self.operations = []
    
    def record(self, function_name: str, cost: float, tokens: int, model: str = None):
        """Record an API call cost"""
        self.total_cost += cost
        self.total_tokens += tokens
        self.operations.append({
            'function': function_name,
            'cost': cost,
            'tokens': tokens,
            'model': model
        })
    
    def get_total(self) -> float:
        """Get total cost"""
        return self.total_cost
    
    def get_total_tokens(self) -> int:
        """Get total tokens"""
        return self.total_tokens
    
    def get_operations(self) -> List[Dict]:
        """Get all operations"""
        return self.operations