agent-governance/teams/framework/team.py

#!/usr/bin/env python3
"""
Hierarchical Agent Team Framework

Provides team-based orchestration with:
- Team Lead coordination
- Specialized sub-teams (Research, Implementation, Review)
- Task delegation and result aggregation
- Consensus mechanisms across teams
"""

import asyncio
import json
import redis
import sqlite3
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Dict, List, Any, Optional, Callable
import hashlib


REDIS_HOST = "127.0.0.1"
REDIS_PORT = 6379
REDIS_PASSWORD = "governance2026"
LEDGER_PATH = Path("/opt/agent-governance/ledger/governance.db")


class TeamRole(Enum):
    """Roles within a hierarchical team"""
    LEAD = "lead"
    RESEARCH = "research"
    IMPLEMENTATION = "implementation"
    REVIEW = "review"
    SPECIALIST = "specialist"


class TaskStatus(Enum):
    """Status of a delegated task"""
    PENDING = "pending"
    ASSIGNED = "assigned"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    BLOCKED = "blocked"


@dataclass
class TeamMember:
    """Represents a member of a team"""
    agent_id: str
    role: TeamRole
    tier: int
    capabilities: List[str]
    status: str = "idle"
    current_task: Optional[str] = None

    def to_dict(self) -> Dict:
        return {
            "agent_id": self.agent_id,
            "role": self.role.value,
            "tier": self.tier,
            "capabilities": self.capabilities,
            "status": self.status,
            "current_task": self.current_task
        }


@dataclass
class DelegatedTask:
    """A task delegated to a team member"""
    task_id: str
    description: str
    assigned_to: Optional[str] = None
    assigned_team: Optional[TeamRole] = None
    status: TaskStatus = TaskStatus.PENDING
    priority: int = 5
    dependencies: List[str] = field(default_factory=list)
    result: Optional[Dict] = None
    created_at: datetime = field(default_factory=datetime.utcnow)
    completed_at: Optional[datetime] = None

    def to_dict(self) -> Dict:
        return {
            "task_id": self.task_id,
            "description": self.description,
            "assigned_to": self.assigned_to,
            "assigned_team": self.assigned_team.value if self.assigned_team else None,
            "status": self.status.value,
            "priority": self.priority,
            "dependencies": self.dependencies,
            "result": self.result,
            "created_at": self.created_at.isoformat(),
            "completed_at": self.completed_at.isoformat() if self.completed_at else None
        }


@dataclass
class TeamResult:
    """Aggregated result from a team"""
    team_role: TeamRole
    tasks_completed: int
    tasks_failed: int
    findings: List[Dict]
    recommendations: List[str]
    confidence: float
    duration_seconds: float


class BaseTeam(ABC):
    """Base class for specialized teams"""

    def __init__(self, team_id: str, role: TeamRole, max_members: int = 3):
        self.team_id = team_id
        self.role = role
        self.max_members = max_members
        self.members: List[TeamMember] = []
        self.tasks: Dict[str, DelegatedTask] = {}
        self.results: List[Dict] = []

    def add_member(self, member: TeamMember) -> bool:
        """Add a member to the team"""
        if len(self.members) >= self.max_members:
            return False
        if member.role != self.role and member.role != TeamRole.SPECIALIST:
            return False
        self.members.append(member)
        return True

    def assign_task(self, task: DelegatedTask) -> Optional[str]:
        """Assign a task to an available member"""
        for member in self.members:
            if member.status == "idle":
                member.status = "working"
                member.current_task = task.task_id
                task.assigned_to = member.agent_id
                task.assigned_team = self.role
                task.status = TaskStatus.ASSIGNED
                self.tasks[task.task_id] = task
                return member.agent_id
        return None

    def complete_task(self, task_id: str, result: Dict) -> bool:
        """Mark a task as completed"""
        if task_id not in self.tasks:
            return False
        task = self.tasks[task_id]
        task.status = TaskStatus.COMPLETED
        task.result = result
        task.completed_at = datetime.utcnow()

        # Free up the member
        for member in self.members:
            if member.current_task == task_id:
                member.status = "idle"
                member.current_task = None
                break

        self.results.append(result)
        return True

    @abstractmethod
    async def process_objective(self, objective: Dict) -> TeamResult:
        """Process an objective assigned to this team"""
        pass

    def get_status(self) -> Dict:
        """Get team status"""
        return {
            "team_id": self.team_id,
            "role": self.role.value,
            "members": [m.to_dict() for m in self.members],
            "active_tasks": len([t for t in self.tasks.values()
                                if t.status == TaskStatus.IN_PROGRESS]),
            "completed_tasks": len([t for t in self.tasks.values()
                                   if t.status == TaskStatus.COMPLETED])
        }


