60 lines
3 KiB
Python
60 lines
3 KiB
Python
# modeBase.py
|
|
# Abstract base class for workflow modes
|
|
|
|
from abc import ABC, abstractmethod
|
|
import logging
|
|
from typing import List, Dict, Any
|
|
from modules.datamodels.datamodelWorkflow import TaskStep, TaskContext, TaskResult, TaskAction
|
|
from modules.datamodels.datamodelChat import ChatWorkflow
|
|
from modules.workflows.processing.core.taskPlanner import TaskPlanner
|
|
from modules.workflows.processing.core.actionExecutor import ActionExecutor
|
|
from modules.workflows.processing.core.messageCreator import MessageCreator
|
|
from modules.workflows.processing.core.validator import WorkflowValidator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class BaseMode(ABC):
|
|
"""Abstract base class for workflow execution modes"""
|
|
|
|
def __init__(self, services, workflow):
|
|
self.services = services
|
|
self.workflow = workflow
|
|
self.taskPlanner = TaskPlanner(services)
|
|
self.actionExecutor = ActionExecutor(services)
|
|
self.messageCreator = MessageCreator(services)
|
|
self.validator = WorkflowValidator(services)
|
|
|
|
def _checkWorkflowStopped(self, workflow):
|
|
"""Check if workflow has been stopped by user and raise exception if so"""
|
|
try:
|
|
# Get the current workflow status from the database to avoid stale data
|
|
current_workflow = self.services.interfaceDbChat.getWorkflow(workflow.id)
|
|
if current_workflow and current_workflow.status == "stopped":
|
|
logger.info("Workflow stopped by user, aborting execution")
|
|
raise Exception("Workflow was stopped by user")
|
|
except Exception as e:
|
|
# If we can't get the current status due to other database issues, fall back to the in-memory object
|
|
logger.warning(f"Could not check current workflow status from database: {str(e)}")
|
|
if workflow and workflow.status == "stopped":
|
|
logger.info("Workflow stopped by user (from in-memory object), aborting execution")
|
|
raise Exception("Workflow was stopped by user")
|
|
|
|
@abstractmethod
|
|
async def executeTask(self, taskStep: TaskStep, workflow: ChatWorkflow, context: TaskContext,
|
|
taskIndex: int = None, totalTasks: int = None) -> TaskResult:
|
|
"""Execute a task step - must be implemented by concrete modes"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def generateTaskActions(self, taskStep: TaskStep, workflow: ChatWorkflow,
|
|
previousResults: List = None, enhancedContext: TaskContext = None) -> List[TaskAction]:
|
|
"""Generate actions for a task step - must be implemented by concrete modes"""
|
|
pass
|
|
|
|
async def generateTaskPlan(self, userInput: str, workflow: ChatWorkflow):
|
|
"""Generate task plan - common to all modes"""
|
|
return await self.taskPlanner.generateTaskPlan(userInput, workflow)
|
|
|
|
async def createTaskPlanMessage(self, taskPlan, workflow: ChatWorkflow):
|
|
"""Create task plan message - common to all modes"""
|
|
return await self.messageCreator.createTaskPlanMessage(taskPlan, workflow)
|