418 lines
20 KiB
Python
418 lines
20 KiB
Python
# modeAutomation.py
|
|
# Automation mode implementation for workflows with predefined action plans
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from typing import List, Dict, Any, Optional
|
|
from modules.datamodels.datamodelChat import (
|
|
TaskStep, TaskContext, TaskResult, ActionItem, TaskStatus,
|
|
TaskPlan, ActionResult
|
|
)
|
|
from modules.datamodels.datamodelChat import ChatWorkflow
|
|
from modules.workflows.processing.modes.modeBase import BaseMode
|
|
from modules.workflows.processing.shared.stateTools import checkWorkflowStopped
|
|
from modules.shared.timeUtils import parseTimestamp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class AutomationMode(BaseMode):
|
|
"""Automation mode implementation - executes workflows from predefined plans"""
|
|
|
|
def __init__(self, services):
|
|
super().__init__(services)
|
|
# Store action lists for each task (mapped by task ID)
|
|
self.taskActionMap: Dict[str, List[Dict[str, Any]]] = {}
|
|
logger.info("AutomationMode initialized - will use predefined plan from workflow")
|
|
|
|
async def generateTaskPlan(self, userInput: str, workflow: ChatWorkflow) -> TaskPlan:
|
|
"""
|
|
Generate task plan from JSON plan in userInput (no AI planning needed).
|
|
AUTOMATION mode ALWAYS requires a JSON plan to be provided in the user input.
|
|
The plan can be:
|
|
- Embedded between <!--TEMPLATE_PLAN_START--> and <!--TEMPLATE_PLAN_END-->
|
|
- Or as direct JSON in userInput
|
|
"""
|
|
try:
|
|
# AUTOMATION mode ALWAYS requires a JSON plan to be provided in userInput
|
|
# Try to extract plan from userInput (embedded JSON or direct JSON)
|
|
templatePlan = None
|
|
try:
|
|
# Look for embedded plan in prompt (between <!--TEMPLATE_PLAN_START--> and <!--TEMPLATE_PLAN_END-->)
|
|
startMarker = "<!--TEMPLATE_PLAN_START-->"
|
|
endMarker = "<!--TEMPLATE_PLAN_END-->"
|
|
startIdx = userInput.find(startMarker)
|
|
endIdx = userInput.find(endMarker)
|
|
|
|
if startIdx >= 0 and endIdx > startIdx:
|
|
planJson = userInput[startIdx + len(startMarker):endIdx].strip()
|
|
templatePlan = json.loads(planJson)
|
|
logger.info("Extracted template plan from embedded JSON in prompt")
|
|
elif '{' in userInput and '"tasks"' in userInput:
|
|
# Try parsing entire userInput as JSON (fallback)
|
|
jsonStart = userInput.find('{')
|
|
jsonEnd = userInput.rfind('}') + 1
|
|
if jsonStart >= 0 and jsonEnd > jsonStart:
|
|
templatePlan = json.loads(userInput[jsonStart:jsonEnd])
|
|
logger.info("Parsed template plan from userInput JSON (fallback)")
|
|
else:
|
|
raise ValueError("No template plan found in userInput. AUTOMATION mode requires a JSON plan to be provided in the user input.")
|
|
except (json.JSONDecodeError, ValueError) as e:
|
|
logger.error(f"Could not parse template plan: {str(e)}")
|
|
raise ValueError(f"AUTOMATION mode requires a predefined JSON plan with 'tasks' array, but none was found. Please provide the plan in the user input (embedded between <!--TEMPLATE_PLAN_START--> and <!--TEMPLATE_PLAN_END--> or as direct JSON). Error: {str(e)}")
|
|
|
|
logger.info(f"Using template plan with {len(templatePlan.get('tasks', []))} tasks")
|
|
|
|
# Convert plan tasks to TaskStep objects
|
|
tasks = []
|
|
for i, taskDict in enumerate(templatePlan.get('tasks', [])):
|
|
if not isinstance(taskDict, dict):
|
|
logger.warning(f"Skipping invalid task {i+1}: not a dictionary")
|
|
continue
|
|
|
|
# Store actionList for this task (before converting to TaskStep)
|
|
taskId = taskDict.get('id', f"task_{i+1}")
|
|
actionList = taskDict.get('actionList', [])
|
|
if actionList:
|
|
self.taskActionMap[taskId] = actionList
|
|
logger.info(f"Stored {len(actionList)} actions for task {taskId}")
|
|
|
|
# Remove actionList from taskDict (TaskStep doesn't have this field)
|
|
taskDictWithoutActions = {k: v for k, v in taskDict.items() if k != 'actionList'}
|
|
|
|
try:
|
|
task = TaskStep(**taskDictWithoutActions)
|
|
tasks.append(task)
|
|
except Exception as e:
|
|
logger.warning(f"Skipping invalid task {i+1}: {str(e)}")
|
|
continue
|
|
|
|
if not tasks:
|
|
raise ValueError("No valid tasks found in template plan")
|
|
|
|
taskPlan = TaskPlan(
|
|
overview=templatePlan.get('overview', 'Automated workflow execution'),
|
|
tasks=tasks,
|
|
userMessage=templatePlan.get('userMessage', '')
|
|
)
|
|
|
|
logger.info(f"Generated task plan from template with {len(tasks)} tasks")
|
|
|
|
return taskPlan
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating task plan from template: {str(e)}")
|
|
raise
|
|
|
|
async def generateActionItems(self, taskStep: TaskStep, workflow: ChatWorkflow,
|
|
previousResults: List = None, enhancedContext: TaskContext = None) -> List[ActionItem]:
|
|
"""
|
|
Generate actions from predefined actionList in template plan.
|
|
No AI planning needed - actions are already defined in the plan.
|
|
"""
|
|
try:
|
|
taskId = taskStep.id
|
|
|
|
# Get actionList from stored map
|
|
actionList = self.taskActionMap.get(taskId, [])
|
|
|
|
if not actionList:
|
|
logger.warning(f"No actionList found for task {taskId} in template plan")
|
|
return []
|
|
|
|
logger.info(f"Generating {len(actionList)} actions for task {taskId} from template")
|
|
|
|
# Convert actionList to ActionItem objects
|
|
actionItems = []
|
|
for i, actionDict in enumerate(actionList):
|
|
if not isinstance(actionDict, dict):
|
|
logger.warning(f"Skipping invalid action {i+1} in task {taskId}: not a dictionary")
|
|
continue
|
|
|
|
# Extract action fields
|
|
execMethod = actionDict.get('execMethod', '')
|
|
execAction = actionDict.get('execAction', '')
|
|
execParameters = actionDict.get('execParameters', {})
|
|
execResultLabel = actionDict.get('execResultLabel', '')
|
|
|
|
if not execMethod or not execAction:
|
|
logger.warning(f"Skipping invalid action {i+1} in task {taskId}: missing execMethod or execAction")
|
|
continue
|
|
|
|
# Create ActionItem
|
|
actionItem = self._createActionItem({
|
|
"id": f"action_{uuid.uuid4()}",
|
|
"execMethod": execMethod,
|
|
"execAction": execAction,
|
|
"execParameters": execParameters,
|
|
"execResultLabel": execResultLabel,
|
|
"expectedDocumentFormats": actionDict.get('expectedDocumentFormats', None),
|
|
"status": TaskStatus.PENDING,
|
|
"userMessage": actionDict.get('userMessage', None)
|
|
})
|
|
|
|
if actionItem:
|
|
actionItems.append(actionItem)
|
|
else:
|
|
logger.warning(f"Failed to create ActionItem for action {i+1} in task {taskId}")
|
|
|
|
logger.info(f"Generated {len(actionItems)} action items for task {taskId}")
|
|
return actionItems
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating action items from template: {str(e)}")
|
|
return []
|
|
|
|
async def executeTask(self, taskStep: TaskStep, workflow: ChatWorkflow, context: TaskContext,
|
|
taskIndex: int = None, totalTasks: int = None) -> TaskResult:
|
|
"""
|
|
Execute task using Automation mode - executes predefined actions directly.
|
|
No AI planning or review phases - actions are executed sequentially as defined.
|
|
"""
|
|
# Get task index from workflow state for consistency
|
|
if taskIndex is None:
|
|
taskIndex = workflow.getTaskIndex()
|
|
|
|
logger.info(f"=== STARTING TASK {taskIndex or '?'}: {taskStep.objective} ===")
|
|
|
|
try:
|
|
# Check workflow status
|
|
checkWorkflowStopped(self.services)
|
|
|
|
# Update workflow before executing task
|
|
if taskIndex is not None:
|
|
self._updateWorkflowBeforeExecutingTask(taskIndex)
|
|
|
|
# Create task start message
|
|
await self.messageCreator.createTaskStartMessage(taskStep, workflow, taskIndex, totalTasks)
|
|
|
|
# Generate actions from template plan
|
|
actions = await self.generateActionItems(taskStep, workflow,
|
|
previousResults=context.previousResults if context else None,
|
|
enhancedContext=context)
|
|
|
|
if not actions:
|
|
logger.error(f"No actions found for task {taskIndex}, aborting")
|
|
return TaskResult(
|
|
taskId=taskStep.id,
|
|
status=TaskStatus.FAILED,
|
|
success=False,
|
|
feedback="No actions defined in template plan for this task",
|
|
error="No actions found in template plan"
|
|
)
|
|
|
|
# Update workflow after action planning
|
|
totalActions = len(actions)
|
|
self._updateWorkflowAfterActionPlanning(totalActions)
|
|
self._setWorkflowTotals(totalActions=totalActions)
|
|
|
|
logger.info(f"Task {taskIndex} has {totalActions} actions to execute")
|
|
|
|
# Execute all actions sequentially
|
|
actionResults = []
|
|
for actionIdx, action in enumerate(actions):
|
|
# Check workflow status before each action
|
|
checkWorkflowStopped(self.services)
|
|
|
|
# Update workflow before executing action
|
|
actionNumber = actionIdx + 1
|
|
self._updateWorkflowBeforeExecutingAction(actionNumber)
|
|
|
|
logger.info(f"Task {taskIndex} - Starting action {actionNumber}/{totalActions}: {action.execMethod}.{action.execAction}")
|
|
|
|
# Create action start message
|
|
actionStartMessage = {
|
|
"workflowId": workflow.id,
|
|
"role": "assistant",
|
|
"message": f"⚡ **Action {actionNumber}** (Method {action.execMethod}.{action.execAction})",
|
|
"status": "step",
|
|
"sequenceNr": len(workflow.messages) + 1,
|
|
"publishedAt": self.services.utils.timestampGetUtc(),
|
|
"documentsLabel": f"action_{actionNumber}_start",
|
|
"documents": [],
|
|
"actionProgress": "running",
|
|
"roundNumber": workflow.currentRound,
|
|
"taskNumber": taskIndex,
|
|
"actionNumber": actionNumber
|
|
}
|
|
|
|
# Add user-friendly message if available
|
|
if action.userMessage:
|
|
actionStartMessage["message"] += f"\n\n💬 {action.userMessage}"
|
|
|
|
self.services.chat.storeMessageWithDocuments(workflow, actionStartMessage, [])
|
|
|
|
# Execute action
|
|
result = await self.actionExecutor.executeSingleAction(
|
|
action, workflow, taskStep
|
|
)
|
|
actionResults.append(result)
|
|
|
|
if result.success:
|
|
logger.info(f"Action {actionNumber} completed successfully")
|
|
else:
|
|
logger.warning(f"Action {actionNumber} failed: {result.error}")
|
|
|
|
# Check if all actions succeeded
|
|
allSucceeded = all(r.success for r in actionResults)
|
|
failedCount = sum(1 for r in actionResults if not r.success)
|
|
|
|
if allSucceeded:
|
|
logger.info(f"=== TASK {taskIndex or '?'} COMPLETED SUCCESSFULLY: {taskStep.objective} ===")
|
|
|
|
# Create task completion message
|
|
await self.messageCreator.createTaskCompletionMessage(
|
|
taskStep, workflow, taskIndex, totalTasks, None
|
|
)
|
|
|
|
return TaskResult(
|
|
taskId=taskStep.id,
|
|
status=TaskStatus.COMPLETED,
|
|
success=True,
|
|
feedback=f"Task completed successfully with {totalActions} actions",
|
|
error=None
|
|
)
|
|
else:
|
|
logger.error(f"=== TASK {taskIndex or '?'} FAILED: {taskStep.objective} ({failedCount}/{totalActions} actions failed) ===")
|
|
|
|
errorMessages = [r.error for r in actionResults if r.error]
|
|
errorSummary = "; ".join(errorMessages[:3]) # Limit to first 3 errors
|
|
|
|
await self.messageCreator.createErrorMessage(
|
|
taskStep, workflow, taskIndex, errorSummary
|
|
)
|
|
|
|
return TaskResult(
|
|
taskId=taskStep.id,
|
|
status=TaskStatus.FAILED,
|
|
success=False,
|
|
feedback=f"Task failed: {failedCount} out of {totalActions} actions failed",
|
|
error=errorSummary
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error executing task {taskIndex}: {str(e)}")
|
|
await self.messageCreator.createErrorMessage(taskStep, workflow, taskIndex, str(e))
|
|
|
|
return TaskResult(
|
|
taskId=taskStep.id,
|
|
status=TaskStatus.FAILED,
|
|
success=False,
|
|
feedback="Task execution failed",
|
|
error=str(e)
|
|
)
|
|
|
|
def _createActionItem(self, actionData: Dict[str, Any]) -> Optional[ActionItem]:
|
|
"""Create ActionItem from action data"""
|
|
try:
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
# Ensure ID is present
|
|
if "id" not in actionData or not actionData["id"]:
|
|
actionData["id"] = f"action_{uuid.uuid4()}"
|
|
|
|
# Ensure required fields
|
|
if "status" not in actionData:
|
|
actionData["status"] = TaskStatus.PENDING
|
|
|
|
if "execMethod" not in actionData:
|
|
logger.error("execMethod is required for task action")
|
|
return None
|
|
|
|
if "execAction" not in actionData:
|
|
logger.error("execAction is required for task action")
|
|
return None
|
|
|
|
if "execParameters" not in actionData:
|
|
actionData["execParameters"] = {}
|
|
|
|
# Use generic field separation based on ActionItem model
|
|
simpleFields, objectFields = self.services.interfaceDbChat._separateObjectFields(ActionItem, actionData)
|
|
|
|
# Create action in database
|
|
createdAction = self.services.interfaceDbChat.db.recordCreate(ActionItem, simpleFields)
|
|
|
|
# Convert to ActionItem model
|
|
return ActionItem(
|
|
id=createdAction["id"],
|
|
execMethod=createdAction["execMethod"],
|
|
execAction=createdAction["execAction"],
|
|
execParameters=createdAction.get("execParameters", {}),
|
|
execResultLabel=createdAction.get("execResultLabel"),
|
|
expectedDocumentFormats=createdAction.get("expectedDocumentFormats"),
|
|
status=createdAction.get("status", TaskStatus.PENDING),
|
|
error=createdAction.get("error"),
|
|
retryCount=createdAction.get("retryCount", 0),
|
|
retryMax=createdAction.get("retryMax", 3),
|
|
processingTime=createdAction.get("processingTime"),
|
|
timestamp=parseTimestamp(createdAction.get("timestamp"), default=self.services.utils.timestampGetUtc()),
|
|
result=createdAction.get("result"),
|
|
userMessage=createdAction.get("userMessage")
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating task action: {str(e)}")
|
|
return None
|
|
|
|
def _updateWorkflowBeforeExecutingTask(self, taskNumber: int):
|
|
"""Update workflow object before executing a task"""
|
|
try:
|
|
workflow = self.services.workflow
|
|
updateData = {
|
|
"currentTask": taskNumber,
|
|
"currentAction": 0,
|
|
"totalActions": 0
|
|
}
|
|
|
|
workflow.currentTask = taskNumber
|
|
workflow.currentAction = 0
|
|
workflow.totalActions = 0
|
|
|
|
self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
|
|
logger.info(f"Updated workflow {workflow.id} before executing task {taskNumber}")
|
|
except Exception as e:
|
|
logger.error(f"Error updating workflow before executing task: {str(e)}")
|
|
|
|
def _updateWorkflowAfterActionPlanning(self, totalActions: int):
|
|
"""Update workflow object after action planning"""
|
|
try:
|
|
workflow = self.services.workflow
|
|
updateData = {"totalActions": totalActions}
|
|
workflow.totalActions = totalActions
|
|
self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
|
|
logger.info(f"Updated workflow {workflow.id} after action planning: {updateData}")
|
|
except Exception as e:
|
|
logger.error(f"Error updating workflow after action planning: {str(e)}")
|
|
|
|
def _updateWorkflowBeforeExecutingAction(self, actionNumber: int):
|
|
"""Update workflow object before executing an action"""
|
|
try:
|
|
workflow = self.services.workflow
|
|
updateData = {"currentAction": actionNumber}
|
|
workflow.currentAction = actionNumber
|
|
self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
|
|
logger.info(f"Updated workflow {workflow.id} before executing action {actionNumber}")
|
|
except Exception as e:
|
|
logger.error(f"Error updating workflow before executing action: {str(e)}")
|
|
|
|
def _setWorkflowTotals(self, totalTasks: int = None, totalActions: int = None):
|
|
"""Set total counts for workflow progress tracking"""
|
|
try:
|
|
workflow = self.services.workflow
|
|
updateData = {}
|
|
|
|
if totalTasks is not None:
|
|
workflow.totalTasks = totalTasks
|
|
updateData["totalTasks"] = totalTasks
|
|
|
|
if totalActions is not None:
|
|
workflow.totalActions = totalActions
|
|
updateData["totalActions"] = totalActions
|
|
|
|
if updateData:
|
|
self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
|
|
logger.info(f"Updated workflow {workflow.id} totals: {updateData}")
|
|
except Exception as e:
|
|
logger.error(f"Error setting workflow totals: {str(e)}")
|
|
|