class ResearchTeam(BaseTeam):
    """Team specialized in research and analysis"""

    def __init__(self, team_id: str):
        super().__init__(team_id, TeamRole.RESEARCH, max_members=3)

    async def process_objective(self, objective: Dict) -> TeamResult:
        """Research and analyze the objective"""
        start_time = datetime.utcnow()
        findings = []
        recommendations = []

        # Simulate research tasks
        research_tasks = [
            "Analyze requirements",
            "Review existing solutions",
            "Identify constraints",
            "Assess risks"
        ]

        completed = 0
        failed = 0

        for task_desc in research_tasks:
            task = DelegatedTask(
                task_id=f"{self.team_id}-research-{len(self.tasks)}",
                description=task_desc,
                priority=objective.get("priority", 5)
            )

            agent_id = self.assign_task(task)
            if agent_id:
                # Simulate task execution
                await asyncio.sleep(0.1)

                result = {
                    "task": task_desc,
                    "findings": [f"Finding from {task_desc}"],
                    "confidence": 0.85
                }
                self.complete_task(task.task_id, result)
                findings.extend(result["findings"])
                completed += 1
            else:
                failed += 1

        recommendations = [
            f"Based on research: {objective.get('description', 'objective')[:50]}...",
            "Recommend proceeding with implementation",
            "Consider edge cases identified in analysis"
        ]

        duration = (datetime.utcnow() - start_time).total_seconds()

        return TeamResult(
            team_role=self.role,
            tasks_completed=completed,
            tasks_failed=failed,
            findings=findings,
            recommendations=recommendations,
            confidence=0.85 if failed == 0 else 0.6,
            duration_seconds=duration
        )


class ImplementationTeam(BaseTeam):
    """Team specialized in implementation"""

    def __init__(self, team_id: str):
        super().__init__(team_id, TeamRole.IMPLEMENTATION, max_members=2)

    async def process_objective(self, objective: Dict) -> TeamResult:
        """Implement the solution"""
        start_time = datetime.utcnow()
        findings = []

        # Implementation steps based on objective
        steps = objective.get("steps", [
            "Setup environment",
            "Implement core logic",
            "Add error handling",
            "Create tests"
        ])

        completed = 0
        failed = 0

        for step in steps:
            task = DelegatedTask(
                task_id=f"{self.team_id}-impl-{len(self.tasks)}",
                description=step,
                priority=objective.get("priority", 5)
            )

            agent_id = self.assign_task(task)
            if agent_id:
                await asyncio.sleep(0.1)

                result = {
                    "step": step,
                    "status": "completed",
                    "artifacts": [f"artifact_{step.replace(' ', '_')}"]
                }
                self.complete_task(task.task_id, result)
                findings.append({"step": step, "result": "success"})
                completed += 1
            else:
                failed += 1
                findings.append({"step": step, "result": "skipped"})

        duration = (datetime.utcnow() - start_time).total_seconds()

        return TeamResult(
            team_role=self.role,
            tasks_completed=completed,
            tasks_failed=failed,
            findings=findings,
            recommendations=["Implementation complete, ready for review"],
            confidence=0.9 if failed == 0 else 0.5,
            duration_seconds=duration
        )


