""" 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] )