463 lines
No EOL
18 KiB
Python
463 lines
No EOL
18 KiB
Python
"""
|
|
Refactored WorkflowManager class for the Agentservice (continued).
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
import asyncio
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import List, Dict, Any, Optional, Tuple, Union
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class WorkflowManager:
|
|
# Previous code is in the first part
|
|
|
|
async def list_workflows(self, mandate_id: int = None, user_id: int = None) -> List[Dict[str, Any]]:
|
|
"""
|
|
List all available workflows.
|
|
|
|
Args:
|
|
mandate_id: Optional mandate ID for filtering
|
|
user_id: Optional user ID for filtering
|
|
|
|
Returns:
|
|
List of workflow summaries
|
|
"""
|
|
workflows = []
|
|
|
|
# Load from database if available
|
|
if self.lucydom_interface:
|
|
try:
|
|
# Get all workflows for the user
|
|
if user_id is not None:
|
|
user_workflows = self.lucydom_interface.get_workflows_by_user(user_id)
|
|
else:
|
|
user_workflows = self.lucydom_interface.get_all_workflows()
|
|
|
|
# Filter by mandate if specified
|
|
if mandate_id is not None:
|
|
user_workflows = [wf for wf in user_workflows if wf.get("mandate_id") == mandate_id]
|
|
|
|
# Create workflow summaries
|
|
for workflow in user_workflows:
|
|
summary = {
|
|
"id": workflow.get("id"),
|
|
"name": workflow.get("name", f"Workflow {workflow.get('id')}"),
|
|
"status": workflow.get("status"),
|
|
"started_at": workflow.get("started_at"),
|
|
"last_activity": workflow.get("last_activity"),
|
|
"completed_at": workflow.get("completed_at")
|
|
}
|
|
|
|
# Add message count if available
|
|
messages = self.lucydom_interface.get_workflow_messages(workflow.get("id"))
|
|
if messages:
|
|
summary["message_count"] = len(messages)
|
|
|
|
workflows.append(summary)
|
|
|
|
logger.info(f"Loaded {len(workflows)} workflows from database")
|
|
|
|
# Sort by last activity (newest first)
|
|
return sorted(workflows, key=lambda w: w.get("last_activity", ""), reverse=True)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error retrieving workflows from database: {str(e)}")
|
|
|
|
# Load from files if no database or error occurred
|
|
try:
|
|
for filename in os.listdir(self.results_dir):
|
|
if filename.startswith("workflow_") and filename.endswith(".json"):
|
|
workflow_path = os.path.join(self.results_dir, filename)
|
|
|
|
try:
|
|
import json
|
|
with open(workflow_path, 'r', encoding='utf-8') as f:
|
|
workflow = json.load(f)
|
|
|
|
# Check if mandate and user ID match filters
|
|
if mandate_id is not None and workflow.get("mandate_id") != mandate_id:
|
|
continue
|
|
|
|
if user_id is not None and workflow.get("user_id") != user_id:
|
|
continue
|
|
|
|
# Create workflow summary
|
|
summary = {
|
|
"id": workflow.get("id"),
|
|
"name": workflow.get("name", f"Workflow {workflow.get('id')}"),
|
|
"status": workflow.get("status"),
|
|
"started_at": workflow.get("started_at"),
|
|
"last_activity": workflow.get("last_activity"),
|
|
"message_count": len(workflow.get("messages", []))
|
|
}
|
|
|
|
workflows.append(summary)
|
|
except Exception as e:
|
|
logger.error(f"Error loading workflow file {filename}: {str(e)}")
|
|
|
|
logger.info(f"Loaded {len(workflows)} workflows from files")
|
|
|
|
# Sort by last activity (newest first)
|
|
return sorted(workflows, key=lambda w: w.get("last_activity", ""), reverse=True)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error listing workflows: {str(e)}")
|
|
return []
|
|
|
|
async def delete_workflow(self, workflow_id: str) -> bool:
|
|
"""
|
|
Delete a workflow.
|
|
|
|
Args:
|
|
workflow_id: ID of the workflow
|
|
|
|
Returns:
|
|
True on success, False if workflow not found
|
|
"""
|
|
# Remove from memory
|
|
if workflow_id in self.workflows:
|
|
del self.workflows[workflow_id]
|
|
|
|
# Delete from database
|
|
if self.lucydom_interface:
|
|
try:
|
|
db_success = self.lucydom_interface.delete_workflow(workflow_id)
|
|
logger.info(f"Workflow {workflow_id} deleted from database: {db_success}")
|
|
except Exception as e:
|
|
logger.error(f"Error deleting workflow {workflow_id} from database: {str(e)}")
|
|
|
|
# Delete file
|
|
workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
|
|
|
|
try:
|
|
if os.path.exists(workflow_path):
|
|
os.remove(workflow_path)
|
|
logger.info(f"Workflow {workflow_id} deleted from file: {workflow_path}")
|
|
return True
|
|
else:
|
|
logger.warning(f"Workflow {workflow_id} not found: {workflow_path}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"Error deleting workflow file {workflow_id}: {str(e)}")
|
|
return False
|
|
|
|
def _initialize_workflow(self, workflow_id: str) -> Dict[str, Any]:
|
|
"""
|
|
Initialize a new workflow.
|
|
|
|
Args:
|
|
workflow_id: ID of the workflow
|
|
|
|
Returns:
|
|
The initialized workflow object
|
|
"""
|
|
current_time = datetime.now().isoformat()
|
|
|
|
# Create complete workflow object according to the data model
|
|
workflow = {
|
|
"id": workflow_id,
|
|
"name": f"Workflow {workflow_id}",
|
|
"mandate_id": self.mandate_id,
|
|
"user_id": self.user_id,
|
|
"status": "running",
|
|
"started_at": current_time,
|
|
"last_activity": current_time,
|
|
"current_round": 1,
|
|
|
|
# Complete statistics structure according to DataStats model
|
|
"data_stats": {
|
|
"total_processing_time": 0.0,
|
|
"total_token_count": 0,
|
|
"total_bytes_sent": 0,
|
|
"total_bytes_received": 0
|
|
},
|
|
|
|
# Empty arrays for messages and logs
|
|
"messages": [],
|
|
"logs": []
|
|
}
|
|
|
|
# Log entry for workflow start
|
|
self._add_log(workflow, "Workflow started", "info", "workflow", "Workflow Management")
|
|
|
|
# Save workflow to database
|
|
if self.lucydom_interface:
|
|
try:
|
|
# Direct save of the complete workflow object
|
|
self.lucydom_interface.save_workflow_state(workflow)
|
|
logger.info(f"Workflow {workflow_id} created in database")
|
|
except Exception as e:
|
|
logger.error(f"Error creating workflow {workflow_id} in database: {str(e)}")
|
|
|
|
# Cache workflow in memory
|
|
self.workflows[workflow_id] = workflow
|
|
|
|
return workflow
|
|
|
|
async def stop_workflow(self, workflow_id: str) -> bool:
|
|
"""
|
|
Stop a running workflow.
|
|
|
|
Args:
|
|
workflow_id: ID of the workflow to stop
|
|
|
|
Returns:
|
|
True on success, False if workflow not found or already stopped
|
|
"""
|
|
try:
|
|
workflow = self.workflows.get(workflow_id)
|
|
|
|
if not workflow:
|
|
# Try to load the workflow
|
|
workflow = await self.load_workflow(workflow_id)
|
|
if not workflow:
|
|
return False
|
|
|
|
# If workflow is not running or completed, abort
|
|
if workflow.get("status") not in ["running", "completed"]:
|
|
return False
|
|
|
|
# Set status to stopped
|
|
workflow["status"] = "stopped"
|
|
workflow["last_activity"] = datetime.now().isoformat()
|
|
|
|
self._add_log(workflow, "Workflow was manually stopped", "info", "workflow", "Workflow Management")
|
|
|
|
# Save workflow
|
|
self._save_workflow(workflow)
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error stopping workflow {workflow_id}: {str(e)}")
|
|
return False
|
|
|
|
def _add_log(self, workflow: Dict[str, Any], message: str, log_type: str, agent_id: Optional[str] = None, agent_name: Optional[str] = None) -> None:
|
|
"""Add a log entry to the workflow."""
|
|
# First, check if workflow is a string (ID) instead of dictionary
|
|
if isinstance(workflow, str):
|
|
# Try to load the workflow by ID
|
|
workflow_id = workflow
|
|
workflow = self.workflows.get(workflow_id)
|
|
if not workflow:
|
|
# Just log to the logger and return
|
|
logger.info(f"Log (couldn't add to workflow {workflow_id}): {log_type} - {message}")
|
|
return
|
|
|
|
# Check if workflow is a dictionary
|
|
if not isinstance(workflow, dict):
|
|
logger.error(f"Invalid workflow type: {type(workflow)}. Expected dictionary.")
|
|
# Just log to the logger and return
|
|
logger.info(f"Log (couldn't add to workflow): {log_type} - {message}")
|
|
return
|
|
|
|
# Create log entry
|
|
log_entry = {
|
|
"id": f"log_{uuid.uuid4()}",
|
|
"message": message,
|
|
"type": log_type,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"agent_id": agent_id,
|
|
"agent_name": agent_name
|
|
}
|
|
|
|
# Add log entry to workflow
|
|
if "logs" not in workflow:
|
|
workflow["logs"] = []
|
|
|
|
workflow["logs"].append(log_entry)
|
|
|
|
# Update last activity
|
|
workflow["last_activity"] = log_entry["timestamp"]
|
|
|
|
# Save log entry to database if available
|
|
if self.lucydom_interface:
|
|
try:
|
|
# Add workflow ID to log entry
|
|
log_data = log_entry.copy()
|
|
log_data["workflow_id"] = workflow["id"]
|
|
|
|
self.lucydom_interface.create_workflow_log(log_data)
|
|
logger.debug(f"Log entry for workflow {workflow['id']} saved to database")
|
|
except Exception as e:
|
|
logger.error(f"Error saving log entry for workflow {workflow['id']} to database: {str(e)}")
|
|
|
|
# Also log to standard logger with the category prefix
|
|
category_prefix = f"[{agent_name or agent_id or 'Workflow'}]" if agent_name or agent_id else ""
|
|
log_message = f"{category_prefix} {message}"
|
|
|
|
if log_type == "error":
|
|
logger.error(log_message)
|
|
elif log_type == "warning":
|
|
logger.warning(log_message)
|
|
else:
|
|
logger.info(log_message)
|
|
|
|
def get_workflow_status(self, workflow_id: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get the status of a workflow.
|
|
|
|
Args:
|
|
workflow_id: ID of the workflow
|
|
|
|
Returns:
|
|
Dictionary with status information or None if workflow not found
|
|
"""
|
|
# Get from memory
|
|
workflow = self.workflows.get(workflow_id)
|
|
|
|
# If not in memory, load from database or file
|
|
if not workflow:
|
|
# Load from database if available
|
|
if self.lucydom_interface:
|
|
try:
|
|
workflow_data = self.lucydom_interface.get_workflow(workflow_id)
|
|
if workflow_data:
|
|
workflow = workflow_data
|
|
except Exception as e:
|
|
logger.error(f"Error loading workflow status from database: {str(e)}")
|
|
|
|
# If not in database, load from file
|
|
if not workflow:
|
|
try:
|
|
import json
|
|
workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
|
|
if os.path.exists(workflow_path):
|
|
with open(workflow_path, 'r', encoding='utf-8') as f:
|
|
workflow = json.load(f)
|
|
except Exception as e:
|
|
logger.error(f"Error loading workflow status from file: {str(e)}")
|
|
return None
|
|
|
|
if not workflow:
|
|
return None
|
|
|
|
# Extract status information
|
|
status_info = {
|
|
"id": workflow.get("id"),
|
|
"name": workflow.get("name", f"Workflow {workflow_id}"),
|
|
"status": workflow.get("status"),
|
|
"progress": 1.0 if workflow.get("status") in ["completed", "failed", "stopped"] else 0.5,
|
|
"started_at": workflow.get("started_at"),
|
|
"last_activity": workflow.get("last_activity"),
|
|
"workflow_complete": workflow.get("status") == "completed",
|
|
"current_round": workflow.get("current_round", 1),
|
|
"data_stats": workflow.get("data_stats", {
|
|
"total_processing_time": 0.0,
|
|
"total_token_count": 0,
|
|
"total_bytes_sent": 0,
|
|
"total_bytes_received": 0
|
|
})
|
|
}
|
|
|
|
return status_info
|
|
|
|
def get_workflow_logs(self, workflow_id: str) -> Optional[List[Dict[str, Any]]]:
|
|
"""
|
|
Get logs for a workflow.
|
|
|
|
Args:
|
|
workflow_id: ID of the workflow
|
|
|
|
Returns:
|
|
List of logs or None if workflow not found
|
|
"""
|
|
# Get from memory
|
|
workflow = self.workflows.get(workflow_id)
|
|
|
|
# If not in memory, load from database
|
|
if not workflow and self.lucydom_interface:
|
|
try:
|
|
logs = self.lucydom_interface.get_workflow_logs(workflow_id)
|
|
return logs
|
|
except Exception as e:
|
|
logger.error(f"Error loading workflow logs from database: {str(e)}")
|
|
|
|
# If not in database or no interface available, load from file
|
|
if not workflow:
|
|
try:
|
|
import json
|
|
workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
|
|
if os.path.exists(workflow_path):
|
|
with open(workflow_path, 'r', encoding='utf-8') as f:
|
|
workflow = json.load(f)
|
|
except Exception as e:
|
|
logger.error(f"Error loading workflow logs from file: {str(e)}")
|
|
return None
|
|
|
|
return workflow.get("logs", []) if workflow else None
|
|
|
|
def get_workflow_messages(self, workflow_id: str) -> Optional[List[Dict[str, Any]]]:
|
|
"""
|
|
Get messages for a workflow.
|
|
|
|
Args:
|
|
workflow_id: ID of the workflow
|
|
|
|
Returns:
|
|
List of messages or None if workflow not found
|
|
"""
|
|
# Get from memory
|
|
workflow = self.workflows.get(workflow_id)
|
|
|
|
# If not in memory, load from database
|
|
if not workflow and self.lucydom_interface:
|
|
try:
|
|
messages = self.lucydom_interface.get_workflow_messages(workflow_id)
|
|
return messages
|
|
except Exception as e:
|
|
logger.error(f"Error loading workflow messages from database: {str(e)}")
|
|
|
|
# If not in database or no interface available, load from file
|
|
if not workflow:
|
|
try:
|
|
import json
|
|
workflow_path = os.path.join(self.results_dir, f"workflow_{workflow_id}.json")
|
|
if os.path.exists(workflow_path):
|
|
with open(workflow_path, 'r', encoding='utf-8') as f:
|
|
workflow = json.load(f)
|
|
except Exception as e:
|
|
logger.error(f"Error loading workflow messages from file: {str(e)}")
|
|
return None
|
|
|
|
return workflow.get("messages", []) if workflow else None
|
|
|
|
# Factory function for WorkflowManager
|
|
def get_workflow_manager(mandate_id: int = None, user_id: int = None, ai_service = None):
|
|
"""
|
|
Get a WorkflowManager instance for the specified context.
|
|
Reuses existing instances.
|
|
|
|
Args:
|
|
mandate_id: Mandate ID
|
|
user_id: User ID
|
|
ai_service: AI service
|
|
|
|
Returns:
|
|
WorkflowManager instance
|
|
"""
|
|
from modules.lucydom_interface import get_lucydom_interface
|
|
|
|
context_key = f"{mandate_id}_{user_id}"
|
|
|
|
# LucyDOM interface for database access
|
|
lucydom_interface = get_lucydom_interface(mandate_id, user_id)
|
|
|
|
if context_key not in _workflow_managers:
|
|
_workflow_managers[context_key] = WorkflowManager(
|
|
mandate_id,
|
|
user_id,
|
|
ai_service,
|
|
lucydom_interface
|
|
)
|
|
|
|
# Update services if changed
|
|
if ai_service is not None:
|
|
_workflow_managers[context_key].ai_service = ai_service
|
|
|
|
return _workflow_managers[context_key]
|
|
|
|
# Singleton factory for WorkflowManager instances per context
|
|
_workflow_managers = {} |