class ReviewTeam(BaseTeam):
    """Team specialized in review and validation"""

    def __init__(self, team_id: str):
        super().__init__(team_id, TeamRole.REVIEW, max_members=2)

    async def process_objective(self, objective: Dict) -> TeamResult:
        """Review and validate the implementation"""
        start_time = datetime.utcnow()
        findings = []

        review_aspects = [
            "Code quality review",
            "Security review",
            "Governance compliance check",
            "Performance assessment"
        ]

        completed = 0
        failed = 0
        issues_found = []

        for aspect in review_aspects:
            task = DelegatedTask(
                task_id=f"{self.team_id}-review-{len(self.tasks)}",
                description=aspect,
                priority=objective.get("priority", 5)
            )

            agent_id = self.assign_task(task)
            if agent_id:
                await asyncio.sleep(0.1)

                # Simulate review findings
                result = {
                    "aspect": aspect,
                    "passed": True,
                    "issues": [],
                    "score": 0.9
                }
                self.complete_task(task.task_id, result)
                findings.append(result)
                completed += 1
            else:
                failed += 1

        # Calculate overall score
        scores = [f.get("score", 0) for f in findings if "score" in f]
        avg_score = sum(scores) / len(scores) if scores else 0

        duration = (datetime.utcnow() - start_time).total_seconds()

        return TeamResult(
            team_role=self.role,
            tasks_completed=completed,
            tasks_failed=failed,
            findings=findings,
            recommendations=[
                f"Overall review score: {avg_score:.2f}",
                "Approved for deployment" if avg_score > 0.7 else "Requires fixes"
            ],
            confidence=avg_score,
            duration_seconds=duration
        )


class TeamLead:
    """
    Team Lead that coordinates multiple specialized teams.

    Responsibilities:
    - Break down objectives into team tasks
    - Delegate to appropriate teams
    - Aggregate and synthesize results
    - Make final decisions
    """

    def __init__(self, lead_id: str, objective: Dict):
        self.lead_id = lead_id
        self.objective = objective
        self.teams: Dict[TeamRole, BaseTeam] = {}
        self.team_results: Dict[TeamRole, TeamResult] = {}
        self.redis = redis.Redis(
            host=REDIS_HOST,
            port=REDIS_PORT,
            password=REDIS_PASSWORD,
            decode_responses=True
        )
        self.start_time = None
        self.end_time = None

    def spawn_team(self, role: TeamRole) -> BaseTeam:
        """Spawn a specialized team"""
        team_id = f"{self.lead_id}-{role.value}"

        if role == TeamRole.RESEARCH:
            team = ResearchTeam(team_id)
        elif role == TeamRole.IMPLEMENTATION:
            team = ImplementationTeam(team_id)
        elif role == TeamRole.REVIEW:
            team = ReviewTeam(team_id)
        else:
            raise ValueError(f"Unknown team role: {role}")

        # Add default members
        for i in range(team.max_members):
            member = TeamMember(
                agent_id=f"{team_id}-agent-{i}",
                role=role,
                tier=1,
                capabilities=self._get_capabilities_for_role(role)
            )
            team.add_member(member)

        self.teams[role] = team
        return team

    def _get_capabilities_for_role(self, role: TeamRole) -> List[str]:
        """Get default capabilities for a role"""
        capabilities = {
            TeamRole.RESEARCH: ["analyze", "search", "summarize", "compare"],
            TeamRole.IMPLEMENTATION: ["code", "configure", "deploy", "test"],
            TeamRole.REVIEW: ["review", "validate", "audit", "approve"]
        }
        return capabilities.get(role, ["generic"])

    async def execute(self) -> Dict:
        """Execute the full team workflow"""
        self.start_time = datetime.utcnow()

        # Store state in DragonflyDB
        self._update_state("starting")

        # Phase 1: Research
        print(f"  [{self.lead_id}] Starting Research phase...")
        research_team = self.spawn_team(TeamRole.RESEARCH)
        research_result = await research_team.process_objective(self.objective)
        self.team_results[TeamRole.RESEARCH] = research_result
        self._update_state("research_complete")

        # Phase 2: Implementation (if research recommends proceeding)
        if research_result.confidence >= 0.6:
            print(f"  [{self.lead_id}] Starting Implementation phase...")
            impl_team = self.spawn_team(TeamRole.IMPLEMENTATION)

            # Enhance objective with research findings
            enhanced_objective = {
                **self.objective,
                "research_findings": research_result.findings,
                "recommendations": research_result.recommendations
            }

            impl_result = await impl_team.process_objective(enhanced_objective)
            self.team_results[TeamRole.IMPLEMENTATION] = impl_result
            self._update_state("implementation_complete")

            # Phase 3: Review
            if impl_result.confidence >= 0.5:
                print(f"  [{self.lead_id}] Starting Review phase...")
                review_team = self.spawn_team(TeamRole.REVIEW)

                review_objective = {
                    **enhanced_objective,
                    "implementation_results": impl_result.findings
                }

                review_result = await review_team.process_objective(review_objective)
                self.team_results[TeamRole.REVIEW] = review_result
                self._update_state("review_complete")

        self.end_time = datetime.utcnow()

        # Aggregate final result
        final_result = self._aggregate_results()
        self._update_state("completed")
        self._record_to_ledger(final_result)

        return final_result

    def _aggregate_results(self) -> Dict:
        """Aggregate results from all teams"""
        all_findings = []
        all_recommendations = []
        total_tasks = 0
        total_failed = 0

        for role, result in self.team_results.items():
            all_findings.extend([
                {"team": role.value, "finding": f}
                for f in result.findings
            ])
            all_recommendations.extend(result.recommendations)
            total_tasks += result.tasks_completed
            total_failed += result.tasks_failed

        # Calculate overall confidence
        confidences = [r.confidence for r in self.team_results.values()]
        overall_confidence = sum(confidences) / len(confidences) if confidences else 0

        # Make final decision
        success = (
            overall_confidence >= 0.7 and
            total_failed == 0 and
            TeamRole.REVIEW in self.team_results
        )

        duration = (self.end_time - self.start_time).total_seconds()

        return {
            "lead_id": self.lead_id,
            "objective": self.objective.get("description", "Unknown"),
            "success": success,
            "decision": "APPROVED" if success else "NEEDS_REVIEW",
            "overall_confidence": overall_confidence,
            "teams_involved": [r.value for r in self.team_results.keys()],
            "total_tasks_completed": total_tasks,
            "total_tasks_failed": total_failed,
            "findings_count": len(all_findings),
            "recommendations": all_recommendations[:5],  # Top 5
            "duration_seconds": duration,
            "completed_at": self.end_time.isoformat()
        }

    def _update_state(self, status: str):
        """Update state in DragonflyDB"""
        state = {
            "lead_id": self.lead_id,
            "status": status,
            "teams": json.dumps([r.value for r in self.teams.keys()]),
            "updated_at": datetime.utcnow().isoformat()
        }
        self.redis.hset(f"team:{self.lead_id}:state", mapping=state)

    def _record_to_ledger(self, result: Dict):
        """Record execution to ledger"""
        try:
            conn = sqlite3.connect(LEDGER_PATH)
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO agent_actions
                (timestamp, agent_id, agent_version, tier, action, decision,
                 confidence, success, session_id, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                datetime.utcnow().isoformat(),
                self.lead_id,
                "1.0.0",
                2,  # Team leads are tier 2
                "team_execution",
                result["decision"],
                result["overall_confidence"],
                1 if result["success"] else 0,
                self.lead_id,
                datetime.utcnow().isoformat()
            ))
            conn.commit()
            conn.close()
        except Exception as e:
            print(f"Warning: Failed to record to ledger: {e}")


