156 lines
5.5 KiB
Python
156 lines
5.5 KiB
Python
"""
|
|
Workflow Manager Module for state machine-based backend chat workflow.
|
|
Implements the state machine as defined in the documentation.
|
|
"""
|
|
|
|
import logging
|
|
from datetime import datetime, UTC, timedelta
|
|
import time
|
|
|
|
from modules.workflow.chatManager import getChatManager
|
|
from modules.interfaces.serviceChatModel import ( UserInputRequest, ChatWorkflow )
|
|
|
|
# Configure logger
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class WorkflowStoppedException(Exception):
|
|
"""Exception raised when a workflow is forcibly stopped with function checkExitCriteria() """
|
|
pass
|
|
|
|
class WorkflowManager:
|
|
"""Manages the execution of workflows and their associated agents."""
|
|
|
|
def __init__(self):
|
|
"""Initialize the workflow manager."""
|
|
self.chatManager = getChatManager()
|
|
|
|
def initialize(self, workflow: ChatWorkflow):
|
|
"""
|
|
Initialize the workflow manager with a workflow object.
|
|
|
|
Args:
|
|
workflow: ChatWorkflow object to initialize with
|
|
"""
|
|
self.chatManager.initialize(workflow)
|
|
return True
|
|
|
|
# Main function to start workflow process
|
|
|
|
async def workflowProcess(self, userInput: UserInputRequest, workflow: ChatWorkflow) -> ChatWorkflow:
|
|
"""
|
|
Main processing function that implements the workflow state machine.
|
|
Handles the complete workflow process from user input to final response.
|
|
|
|
Args:
|
|
userInput: User input with prompt and optional file list
|
|
workflow: Current ChatWorkflow object
|
|
|
|
Returns:
|
|
Updated ChatWorkflow object with processing results
|
|
"""
|
|
startTime = time.time()
|
|
try:
|
|
# Initialize workflow
|
|
self.initialize(workflow)
|
|
|
|
# Check if workflow should exit
|
|
self.checkExitCriteria(workflow)
|
|
|
|
# Create initial handover
|
|
handover = self.chatManager.createInitialHandover(userInput)
|
|
|
|
# Process agents until completion or failure
|
|
while True:
|
|
self.checkExitCriteria(workflow)
|
|
|
|
# Define next handover
|
|
nextHandover = await self.chatManager.defineNextHandover(handover)
|
|
if not nextHandover:
|
|
break
|
|
|
|
# Process next agent
|
|
handover = await self.chatManager.processNextAgent(nextHandover)
|
|
|
|
# Check if we should continue
|
|
if handover.status in ["failed", "retry"]:
|
|
break
|
|
|
|
# Send final message
|
|
finalMessage = await self.chatManager.sendFinalMessage(handover)
|
|
workflow.messages.append(finalMessage)
|
|
|
|
# Update workflow stats
|
|
endTime = time.time()
|
|
workflow.stats.processingTime = endTime - startTime
|
|
|
|
# Update workflow status
|
|
workflow.status = "completed"
|
|
workflow.lastActivity = datetime.now(UTC).isoformat()
|
|
|
|
return workflow
|
|
|
|
except WorkflowStoppedException:
|
|
# Handle workflow stop
|
|
workflow.status = "stopped"
|
|
workflow.lastActivity = datetime.now(UTC).isoformat()
|
|
return workflow
|
|
|
|
except Exception as e:
|
|
# Handle workflow failure
|
|
logger.error(f"Workflow processing error: {str(e)}", exc_info=True)
|
|
workflow.status = "failed"
|
|
workflow.lastActivity = datetime.now(UTC).isoformat()
|
|
|
|
# Update processing time even on error
|
|
endTime = time.time()
|
|
workflow.stats.processingTime = endTime - startTime
|
|
|
|
logger.error(f"Workflow failed: {str(e)}", level="error", progress=100)
|
|
return workflow
|
|
|
|
# Workflow state machine functions
|
|
|
|
def checkExitCriteria(self, workflow: ChatWorkflow) -> None:
|
|
"""
|
|
Check if the workflow should exit based on the current state.
|
|
Raises WorkflowStoppedException if workflow should stop.
|
|
|
|
Args:
|
|
workflow: ChatWorkflow object to check
|
|
"""
|
|
if workflow.status in ["stopped", "failed"]:
|
|
logger.info(f"Workflow processing terminated due to status: {workflow.status}")
|
|
# Raise an exception to stop execution
|
|
raise WorkflowStoppedException(f"Workflow execution stopped due to status: {workflow.status}")
|
|
|
|
def workflowFinish(self, workflow: ChatWorkflow) -> ChatWorkflow:
|
|
"""
|
|
Finalizes a workflow and sets the status to 'completed'.
|
|
|
|
Args:
|
|
workflow: ChatWorkflow object
|
|
|
|
Returns:
|
|
Updated ChatWorkflow object
|
|
"""
|
|
# Prepare workflow update data
|
|
workflowUpdate = {
|
|
"status": "completed",
|
|
"lastActivity": datetime.now().isoformat(),
|
|
}
|
|
|
|
# Update the workflow object in memory
|
|
workflow.status = workflowUpdate["status"]
|
|
workflow.lastActivity = workflowUpdate["lastActivity"]
|
|
|
|
# Save workflow state to database
|
|
self.chatManager.updateWorkflow(workflow.id, workflowUpdate)
|
|
|
|
logger.info(f"Workflow completed successfully", level="info", progress=100)
|
|
return workflow
|
|
|
|
async def getWorkflowManager() -> WorkflowManager:
|
|
"""Get or create a workflow manager instance."""
|
|
# Create new instance
|
|
manager = WorkflowManager()
|
|
return manager
|