nas-tools/dashboard/backend/services/agent_executor.py

"""
Agent Executor Service - Executes agent tasks and manages their lifecycle
"""
import asyncio
import logging
from typing import Dict, Any, Optional
from datetime import datetime
from db import opc_db
from services.agents.base_agent import BaseAgent
from services import email_service
import httpx
import os

logger = logging.getLogger(__name__)

# Telegram notification settings
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID")


class AgentExecutor:
    """Manages agent execution lifecycle"""

    def __init__(self):
        self.agents: Dict[str, BaseAgent] = {}
        self.running = False

    async def initialize(self):
        """Load agents from database"""
        agents_data = await opc_db.get_agents(enabled_only=True)

        for agent_data in agents_data:
            agent = BaseAgent(
                agent_id=agent_data["id"],
                name=agent_data["name"],
                role=agent_data["role"],
                system_prompt=agent_data["system_prompt"],
                config=agent_data["config"]
            )
            self.agents[agent_data["id"]] = agent
            logger.info(f"Loaded agent: {agent.name} ({agent.agent_id})")

    async def start(self):
        """Start the executor service"""
        self.running = True
        logger.info("Agent executor service started")

        while self.running:
            try:
                await self._process_pending_executions()
            except Exception as e:
                logger.error(f"Error in executor loop: {e}")

            await asyncio.sleep(5)  # Check every 5 seconds

    async def stop(self):
        """Stop the executor service"""
        self.running = False
        logger.info("Agent executor service stopped")

    async def _process_pending_executions(self):
        """Process pending agent executions"""
        executions = await opc_db.get_agent_executions(status="pending", limit=10)

        for execution in executions:
            try:
                await self._execute_agent(execution)
            except Exception as e:
                logger.error(f"Failed to execute agent {execution['agent_id']}: {e}")
                await self._mark_execution_failed(execution["id"], str(e))

    async def _execute_agent(self, execution: Dict[str, Any]):
        """Execute a single agent task"""
        agent_id = execution["agent_id"]
        task_id = execution["task_id"]

        # Get agent
        agent = self.agents.get(agent_id)
        if not agent:
            raise ValueError(f"Agent {agent_id} not found")

        # Mark as running
        await self._update_execution_status(execution["id"], "running")

        # Get task and context
        tasks = await opc_db.get_tasks(limit=1000)
        task = next((t for t in tasks if t["id"] == task_id), None)
        if not task:
            raise ValueError(f"Task {task_id} not found")

        # Build context
        context = await self._build_context(task)

        # Execute agent
        logger.info(f"Executing agent {agent.name} on task {task['title']}")
        result = await agent.execute(task, context)

        # Check if requires approval
        if execution.get("requires_approval") or result.get("requires_approval"):
            # Mark as pending approval
            await self._mark_pending_approval(execution["id"], result)
            await self._send_approval_notification(agent, task, result)
        else:
            # Execute actions
            await self._execute_actions(task, result["actions"])

            # Mark as completed
            await self._mark_execution_completed(execution["id"], result)

            # Send completion notification
            await self._send_completion_notification(agent, task, result)

    async def _build_context(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """Build context for agent execution"""
        context = {
            "task": task,
            "timestamp": datetime.utcnow().isoformat()
        }

        # Get related tasks from same project
        if task.get("project_id"):
            related_tasks = await opc_db.get_tasks(project_id=task["project_id"], limit=20)
            context["related_tasks"] = related_tasks

        # Get task history
        history = await opc_db.get_task_history(task["id"])
        context["history"] = history[:10]  # Last 10 events

        return context

    async def _execute_actions(self, task: Dict[str, Any], actions: list):
        """Execute proposed actions"""
        for action in actions:
            action_type = action.get("type")

            try:
                if action_type == "create_subtask":
                    await self._action_create_subtask(task, action)
                elif action_type == "update_task_status":
                    await self._action_update_status(task, action)
                elif action_type == "send_notification":
                    await self._action_send_notification(action)
                elif action_type == "add_comment":
                    await self._action_add_comment(task, action)
                else:
                    logger.warning(f"Unknown action type: {action_type}")
            except Exception as e:
                logger.error(f"Failed to execute action {action_type}: {e}")

    async def _action_create_subtask(self, parent_task: Dict[str, Any], action: Dict[str, Any]):
        """Create a subtask"""
        await opc_db.create_task(
            title=action.get("title", "Subtask"),
            description=action.get("description", ""),
            status="parking_lot",
            priority=parent_task.get("priority", "medium"),
            project_id=parent_task.get("project_id"),
            tags=parent_task.get("tags", []),
            actor="agent"
        )
        logger.info(f"Created subtask: {action.get('title')}")

    async def _action_update_status(self, task: Dict[str, Any], action: Dict[str, Any]):
        """Update task status"""
        new_status = action.get("status")
        if new_status:
            await opc_db.update_task(
                task_id=task["id"],
                updates={"status": new_status},
                actor="agent"
            )
            logger.info(f"Updated task {task['id']} status to {new_status}")

    async def _action_send_notification(self, action: Dict[str, Any]):
        """Send notification"""
        message = action.get("message", "")
        if message and TELEGRAM_BOT_TOKEN and TELEGRAM_CHAT_ID:
            await self._send_telegram(message)

    async def _action_add_comment(self, task: Dict[str, Any], action: Dict[str, Any]):
        """Add comment to task (stored in history)"""
        comment = action.get("comment", "")
        if comment:
            # Store as task history
            pool = await opc_db.get_pool()
            async with pool.acquire() as conn:
                await conn.execute("""
                    INSERT INTO task_history (task_id, action, actor, actor_type, new_value)
                    VALUES ($1, $2, $3, $4, $5)
                """, task["id"], "comment", "agent", "agent", comment)

    async def _update_execution_status(self, execution_id: int, status: str):
        """Update execution status"""
        pool = await opc_db.get_pool()
        async with pool.acquire() as conn:
            await conn.execute("""
                UPDATE agent_executions
                SET status = $1, started_at = CURRENT_TIMESTAMP
                WHERE id = $2
            """, status, execution_id)

    async def _mark_execution_completed(self, execution_id: int, result: Dict[str, Any]):
        """Mark execution as completed"""
        pool = await opc_db.get_pool()
        async with pool.acquire() as conn:
            await conn.execute("""
                UPDATE agent_executions
                SET status = $1,
                    output_result = $2,
                    actions_proposed = $3,
                    actions_executed = $3,
                    completed_at = CURRENT_TIMESTAMP
                WHERE id = $4
            """, "completed",
                opc_db.json.dumps(result),
                opc_db.json.dumps(result.get("actions", [])),
                execution_id)

    async def _mark_execution_failed(self, execution_id: int, error: str):
        """Mark execution as failed"""
        pool = await opc_db.get_pool()
        async with pool.acquire() as conn:
            await conn.execute("""
                UPDATE agent_executions
                SET status = $1, error_message = $2, completed_at = CURRENT_TIMESTAMP
                WHERE id = $3
            """, "failed", error, execution_id)

    async def _mark_pending_approval(self, execution_id: int, result: Dict[str, Any]):
        """Mark execution as pending approval"""
        pool = await opc_db.get_pool()
        async with pool.acquire() as conn:
            await conn.execute("""
                UPDATE agent_executions
                SET status = $1,
                    output_result = $2,
                    actions_proposed = $3,
                    requires_approval = TRUE
                WHERE id = $4
            """, "pending_approval",
                opc_db.json.dumps(result),
                opc_db.json.dumps(result.get("actions", [])),
                execution_id)

    async def _send_telegram(self, message: str):
        """Send Telegram notification"""
        try:
            proxy_url = os.environ.get("TELEGRAM_PROXY")
            async with httpx.AsyncClient(proxy=proxy_url, timeout=10.0) as client:
                await client.post(
                    f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
                    data={
                        "chat_id": TELEGRAM_CHAT_ID,
                        "text": message,
                        "parse_mode": "Markdown"
                    }
                )
        except Exception as e:
            logger.error(f"Failed to send Telegram notification: {e}")

    async def _send_approval_notification(self, agent: BaseAgent, task: Dict[str, Any], result: Dict[str, Any]):
        """Send notification for approval request"""
        message = f"""🤖 *Agent Approval Required*

Agent: {agent.name}
Task: {task['title']}

Reasoning: {result.get('reasoning', 'No reasoning provided')}

Proposed Actions:
"""
        for i, action in enumerate(result.get("actions", []), 1):
            message += f"{i}. {action.get('type')}: {action.get('title', action.get('message', 'N/A'))}\n"

        message += "\nPlease review and approve in the OPC dashboard."

        # Send Telegram
        await self._send_telegram(message)

        # Send Email
        await email_service.send_agent_approval_email(
            agent_name=agent.name,
            task=task,
            reasoning=result.get('reasoning', 'No reasoning provided'),
            actions=result.get('actions', [])
        )

    async def _send_completion_notification(self, agent: BaseAgent, task: Dict[str, Any], result: Dict[str, Any]):
        """Send notification for completed execution"""
        message = f"""✅ *Agent Task Completed*

Agent: {agent.name}
Task: {task['title']}

Actions Executed: {len(result.get('actions', []))}
"""
        # Send Telegram
        await self._send_telegram(message)

        # Send Email
        await email_service.send_agent_completion_email(
            agent_name=agent.name,
            task=task,
            actions_count=len(result.get('actions', []))
        )


# Global executor instance
_executor: Optional[AgentExecutor] = None


async def get_executor() -> AgentExecutor:
    """Get or create executor instance"""
    global _executor
    if _executor is None:
        _executor = AgentExecutor()
        await _executor.initialize()
    return _executor


async def start_executor():
    """Start the executor service - initializes and starts the background task"""
    print("start_executor() called")
    try:
        executor = await get_executor()
        print(f"Executor obtained: {executor is not None}, running: {executor.running}")
        # Don't await start() - it's an infinite loop!
        # Just call it to set running=True and start the loop
        # The task will be managed by the caller
        asyncio.create_task(executor.start())
        print("Executor background task created")
    except Exception as e:
        print(f"ERROR in start_executor: {e}")
        import traceback
        traceback.print_exc()
        raise