async def run_team_task(objective: Dict) -> Dict:
    """
    Run a task using hierarchical team structure.

    Args:
        objective: Task objective with description, priority, etc.

    Returns:
        Aggregated result from all teams
    """
    # Generate lead ID
    timestamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
    lead_id = f"team-lead-{timestamp}"

    print(f"\n{'='*60}")
    print(f"HIERARCHICAL TEAM EXECUTION")
    print(f"{'='*60}")
    print(f"Lead: {lead_id}")
    print(f"Objective: {objective.get('description', 'Unknown')[:50]}...")
    print()

    lead = TeamLead(lead_id, objective)
    result = await lead.execute()

    print(f"\n{'='*60}")
    print(f"RESULT: {result['decision']}")
    print(f"Confidence: {result['overall_confidence']:.2f}")
    print(f"Teams: {', '.join(result['teams_involved'])}")
    print(f"Tasks: {result['total_tasks_completed']} completed, {result['total_tasks_failed']} failed")
    print(f"Duration: {result['duration_seconds']:.2f}s")
    print(f"{'='*60}\n")

    return result


if __name__ == "__main__":
    # Example usage
    objective = {
        "description": "Implement a new monitoring dashboard feature",
        "priority": 3,
        "constraints": ["Must be backward compatible", "No downtime"],
        "success_criteria": ["All tests pass", "Review approved"]
    }

    result = asyncio.run(run_team_task(objective))
    print(json.dumps(result, indent=2))