@ngocsangairvds/vsaf 3.2.13 → 3.2.15

package/tools/vds-scripts/multi_agent_orchestrator/src/vds_multi_agent_orchestrator/task_tracking.py

@@ -1,550 +0,0 @@
-"""Multi-agent task tracking system with Redis-based coordination."""
-
-import logging
-from datetime import UTC, datetime, timedelta
-from enum import StrEnum
-from typing import Any
-
-from .config import get_settings
-from .models import TaskStatusEnum
-from .redis_client import get_redis_client
-
-logger = logging.getLogger(__name__)
-
-
-class CoordinationPatternEnum(StrEnum):
-    """Coordination pattern enumeration."""
-
-    HIERARCHICAL = "hierarchical"
-    PARALLEL = "parallel"
-    SEQUENTIAL = "sequential"
-
-
-class TaskTracker:
-    """Distributed task tracking with Redis coordination."""
-
-    def __init__(self):
-        self.settings = get_settings()
-        self._redis = None
-
-    async def initialize(self):
-        """Initialize Redis connection."""
-        self._redis = await get_redis_client()
-
-    async def create_task_workflow(self, session_id: str, workflow_config: dict[str, Any]) -> str:
-        """Create a multi-agent task workflow with dependencies."""
-
-        workflow_id = f"workflow_{session_id}_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}"
-
-        # Store workflow configuration
-        await self._redis.hset(
-            f"workflow:{workflow_id}",
-            mapping={
-                "session_id": session_id,
-                "pattern": workflow_config.get("pattern", CoordinationPatternEnum.PARALLEL),
-                "created_at": datetime.now(UTC).isoformat(),
-                "status": "active",
-                "total_tasks": len(workflow_config.get("tasks", [])),
-                "completed_tasks": 0,
-            },
-        )
-
-        # Create task dependencies
-        for task_config in workflow_config.get("tasks", []):
-            task_id = task_config["task_id"]
-            dependencies = task_config.get("dependencies", [])
-
-            # Store task metadata
-            await self._redis.hset(
-                f"task_meta:{task_id}",
-                mapping={
-                    "workflow_id": workflow_id,
-                    "service_id": task_config.get("service_id", ""),
-                    "task_type": task_config.get("task_type", ""),
-                    "required_capabilities": ",".join(task_config.get("required_capabilities", [])),
-                    "estimated_duration_minutes": task_config.get("estimated_duration_minutes", 60),
-                    "priority": task_config.get("priority", "medium"),
-                    "created_at": datetime.now(UTC).isoformat(),
-                    "status": TaskStatusEnum.PENDING,
-                },
-            )
-
-            # Store dependencies
-            if dependencies:
-                await self._redis.sadd(f"task_deps:{task_id}", *dependencies)
-
-            # Add to workflow task list
-            await self._redis.sadd(f"workflow_tasks:{workflow_id}", task_id)
-
-        logger.info(
-            f"Created workflow {workflow_id} with {len(workflow_config.get('tasks', []))} tasks"
-        )
-        return workflow_id
-
-    async def get_ready_tasks(
-        self, workflow_id: str, agent_capabilities: list[str]
-    ) -> list[dict[str, Any]]:
-        """Get tasks that are ready to be executed by an agent with given capabilities."""
-
-        ready_tasks = []
-        workflow_tasks = await self._redis.smembers(f"workflow_tasks:{workflow_id}")
-
-        for task_id in workflow_tasks:
-            task_meta = await self._redis.hgetall(f"task_meta:{task_id}")
-
-            if not task_meta or task_meta.get("status") != TaskStatusEnum.PENDING:
-                continue
-
-            # Check if all dependencies are completed
-            dependencies = await self._redis.smembers(f"task_deps:{task_id}")
-            if dependencies:
-                dependency_statuses = []
-                for dep_id in dependencies:
-                    dep_status = await self._redis.hget(f"task_meta:{dep_id}", "status")
-                    dependency_statuses.append(dep_status == TaskStatusEnum.COMPLETED)
-
-                if not all(dependency_statuses):
-                    continue  # Dependencies not met
-
-            # Check capability requirements
-            required_caps = task_meta.get("required_capabilities", "").split(",")
-            required_caps = [cap.strip() for cap in required_caps if cap.strip()]
-
-            if required_caps and not any(cap in agent_capabilities for cap in required_caps):
-                continue  # Agent doesn't have required capabilities
-
-            ready_tasks.append(
-                {
-                    "task_id": task_id,
-                    "service_id": task_meta.get("service_id"),
-                    "task_type": task_meta.get("task_type"),
-                    "priority": task_meta.get("priority"),
-                    "estimated_duration_minutes": int(
-                        task_meta.get("estimated_duration_minutes", 60)
-                    ),
-                    "required_capabilities": required_caps,
-                }
-            )
-
-        # Sort by priority (high -> medium -> low)
-        priority_order = {"high": 0, "medium": 1, "low": 2}
-        ready_tasks.sort(key=lambda t: priority_order.get(t["priority"], 1))
-
-        return ready_tasks
-
-    async def claim_task(self, task_id: str, agent_id: str) -> bool:
-        """Atomically claim a task for an agent."""
-
-        # Use Redis transaction to ensure atomicity
-        async with self._redis.pipeline(transaction=True) as pipe:
-            while True:
-                try:
-                    # Watch the task status
-                    await pipe.watch(f"task_meta:{task_id}")
-
-                    # Check current status
-                    current_status = await self._redis.hget(f"task_meta:{task_id}", "status")
-
-                    if current_status != TaskStatusEnum.PENDING:
-                        await pipe.unwatch()
-                        return False  # Task already claimed or completed
-
-                    # Start transaction
-                    pipe.multi()
-
-                    # Update task status and assign agent
-                    await pipe.hset(
-                        f"task_meta:{task_id}",
-                        mapping={
-                            "status": TaskStatusEnum.IN_PROGRESS,
-                            "assigned_agent": agent_id,
-                            "claimed_at": datetime.now(UTC).isoformat(),
-                        },
-                    )
-
-                    # Execute transaction
-                    await pipe.execute()
-
-                    logger.info(f"Task {task_id} claimed by agent {agent_id}")
-                    return True
-
-                except Exception as e:
-                    logger.warning(f"Task claim conflict for {task_id}: {e}")
-                    continue  # Retry the transaction
-
-    async def update_task_progress(
-        self,
-        task_id: str,
-        agent_id: str,
-        progress: float,
-        status: TaskStatusEnum | None = None,
-        message: str | None = None,
-    ) -> bool:
-        """Update task progress and status."""
-
-        # Verify agent owns the task
-        assigned_agent = await self._redis.hget(f"task_meta:{task_id}", "assigned_agent")
-        if assigned_agent != agent_id:
-            logger.warning(
-                f"Agent {agent_id} attempted to update task {task_id} owned by {assigned_agent}"
-            )
-            return False
-
-        update_data = {"progress": str(progress), "last_updated": datetime.now(UTC).isoformat()}
-
-        if status:
-            update_data["status"] = status
-
-        if message:
-            update_data["last_message"] = message
-
-        await self._redis.hset(f"task_meta:{task_id}", mapping=update_data)
-
-        # Store progress history
-        progress_entry = {
-            "timestamp": datetime.now(UTC).isoformat(),
-            "progress": progress,
-            "status": status or TaskStatusEnum.IN_PROGRESS,
-            "message": message or "",
-            "agent_id": agent_id,
-        }
-
-        await self._redis.lpush(f"task_progress:{task_id}", str(progress_entry))
-
-        # Keep only last 50 progress entries
-        await self._redis.ltrim(f"task_progress:{task_id}", 0, 49)
-
-        logger.debug(f"Updated progress for task {task_id}: {progress:.2f}")
-        return True
-
-    async def complete_task(self, task_id: str, agent_id: str, result: dict[str, Any]) -> bool:
-        """Mark task as completed and unlock dependent tasks."""
-
-        # Verify agent owns the task
-        assigned_agent = await self._redis.hget(f"task_meta:{task_id}", "assigned_agent")
-        if assigned_agent != agent_id:
-            logger.warning(
-                f"Agent {agent_id} attempted to complete task {task_id} owned by {assigned_agent}"
-            )
-            return False
-
-        # Update task status
-        await self._redis.hset(
-            f"task_meta:{task_id}",
-            mapping={
-                "status": TaskStatusEnum.COMPLETED,
-                "completed_at": datetime.now(UTC).isoformat(),
-                "progress": "1.0",
-                "result": str(result),
-            },
-        )
-
-        # Update workflow progress
-        workflow_id = await self._redis.hget(f"task_meta:{task_id}", "workflow_id")
-        if workflow_id:
-            await self._redis.hincrby(f"workflow:{workflow_id}", "completed_tasks", 1)
-
-            # Check if workflow is complete
-            workflow_info = await self._redis.hgetall(f"workflow:{workflow_id}")
-            total_tasks = int(workflow_info.get("total_tasks", 0))
-            completed_tasks = int(workflow_info.get("completed_tasks", 0))
-
-            if completed_tasks >= total_tasks:
-                await self._redis.hset(f"workflow:{workflow_id}", "status", "completed")
-                await self._redis.hset(
-                    f"workflow:{workflow_id}", "completed_at", datetime.now(UTC).isoformat()
-                )
-                logger.info(f"Workflow {workflow_id} completed")
-
-        logger.info(f"Task {task_id} completed by agent {agent_id}")
-        return True
-
-    async def fail_task(self, task_id: str, agent_id: str, error_message: str) -> bool:
-        """Mark task as failed and handle reassignment."""
-
-        # Verify agent owns the task
-        assigned_agent = await self._redis.hget(f"task_meta:{task_id}", "assigned_agent")
-        if assigned_agent != agent_id:
-            logger.warning(
-                f"Agent {agent_id} attempted to fail task {task_id} owned by {assigned_agent}"
-            )
-            return False
-
-        # Get current failure count
-        failure_count = await self._redis.hget(f"task_meta:{task_id}", "failure_count") or "0"
-        failure_count = int(failure_count) + 1
-
-        max_retries = 3  # Configurable
-
-        if failure_count >= max_retries:
-            # Mark as permanently failed
-            await self._redis.hset(
-                f"task_meta:{task_id}",
-                mapping={
-                    "status": TaskStatusEnum.FAILED,
-                    "failed_at": datetime.now(UTC).isoformat(),
-                    "failure_count": str(failure_count),
-                    "error_message": error_message,
-                },
-            )
-            logger.error(
-                f"Task {task_id} permanently failed after {failure_count} attempts: {error_message}"
-            )
-        else:
-            # Reset for retry
-            await self._redis.hset(
-                f"task_meta:{task_id}",
-                mapping={
-                    "status": TaskStatusEnum.PENDING,
-                    "assigned_agent": "",
-                    "failure_count": str(failure_count),
-                    "last_error": error_message,
-                    "retry_after": (datetime.now(UTC) + timedelta(minutes=5)).isoformat(),
-                },
-            )
-            logger.warning(
-                f"Task {task_id} failed (attempt {failure_count}), will retry: {error_message}"
-            )
-
-        return True
-
-    async def get_workflow_status(self, workflow_id: str) -> dict[str, Any]:
-        """Get comprehensive workflow status."""
-
-        workflow_info = await self._redis.hgetall(f"workflow:{workflow_id}")
-        if not workflow_info:
-            return {}
-
-        # Get task statuses
-        workflow_tasks = await self._redis.smembers(f"workflow_tasks:{workflow_id}")
-        task_statuses = {}
-
-        for task_id in workflow_tasks:
-            task_meta = await self._redis.hgetall(f"task_meta:{task_id}")
-            task_statuses[task_id] = {
-                "status": task_meta.get("status"),
-                "progress": float(task_meta.get("progress", 0)),
-                "assigned_agent": task_meta.get("assigned_agent"),
-                "service_id": task_meta.get("service_id"),
-                "task_type": task_meta.get("task_type"),
-            }
-
-        # Calculate overall progress
-        total_progress = sum(task["progress"] for task in task_statuses.values())
-        overall_progress = total_progress / len(task_statuses) if task_statuses else 0
-
-        return {
-            "workflow_id": workflow_id,
-            "session_id": workflow_info.get("session_id"),
-            "pattern": workflow_info.get("pattern"),
-            "status": workflow_info.get("status"),
-            "total_tasks": int(workflow_info.get("total_tasks", 0)),
-            "completed_tasks": int(workflow_info.get("completed_tasks", 0)),
-            "overall_progress": overall_progress,
-            "created_at": workflow_info.get("created_at"),
-            "completed_at": workflow_info.get("completed_at"),
-            "task_statuses": task_statuses,
-        }
-
-    async def get_agent_workload(self, agent_id: str) -> dict[str, Any]:
-        """Get current workload for an agent."""
-
-        # Find all tasks assigned to this agent
-        task_keys = await self._redis.keys("task_meta:*")
-        assigned_tasks = []
-
-        for task_key in task_keys:
-            task_meta = await self._redis.hgetall(task_key)
-            if task_meta.get("assigned_agent") == agent_id:
-                task_id = task_key.decode().split(":")[-1]
-                assigned_tasks.append(
-                    {
-                        "task_id": task_id,
-                        "status": task_meta.get("status"),
-                        "progress": float(task_meta.get("progress", 0)),
-                        "service_id": task_meta.get("service_id"),
-                        "task_type": task_meta.get("task_type"),
-                    }
-                )
-
-        # Calculate workload metrics
-        active_tasks = [t for t in assigned_tasks if t["status"] == TaskStatusEnum.IN_PROGRESS]
-        completed_tasks = [t for t in assigned_tasks if t["status"] == TaskStatusEnum.COMPLETED]
-        failed_tasks = [t for t in assigned_tasks if t["status"] == TaskStatusEnum.FAILED]
-
-        return {
-            "agent_id": agent_id,
-            "total_assigned": len(assigned_tasks),
-            "active_tasks": len(active_tasks),
-            "completed_tasks": len(completed_tasks),
-            "failed_tasks": len(failed_tasks),
-            "success_rate": len(completed_tasks) / len(assigned_tasks) if assigned_tasks else 0,
-            "current_load": len(active_tasks),
-            "task_details": assigned_tasks,
-        }
-
-    async def create_progress_dashboard(self, session_id: str) -> dict[str, Any]:
-        """Create real-time progress dashboard for all workflows in a session."""
-
-        # Find all workflows for this session
-        workflow_keys = await self._redis.keys("workflow:*")
-        session_workflows = []
-
-        for workflow_key in workflow_keys:
-            workflow_info = await self._redis.hgetall(workflow_key)
-            if workflow_info.get("session_id") == session_id:
-                workflow_id = workflow_key.decode().split(":")[-1]
-                workflow_status = await self.get_workflow_status(workflow_id)
-                session_workflows.append(workflow_status)
-
-        # Get all active agents
-        agent_keys = await self._redis.keys("agent:*")
-        agent_workloads = []
-
-        for agent_key in agent_keys:
-            agent_id = agent_key.decode().split(":")[-1]
-            workload = await self.get_agent_workload(agent_id)
-            if workload["total_assigned"] > 0:
-                agent_workloads.append(workload)
-
-        # Calculate session-level metrics
-        total_tasks = sum(w["total_tasks"] for w in session_workflows)
-        completed_tasks = sum(w["completed_tasks"] for w in session_workflows)
-        overall_progress = completed_tasks / total_tasks if total_tasks > 0 else 0
-
-        return {
-            "session_id": session_id,
-            "timestamp": datetime.now(UTC).isoformat(),
-            "overall_progress": overall_progress,
-            "total_workflows": len(session_workflows),
-            "total_tasks": total_tasks,
-            "completed_tasks": completed_tasks,
-            "active_agents": len([a for a in agent_workloads if a["active_tasks"] > 0]),
-            "workflows": session_workflows,
-            "agent_workloads": agent_workloads,
-        }
-
-
-class CoordinationPatternManager:
-    """Manages different coordination patterns for multi-agent workflows."""
-
-    @staticmethod
-    def create_hierarchical_workflow(
-        services: list[dict[str, Any]], supervisor_agent: str
-    ) -> dict[str, Any]:
-        """Create hierarchical coordination workflow."""
-
-        tasks = []
-
-        # Supervisor task to coordinate others
-        tasks.append(
-            {
-                "task_id": "supervisor_coordination",
-                "task_type": "supervisor_coordination",
-                "required_capabilities": ["supervision", "coordination"],
-                "dependencies": [],
-                "priority": "high",
-                "estimated_duration_minutes": 30,
-            }
-        )
-
-        # Specialized agent tasks
-        for service in services:
-            service_id = service["service_id"]
-
-            # BRD analysis task
-            tasks.append(
-                {
-                    "task_id": f"brd_analysis_{service_id}",
-                    "service_id": service_id,
-                    "task_type": "brd_analysis",
-                    "required_capabilities": ["brd_analysis"],
-                    "dependencies": ["supervisor_coordination"],
-                    "priority": service.get("priority", "medium"),
-                    "estimated_duration_minutes": 45,
-                }
-            )
-
-            # Diagram generation task (depends on BRD)
-            tasks.append(
-                {
-                    "task_id": f"diagram_generation_{service_id}",
-                    "service_id": service_id,
-                    "task_type": "diagram_generation",
-                    "required_capabilities": ["diagram_generation"],
-                    "dependencies": [f"brd_analysis_{service_id}"],
-                    "priority": service.get("priority", "medium"),
-                    "estimated_duration_minutes": 60,
-                }
-            )
-
-        return {
-            "pattern": CoordinationPatternEnum.HIERARCHICAL,
-            "tasks": tasks,
-            "supervisor_agent": supervisor_agent,
-        }
-
-    @staticmethod
-    def create_parallel_workflow(services: list[dict[str, Any]]) -> dict[str, Any]:
-        """Create parallel coordination workflow."""
-
-        tasks = []
-
-        for service in services:
-            service_id = service["service_id"]
-
-            # Independent tasks that can run in parallel
-            for task_type in ["component_diagram", "sequence_diagrams", "use_case_diagram"]:
-                tasks.append(
-                    {
-                        "task_id": f"{task_type}_{service_id}",
-                        "service_id": service_id,
-                        "task_type": task_type,
-                        "required_capabilities": ["diagram_generation"],
-                        "dependencies": [],  # No dependencies - fully parallel
-                        "priority": service.get("priority", "medium"),
-                        "estimated_duration_minutes": 30,
-                    }
-                )
-
-        return {"pattern": CoordinationPatternEnum.PARALLEL, "tasks": tasks}
-
-    @staticmethod
-    def create_sequential_workflow(services: list[dict[str, Any]]) -> dict[str, Any]:
-        """Create sequential coordination workflow."""
-
-        tasks = []
-        previous_task = None
-
-        for service in services:
-            service_id = service["service_id"]
-
-            # Sequential tasks with dependencies
-            task_sequence = [
-                "brd_update",
-                "architecture_doc",
-                "service_spec",
-                "diagram_generation",
-                "validation",
-            ]
-
-            for task_type in task_sequence:
-                task_id = f"{task_type}_{service_id}"
-
-                dependencies = [previous_task] if previous_task else []
-
-                tasks.append(
-                    {
-                        "task_id": task_id,
-                        "service_id": service_id,
-                        "task_type": task_type,
-                        "required_capabilities": [task_type.replace("_", "")],
-                        "dependencies": dependencies,
-                        "priority": service.get("priority", "medium"),
-                        "estimated_duration_minutes": 45,
-                    }
-                )
-
-                previous_task = task_id
-
-        return {"pattern": CoordinationPatternEnum.SEQUENTIAL, "tasks": tasks}