gateway/modules/chat/handling/handlingTasks.py
2025-07-18 14:20:11 +02:00

291 lines
No EOL
14 KiB
Python

# handlingTasks.py
# Refactored for clarity and consolidation
import asyncio
import logging
import json
import time
from typing import Dict, Any, Optional, List, Union
from datetime import datetime, UTC
from modules.interfaces.interfaceChatModel import (
TaskStatus, TaskStep, TaskContext, TaskAction, ActionExecutionResult, ReviewResult, TaskPlan, WorkflowResult, TaskResult, ReviewContext
)
from .executionState import TaskExecutionState
from .handlingActions import HandlingActions
from .promptFactory import createTaskPlanningPrompt, createActionDefinitionPrompt, createResultReviewPrompt
logger = logging.getLogger(__name__)
class HandlingTasks:
def __init__(self, chatInterface, service, workflow=None):
self.chatInterface = chatInterface
self.service = service
self.workflow = workflow
self.handlingActions = HandlingActions(service, chatInterface)
async def generateTaskPlan(self, userInput: str, workflow) -> TaskPlan:
"""Generate a high-level task plan for the workflow."""
try:
logger.info(f"Generating task plan for workflow {workflow.id}")
prompt = await self.service.callAiTextAdvanced(
createTaskPlanningPrompt(self, {
'user_request': userInput,
'available_documents': self._getAvailableDocuments(workflow),
'workflow_id': workflow.id
})
)
task_plan_dict = self._parseTaskPlanResponse(prompt)
if not self._validateTaskPlan(task_plan_dict):
logger.error("Generated task plan failed validation")
raise Exception("AI-generated task plan failed validation - AI is required for task planning")
tasks = [TaskStep(**task_dict) for task_dict in task_plan_dict.get('tasks', [])]
return TaskPlan(
overview=task_plan_dict.get('overview', ''),
tasks=tasks
)
except Exception as e:
logger.error(f"Error in generateTaskPlan: {str(e)}")
raise
async def generateTaskActions(self, task_step, workflow, previous_results=None, enhanced_context=None) -> List[TaskAction]:
"""Generate actions for a given task step."""
try:
logger.info(f"Generating actions for task: {task_step.description}")
context = enhanced_context or TaskContext(
task_step=task_step,
workflow=workflow,
workflow_id=workflow.id,
available_documents=self._getAvailableDocuments(workflow),
previous_results=previous_results or [],
improvements=[],
retry_count=0,
previous_action_results=[],
previous_review_result=None,
is_regeneration=False,
failure_patterns=[],
failed_actions=[],
successful_actions=[]
)
prompt = await self.service.callAiTextAdvanced(
createActionDefinitionPrompt(self, context)
)
actions = self.handlingActions.parseActionResponse(prompt)
if not self._validateActions(actions, context):
logger.error("Generated actions failed validation")
raise Exception("AI-generated actions failed validation - AI is required for action generation")
# Convert to TaskAction objects
task_actions = [self.chatInterface.createTaskAction({
"execMethod": a.get('method', 'unknown'),
"execAction": a.get('action', 'unknown'),
"execParameters": a.get('parameters', {}),
"execResultLabel": a.get('resultLabel', ''),
"expectedDocumentFormats": a.get('expectedDocumentFormats', None),
"status": TaskStatus.PENDING
}) for a in actions]
return [ta for ta in task_actions if ta]
except Exception as e:
logger.error(f"Error in generateTaskActions: {str(e)}")
return []
async def executeTask(self, task_step, workflow, context) -> TaskResult:
"""Execute all actions for a task step, with state management and retries."""
logger.info(f"Executing task: {task_step.description}")
state = TaskExecutionState(task_step)
retry_context = context
max_retries = state.max_retries
for attempt in range(max_retries):
logger.info(f"Task execution attempt {attempt+1}/{max_retries}")
actions = await self.generateTaskActions(task_step, workflow, previous_results=retry_context.previous_results, enhanced_context=retry_context)
if not actions:
logger.error("No actions defined for task step, aborting task execution")
break
action_results = []
for action in actions:
result = await self.handlingActions.executeSingleAction(action, workflow)
action_results.append(result)
if result.success:
state.addSuccessfulAction(result)
else:
state.addFailedAction(result)
review_result = await self.reviewTaskCompletion(task_step, actions, action_results, workflow)
success = review_result.status == 'success'
feedback = review_result.reason
error = None if success else review_result.reason
if success:
logger.info(f"Task step '{task_step.description}' completed successfully")
return TaskResult(
taskId=task_step.id,
status=TaskStatus.COMPLETED,
success=True,
feedback=feedback,
error=None
)
elif review_result.status == 'retry' and state.canRetry():
logger.warning(f"Task step '{task_step.description}' requires retry: {review_result.improvements}")
state.incrementRetryCount()
retry_context.retry_count = state.retry_count
retry_context.improvements = review_result.improvements
retry_context.previous_action_results = action_results
retry_context.previous_review_result = review_result
retry_context.is_regeneration = True
retry_context.failure_patterns = state.getFailurePatterns()
retry_context.failed_actions = state.failed_actions
retry_context.successful_actions = state.successful_actions
continue
else:
logger.error(f"Task step '{task_step.description}' failed after {attempt+1} attempts")
return TaskResult(
taskId=task_step.id,
status=TaskStatus.FAILED,
success=False,
feedback=feedback,
error=error
)
logger.error(f"Task step '{task_step.description}' failed after all retries")
return TaskResult(
taskId=task_step.id,
status=TaskStatus.FAILED,
success=False,
feedback="Task failed after all retries.",
error="Task failed after all retries."
)
async def reviewTaskCompletion(self, task_step, task_actions, action_results, workflow):
try:
review_context = ReviewContext(
task_step=task_step,
action_results=action_results,
workflow=workflow,
step_result={
'successful_actions': sum(1 for result in action_results if result.success),
'total_actions': len(action_results),
'results': [result.data.get('result', '') for result in action_results if result.success],
'errors': [result.error for result in action_results if not result.success]
}
)
# Use promptFactory for review prompt
prompt = await createResultReviewPrompt(self, review_context)
response = await self.service.callAiTextAdvanced(prompt)
review_dict = self.handlingActions.parseReviewResponse(response)
review_dict.setdefault('status', 'unknown')
review_dict.setdefault('reason', 'No reason provided')
review_dict.setdefault('quality_score', 5)
return ReviewResult(
status=review_dict.get('status', 'unknown'),
reason=review_dict.get('reason', 'No reason provided'),
improvements=review_dict.get('improvements', []),
quality_score=review_dict.get('quality_score', 5),
missing_outputs=review_dict.get('missing_outputs', []),
met_criteria=review_dict.get('met_criteria', []),
unmet_criteria=review_dict.get('unmet_criteria', []),
confidence=review_dict.get('confidence', 0.5)
)
except Exception as e:
logger.error(f"Error in reviewTaskCompletion: {str(e)}")
return ReviewResult(
status='failed',
reason=str(e),
quality_score=0
)
async def prepareTaskHandover(self, task_step, task_actions, review_result, workflow):
try:
handover_data = {
'task_id': task_step.id,
'task_description': task_step.description,
'actions': [action.to_dict() for action in task_actions],
'review_result': review_result.to_dict() if hasattr(review_result, 'to_dict') else review_result,
'workflow_id': workflow.id,
'handover_time': datetime.now(UTC).isoformat()
}
logger.info(f"Prepared handover for task {task_step.id} in workflow {workflow.id}")
return handover_data
except Exception as e:
logger.error(f"Error in prepareTaskHandover: {str(e)}")
return {'error': str(e)}
# --- Helper and validation methods (unchanged, but can be inlined or made private) ---
def _getAvailableDocuments(self, workflow):
documents = []
for message in workflow.messages:
for doc in message.documents:
documents.append(doc.filename)
return documents
def _parseTaskPlanResponse(self, response: str) -> dict:
try:
json_start = response.find('{')
json_end = response.rfind('}') + 1
if json_start == -1 or json_end == 0:
raise ValueError("No JSON found in response")
json_str = response[json_start:json_end]
task_plan = json.loads(json_str)
if 'tasks' not in task_plan:
raise ValueError("Task plan missing 'tasks' field")
return task_plan
except Exception as e:
logger.error(f"Error parsing task plan response: {str(e)}")
return {'tasks': []}
def _validateTaskPlan(self, task_plan: Dict[str, Any]) -> bool:
try:
if not isinstance(task_plan, dict):
return False
if 'tasks' not in task_plan or not isinstance(task_plan['tasks'], list):
return False
task_ids = set()
for task in task_plan['tasks']:
if not isinstance(task, dict):
return False
required_fields = ['id', 'description', 'expected_outputs', 'success_criteria']
if not all(field in task for field in required_fields):
return False
if task['id'] in task_ids:
return False
task_ids.add(task['id'])
dependencies = task.get('dependencies', [])
if not isinstance(dependencies, list):
return False
for dep in dependencies:
if dep not in task_ids and dep != 'task_0':
return False
if 'ai_prompt' in task and not isinstance(task['ai_prompt'], str):
return False
return True
except Exception as e:
logger.error(f"Error validating task plan: {str(e)}")
return False
def _validateActions(self, actions: List[Dict[str, Any]], context) -> bool:
try:
if not isinstance(actions, list):
logger.error("Actions must be a list")
return False
if len(actions) == 0:
logger.warning("No actions generated")
return False
for i, action in enumerate(actions):
if not isinstance(action, dict):
logger.error(f"Action {i} must be a dictionary")
return False
required_fields = ['method', 'action', 'parameters', 'resultLabel']
missing_fields = []
for field in required_fields:
if field not in action or not action[field]:
missing_fields.append(field)
if missing_fields:
logger.error(f"Action {i} missing required fields: {missing_fields}")
return False
result_label = action.get('resultLabel', '')
if not result_label.startswith('task'):
logger.error(f"Action {i} result label must start with 'task': {result_label}")
return False
parameters = action.get('parameters', {})
if not isinstance(parameters, dict):
logger.error(f"Action {i} parameters must be a dictionary")
return False
logger.info(f"Successfully validated {len(actions)} actions")
return True
except Exception as e:
logger.error(f"Error validating actions: {str(e)}")
return False