breakpilot-core/voice-service/services/enhanced_task_helpers.py

"""
Enhanced Task Orchestrator - Helper Mixin

Provides routing, quality checks, memory storage, message handling,
and session recovery logic for EnhancedTaskOrchestrator.

Separated from enhanced_task_orchestrator.py for the 500 LOC budget.
"""

import structlog
import asyncio
from typing import Optional, Dict, Any
from datetime import datetime

from models.task import Task, TaskState

from sessions.session_manager import AgentSession, SessionState
from sessions.heartbeat import HeartbeatClient
from orchestrator.message_bus import AgentMessage, MessagePriority
from orchestrator.task_router import RoutingStrategy

logger = structlog.get_logger(__name__)


class EnhancedTaskHelpersMixin:
    """
    Mixin providing internal helper methods for EnhancedTaskOrchestrator.

    This class should not be instantiated directly. It is mixed into
    EnhancedTaskOrchestrator to provide:
    - Agent routing (_route_to_agent, _needs_specialized_agent)
    - Quality checks (_run_quality_check, _needs_quality_check)
    - Memory storage (_store_task_result)
    - Message handling (_handle_agent_message)
    - Session recovery (recover_session, _recover_pending_tasks)
    - System prompt (_get_system_prompt)

    All attributes (message_bus, task_router, memory_store, heartbeat,
    session_manager, _voice_sessions, _heartbeat_clients, _tasks) are
    expected to be set by the main class __init__.
    """

    def _needs_specialized_agent(self, task: Task) -> bool:
        """Check if task needs routing to a specialized agent"""
        from models.task import TaskType

        # Tasks that benefit from specialized agents
        specialized_types = [
            TaskType.PARENT_LETTER,      # Could use grader for tone
            TaskType.FEEDBACK_SUGGEST,   # Quality judge for appropriateness
        ]

        return task.type in specialized_types

    def _needs_quality_check(self, task: Task) -> bool:
        """Check if task result needs quality validation"""
        from models.task import TaskType

        # Tasks that generate content should be checked
        content_types = [
            TaskType.PARENT_LETTER,
            TaskType.CLASS_MESSAGE,
            TaskType.FEEDBACK_SUGGEST,
            TaskType.WORKSHEET_GENERATE,
        ]

        return task.type in content_types

    async def _route_to_agent(
        self,
        task: Task,
        session: Optional[AgentSession]
    ) -> None:
        """Routes a task to a specialized agent"""
        # Determine target agent
        intent = f"task_{task.type.value}"
        routing_result = await self.task_router.route(
            intent=intent,
            context={"task": task.parameters},
            strategy=RoutingStrategy.LEAST_LOADED
        )

        if not routing_result.success:
            # Fall back to local processing
            logger.warning(
                "No agent available for task, using local processing",
                task_id=task.id[:8],
                reason=routing_result.reason
            )
            await super().process_task(task)
            return

        # Send to agent via message bus
        try:
            response = await self.message_bus.request(
                AgentMessage(
                    sender="voice-orchestrator",
                    receiver=routing_result.agent_id,
                    message_type=f"process_{task.type.value}",
                    payload={
                        "task_id": task.id,
                        "task_type": task.type.value,
                        "parameters": task.parameters,
                        "session_id": session.session_id if session else None
                    },
                    priority=MessagePriority.NORMAL
                ),
                timeout=30.0
            )

            task.result_ref = response.get("result", "")
            task.transition_to(TaskState.READY, "agent_processed")

        except asyncio.TimeoutError:
            logger.error(
                "Agent timeout, falling back to local",
                task_id=task.id[:8],
                agent=routing_result.agent_id
            )
            await super().process_task(task)

    async def _run_quality_check(
        self,
        task: Task,
        session: Optional[AgentSession]
    ) -> None:
        """Runs quality check on task result via quality judge"""
        try:
            response = await self.message_bus.request(
                AgentMessage(
                    sender="voice-orchestrator",
                    receiver="quality-judge",
                    message_type="evaluate_response",
                    payload={
                        "task_id": task.id,
                        "task_type": task.type.value,
                        "response": task.result_ref,
                        "context": task.parameters
                    },
                    priority=MessagePriority.NORMAL
                ),
                timeout=10.0
            )

            quality_score = response.get("composite_score", 0)

            if quality_score < 60:
                # Mark for review
                task.error_message = f"Quality check failed: {quality_score}"
                logger.warning(
                    "Task failed quality check",
                    task_id=task.id[:8],
                    score=quality_score
                )

        except asyncio.TimeoutError:
            # Quality check timeout is non-fatal
            logger.warning(
                "Quality check timeout",
                task_id=task.id[:8]
            )

    async def _store_task_result(self, task: Task) -> None:
        """Stores task result in memory for learning"""
        await self.memory_store.remember(
            key=f"task:{task.type.value}:{task.id}",
            value={
                "result": task.result_ref,
                "parameters": task.parameters,
                "completed_at": datetime.utcnow().isoformat()
            },
            agent_id="voice-orchestrator",
            ttl_days=30
        )

    async def _handle_agent_message(
        self,
        message: AgentMessage
    ) -> Optional[Dict[str, Any]]:
        """Handles incoming messages from other agents"""
        logger.debug(
            "Received agent message",
            sender=message.sender,
            type=message.message_type
        )

        if message.message_type == "task_status_update":
            # Handle task status updates
            task_id = message.payload.get("task_id")
            if task_id in self._tasks:
                task = self._tasks[task_id]
                new_state = message.payload.get("state")
                if new_state:
                    task.transition_to(TaskState(new_state), "agent_update")

        return None

    def _get_system_prompt(self) -> str:
        """Returns the system prompt for the voice assistant"""
        return """Du bist ein hilfreicher Assistent für Lehrer in der Breakpilot-App.

Deine Aufgaben:
- Hilf beim Erstellen von Arbeitsblättern
- Unterstütze bei der Korrektur
- Erstelle Elternbriefe und Klassennachrichten
- Dokumentiere Beobachtungen und Erinnerungen

Halte dich kurz und präzise. Nutze einfache, klare Sprache.
Bei Unklarheiten frage nach."""

    # Recovery methods

    async def recover_session(
        self,
        voice_session_id: str,
        session_id: str
    ) -> Optional[AgentSession]:
        """
        Recovers a session from checkpoint.

        Args:
            voice_session_id: The voice session ID
            session_id: The agent session ID to recover

        Returns:
            The recovered session or None
        """
        session = await self.session_manager.get_session(session_id)

        if not session:
            logger.warning(
                "Session not found for recovery",
                session_id=session_id
            )
            return None

        if session.state != SessionState.ACTIVE:
            logger.warning(
                "Session not active for recovery",
                session_id=session_id,
                state=session.state.value
            )
            return None

        # Resume session
        session.resume()

        # Restore heartbeat
        heartbeat_client = HeartbeatClient(
            session_id=session.session_id,
            monitor=self.heartbeat,
            interval_seconds=10
        )
        await heartbeat_client.start()
        self.heartbeat.register(session.session_id, "voice-orchestrator")

        # Store references
        self._voice_sessions[voice_session_id] = session
        self._heartbeat_clients[session.session_id] = heartbeat_client

        # Recover pending tasks from checkpoints
        await self._recover_pending_tasks(session)

        logger.info(
            "Recovered session",
            session_id=session.session_id[:8],
            checkpoints=len(session.checkpoints)
        )

        return session

    async def _recover_pending_tasks(self, session: AgentSession) -> None:
        """Recovers pending tasks from session checkpoints"""
        for checkpoint in reversed(session.checkpoints):
            if checkpoint.name == "task_queued":
                task_id = checkpoint.data.get("task_id")
                if task_id and task_id in self._tasks:
                    task = self._tasks[task_id]
                    if task.state == TaskState.QUEUED:
                        # Re-process queued task
                        await self.process_task(task)
                        logger.info(
                            "Recovered pending task",
                            task_id=task_id[:8]
                        )