gateway/gwserver/modules/agentservice_workflow_execution.py
2025-04-15 01:04:38 +02:00

660 lines
No EOL
27 KiB
Python

"""
Refactored architecture for the Agentservice multi-agent system.
This module defines the revised workflow execution with improved agent handovers.
"""
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__)
logging.getLogger('matplotlib.font_manager').setLevel(logging.INFO)
class WorkflowExecution:
"""
Handles the execution of workflows with improved agent collaboration.
Integrates planning and execution phases for better context awareness.
"""
def __init__(self, workflow_manager, workflow_id: str, mandate_id: int, user_id: int, ai_service, lucydom_interface):
"""Initialize the workflow execution"""
self.workflow_manager = workflow_manager
self.workflow_id = workflow_id
self.mandate_id = mandate_id
self.user_id = user_id
self.ai_service = ai_service
self.lucydom_interface = lucydom_interface
# Import necessary modules
from modules.agentservice_utils import WorkflowUtils, MessageUtils, LoggingUtils
from modules.agentservice_registry import AgentRegistry
from modules.agentservice_filemanager import get_workflow_file_manager
# Initialize utilities
self.workflow_utils = WorkflowUtils(workflow_id)
self.message_utils = MessageUtils()
self.logging_utils = LoggingUtils(workflow_id, self._add_log)
# Initialize agent registry
self.agent_registry = AgentRegistry.get_instance()
# Set dependencies for agents
# Initialize file manager
self.file_manager = get_workflow_file_manager(workflow_id, lucydom_interface)
# Import and initialize document handler
from modules.agentservice_document_handler import get_document_handler
self.document_handler = get_document_handler(workflow_id, lucydom_interface, ai_service)
self.agent_registry.set_dependencies(
ai_service=ai_service,
document_handler=self.document_handler,
lucydom_interface=lucydom_interface
)
async def execute(self, message: Dict[str, Any], workflow: Dict[str, Any], files: List[Dict[str, Any]] = None, is_user_input: bool = False):
"""
Execute the workflow with integrated planning and agent selection.
Args:
message: The initiating message (prompt or user input)
workflow: The workflow object
files: Optional list of file metadata
is_user_input: Flag indicating if this is user input
Returns:
Dict with workflow status and result
"""
try:
# 1. Initialize workflow logging
self.logging_utils.info("Starting workflow execution", "workflow", "Workflow initialized")
# 2. Process user message and files
user_message = await self._process_user_message(workflow, message, files)
self.logging_utils.info("User message processed", "workflow", "User input added to workflow")
# 3. Create agent-aware work plan
work_plan = await self._create_agent_aware_work_plan(workflow, user_message)
self.logging_utils.info(f"Created agent-aware work plan with {len(work_plan)} activities", "planning")
self.logging_utils.debug(f"{work_plan}.", "planning")
# 4. Execute the activities in the work plan
results = await self._execute_work_plan(workflow, work_plan)
# 5. Create summary
summary = await self._create_summary(workflow, results)
self.logging_utils.info("Created workflow summary", "summary")
# Set workflow status to completed
workflow["status"] = "completed"
workflow["last_activity"] = datetime.now().isoformat()
# Final save
self.workflow_manager._save_workflow(workflow)
return {
"workflow_id": self.workflow_id,
"status": "completed",
"messages": workflow.get("messages", [])
}
except Exception as e:
self.logging_utils.error(f"Workflow execution failed: {str(e)}", "error")
workflow["status"] = "failed"
self.workflow_manager._save_workflow(workflow)
return {
"workflow_id": self.workflow_id,
"status": "failed",
"error": str(e)
}
async def _process_user_message(self, workflow: Dict[str, Any], message: Dict[str, Any], files: List[Dict[str, Any]] = None) -> Dict[str, Any]:
"""
Process the user message and add it to the workflow.
Args:
workflow: The workflow object
message: The user message
files: Optional list of file metadata
Returns:
The processed user message
"""
# Create a message with user input
user_message = self._create_message(workflow, message.get("role", "user"))
user_message["content"] = message.get("content", "")
# Process files if provided
if files and len(files) > 0:
self.logging_utils.info(f"Processing {len(files)} files", "files")
# Add files to message via file manager instead of _process_files
user_message = await self.file_manager.add_files_to_message(
user_message,
[f.get('id') for f in files],
self._add_log
)
# Add the message to the workflow
if "messages" not in workflow:
workflow["messages"] = []
workflow["messages"].append(user_message)
# Save workflow state
self.workflow_manager._save_workflow(workflow)
return user_message
async def _create_agent_aware_work_plan(self, workflow: Dict[str, Any], message: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Create an agent-aware work plan that integrates agent selection during planning.
Args:
workflow: The workflow object
message: The initiating message
Returns:
List of structured activities with agent assignments
"""
# Extract context information
task = message.get("content", "")
# Get all available agents and their capabilities
agent_infos = self.agent_registry.get_agent_infos()
# Extract documents
documents = message.get("documents", [])
document_info = []
for doc in documents:
source = doc.get("source", {})
document_info.append({
"id": doc.get("id"),
"name": source.get("name", "unnamed"),
"type": source.get("type", "unknown"),
"content_type": source.get("content_type", "unknown")
})
# Create the planning prompt with agent awareness and document handling information
plan_prompt = f"""
As an AI workflow manager, create a detailed agent-aware work plan for the following task:
TASK: {task}
AVAILABLE AGENTS:
{self._format_agent_info(agent_infos)}
AVAILABLE DOCUMENTS:
{document_info if document_info else "No documents provided"}
IMPORTANT: Document extraction happens automatically in the workflow. Documents in the message are already available to all agents. DO NOT assign agent_coder or any other agent specifically for just reading or extracting document content. Only assign agents for tasks that require specific processing beyond what the document handler already provides.
The work plan should include a structured list of activities. Each activity should have:
1. title - A short descriptive title for the activity
2. description - What needs to be done in this activity
3. assigned_agents - List of agent IDs that should handle this activity (can be multiple in sequence)
4. agent_prompts - Specific instructions for each agent (matched by index to assigned_agents)
5. document_requirements - Description of which documents are needed for this activity (these will be automatically extracted)
6. expected_output - The expected output format and content
7. dependencies - List of previous activities this depends on (by index)
IMPORTANT GUIDELINES:
- Each activity should have clear objectives and be assigned to the most appropriate agent(s)
- When multiple agents are assigned to an activity, specify the sequence and how outputs should flow between them
- Documents are processed on-demand by the system's document handler, so only specify which documents are needed, not how to extract them
- DO NOT create activities that only read or extract document content - this happens automatically
- Create a logical sequence where later activities can use outputs from earlier ones
- If no specialized agent is needed for a task, use the default "assistant" agent
- Only use the agent_coder for tasks that require actual coding or complex data analysis, not for simply reading documents
Return the work plan as a JSON array of activity objects, each with the above properties.
"""
self.logging_utils.info("Creating agent-aware work plan", "planning")
# Call AI to generate work plan
try:
plan_response = await self.ai_service.call_api([{"role": "user", "content": plan_prompt}])
# Extract JSON plan
import json
import re
# Look for JSON array in the response
json_pattern = r'\[\s*\{.*\}\s*\]'
json_match = re.search(json_pattern, plan_response, re.DOTALL)
if json_match:
json_str = json_match.group(0)
work_plan = json.loads(json_str)
self.logging_utils.info(f"Work plan created with {len(work_plan)} activities", "planning")
return work_plan
else:
self.logging_utils.warning("Could not extract JSON from AI response", "planning")
# Fallback: Create a simple default work plan
return [{
"title": "Process Task",
"description": "Process the user's request directly",
"assigned_agents": ["assistant"],
"agent_prompts": [task],
"document_requirements": "All available documents may be needed",
"expected_output": "Text",
"dependencies": []
}]
except Exception as e:
self.logging_utils.error(f"Error creating work plan: {str(e)}", "planning")
# Return a minimal fallback plan
return [{
"title": "Process Task (Error Recovery)",
"description": "Process the user's request after planning error",
"assigned_agents": ["assistant"],
"agent_prompts": [task],
"document_requirements": "All available documents may be needed",
"expected_output": "Text",
"dependencies": []
}]
async def _execute_work_plan(self, workflow: Dict[str, Any], work_plan: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Execute all activities in the work plan with proper agent handovers.
Args:
workflow: The workflow object
work_plan: The work plan with activities
Returns:
Results from all activities
"""
results = []
activity_outputs = {} # Store outputs for dependency resolution
for activity_index, activity in enumerate(work_plan):
# Extract activity info
title = activity.get("title", f"Activity {activity_index+1}")
description = activity.get("description", "")
assigned_agents = activity.get("assigned_agents", ["assistant"])
agent_prompts = activity.get("agent_prompts", [description])
doc_requirements = activity.get("document_requirements", "")
expected_output = activity.get("expected_output", "Text")
dependencies = activity.get("dependencies", [])
self.logging_utils.info(f"Starting activity: {title}", "execution")
# Validate assigned_agents and agent_prompts
if len(assigned_agents) > len(agent_prompts):
# Duplicate the last prompt for additional agents
agent_prompts.extend([agent_prompts[-1]] * (len(assigned_agents) - len(agent_prompts)))
elif len(agent_prompts) > len(assigned_agents):
# Truncate excess prompts
agent_prompts = agent_prompts[:len(assigned_agents)]
# Process dependencies first
dependency_context = {}
for dep_index in dependencies:
if dep_index < activity_index and dep_index in activity_outputs:
dep_output = activity_outputs[dep_index]
dependency_context[f"activity_{dep_index+1}"] = dep_output
# Extract required documents if needed
document_content = ""
if doc_requirements:
extracted_data = await self._extract_required_documents(workflow, doc_requirements)
if extracted_data and "extracted_content" in extracted_data:
# Format document content for the prompt
document_content = "\n\n=== EXTRACTED DOCUMENT CONTENT ===\n\n"
for item in extracted_data.get("extracted_content", []):
doc_name = item.get("name", "Unnamed document")
doc_content = item.get("content", "No content available")
document_content += f"--- {doc_name} ---\n{doc_content}\n\n"
# Execute the activity with the assigned agents
activity_result = await self._execute_agent_sequence(
workflow,
assigned_agents,
agent_prompts,
document_content,
dependency_context,
expected_output
)
# Store the result
activity_outputs[activity_index] = activity_result
results.append({
"title": title,
"description": description,
"agents": assigned_agents,
"result": activity_result.get("content", ""),
"output_format": activity_result.get("format", "Text")
})
self.logging_utils.info(f"Completed activity: {title}", "execution")
# Save intermediate state
self.workflow_manager._save_workflow(workflow)
return results
async def _execute_agent_sequence(
self,
workflow: Dict[str, Any],
agent_ids: List[str],
prompts: List[str],
document_content: str,
dependency_context: Dict[str, Any],
expected_output: str
) -> Dict[str, Any]:
"""
Execute a sequence of agents with proper handovers.
Args:
workflow: The workflow object
agent_ids: List of agent IDs to execute in sequence
prompts: List of prompts for each agent
document_content: Extracted document content
dependency_context: Context from dependent activities
expected_output: Expected output format
Returns:
Result of the agent sequence execution
"""
context = {
"workflow_id": self.workflow_id,
"expected_format": expected_output,
"dependency_outputs": dependency_context
}
last_result = None
last_documents = []
for i, agent_id in enumerate(agent_ids):
# Get the agent
agent = self.agent_registry.get_agent(agent_id)
if agent:
# Ensure dependencies are set
if hasattr(agent, 'set_dependencies'):
agent.set_dependencies(
ai_service=self.ai_service,
document_handler=self.document_handler,
lucydom_interface=self.lucydom_interface
)
# Set document handler if agent supports it
if hasattr(agent, 'set_document_handler') and hasattr(self, 'document_handler'):
agent.set_document_handler(self.document_handler)
if not agent:
self.logging_utils.warning(f"Agent '{agent_id}' not found, using assistant instead", "agents")
agent = self.agent_registry.get_agent("assistant")
if not agent:
# If assistant not found, create a minimal agent response
continue
# Get the agent prompt
base_prompt = prompts[i] if i < len(prompts) else prompts[-1]
# Enhance the prompt with context
enhanced_prompt = self._enhance_prompt(
base_prompt,
document_content,
dependency_context,
last_result.get("content", "") if last_result else "",
i > 0 # is_continuation flag
)
# Create the message for this agent
agent_message = self._create_message(workflow, "user")
agent_message["content"] = enhanced_prompt
# IMPORTANT FIX: Document handling logic
# First, check if we have documents from previous agent if this is a continuation
if last_documents and i > 0:
agent_message["documents"] = last_documents
# For the first agent, make sure we pass any documents from the most recent user message
elif i == 0:
# Find the most recent user message with documents
for msg in reversed(workflow.get("messages", [])):
if msg.get("role") == "user" and msg.get("documents"):
agent_message["documents"] = msg.get("documents", [])
self.logging_utils.info(f"Passing {len(agent_message['documents'])} documents from user message to {agent_id}", "agents")
break
# Log agent execution
self.logging_utils.info(f"Executing agent: {agent_id}", "agents")
# Execute the agent
agent_response = await agent.process_message(agent_message, context)
# Create response message
response_message = self._create_message(workflow, "assistant")
response_message["content"] = agent_response.get("content", "")
response_message["agent_type"] = agent_id
response_message["agent_id"] = agent_id
response_message["agent_name"] = agent.name
response_message["result_format"] = agent_response.get("result_format", expected_output)
# Capture documents from response
if "documents" in agent_response:
response_message["documents"] = agent_response["documents"]
last_documents = agent_response["documents"]
self.logging_utils.info(f"Agent {agent_id} produced {len(last_documents)} documents", "agents")
# Add to workflow
workflow["messages"].append(response_message)
# Update last result
last_result = {
"content": agent_response.get("content", ""),
"format": agent_response.get("result_format", expected_output),
"agent_id": agent_id,
"documents": agent_response.get("documents", [])
}
return last_result or {
"content": "No agent response was generated.",
"format": "Text"
}
async def _extract_required_documents(self, workflow: Dict[str, Any], doc_requirements: str) -> Dict[str, Any]:
"""
Extract required documents based on requirements description.
Args:
workflow: The workflow object
doc_requirements: Description of document requirements
Returns:
Extracted document data
"""
# Import for data extraction
from modules.agentservice_dataextraction import data_extraction
# Get all files from the workflow
files = self.workflow_utils.get_files(workflow)
# Get all messages from the workflow
workflow_messages = workflow.get("messages", [])
# Extract data using the dataextraction module
extracted_data = await data_extraction(
prompt=doc_requirements,
files=files,
messages=workflow_messages,
ai_service=self.ai_service,
lucydom_interface=self.lucydom_interface,
workflow_id=self.workflow_id,
add_log_func=self._add_log
)
return extracted_data
async def _create_summary(self, workflow: Dict[str, Any], results: List[Dict[str, Any]]) -> Dict[str, Any]:
"""
Create a summary of the workflow results for the user.
Args:
workflow: The workflow object
results: Results from activity executions
Returns:
Summary message
"""
# Create a summary prompt
summary_prompt = "Create a clear, concise summary of the following workflow results:\n\n"
for i, result in enumerate(results, 1):
title = result.get("title", f"Activity {i}")
description = result.get("description", "")
content = result.get("result", "")
agents = ", ".join(result.get("agents", ["unknown"]))
# Limit content length for the summary prompt
content_preview = content[:500] + "..." if len(content) > 500 else content
summary_prompt += f"""
ACTIVITY {i}: {title}
Description: {description}
Executed by: {agents}
{content_preview}
---
"""
summary_prompt += """
Provide a well-structured summary that:
1. Highlights the key findings and results
2. Connects the results to the original task
3. Presents any conclusions or recommendations
Make sure the summary is clear, concise, and useful to the user.
"""
# Call AI to generate summary
summary_content = await self.ai_service.call_api([{"role": "user", "content": summary_prompt}])
# Create summary message
summary_message = self._create_message(workflow, "assistant")
summary_message["content"] = summary_content
summary_message["agent_type"] = "summary"
summary_message["agent_id"] = "workflow_summary"
summary_message["agent_name"] = "Workflow Summary"
summary_message["result_format"] = "Text"
summary_message["workflow_complete"] = True
# Add to workflow
workflow["messages"].append(summary_message)
return summary_message
def _create_message(self, workflow: Dict[str, Any], role: str) -> Dict[str, Any]:
"""Create a new message object for the workflow"""
message_id = f"msg_{uuid.uuid4()}"
current_time = datetime.now().isoformat()
# Determine sequence number
sequence_no = 1
if "messages" in workflow and workflow["messages"]:
sequence_no = len(workflow["messages"]) + 1
# Create message object
message = {
"id": message_id,
"workflow_id": self.workflow_id,
"parent_message_id": None,
"started_at": current_time,
"finished_at": None,
"sequence_no": sequence_no,
"status": "pending",
"role": role,
"data_stats": {
"processing_time": 0.0,
"token_count": 0,
"bytes_sent": 0,
"bytes_received": 0
},
"documents": [],
"content": None,
"agent_type": None
}
return message
def _add_log(self, workflow_id: str, message: str, log_type: str, agent_id: str = None, agent_name: str = None):
"""Add a log entry to the workflow"""
# This calls back to the workflow manager's log function
self.workflow_manager._add_log(workflow_id, message, log_type, agent_id, agent_name)
def _format_agent_info(self, agent_infos: List[Dict[str, Any]]) -> str:
"""Format agent information for the planning prompt"""
formatted_info = ""
for agent in agent_infos:
formatted_info += f"""
- ID: {agent.get('id', 'unknown')}
Name: {agent.get('name', '')}
Type: {agent.get('type', '')}
Description: {agent.get('description', '')}
Capabilities: {agent.get('capabilities', '')}
Result Format: {agent.get('result_format', 'Text')}
"""
return formatted_info
def _enhance_prompt(
self,
base_prompt: str,
document_content: str,
dependency_context: Dict[str, Any],
previous_result: str,
is_continuation: bool
) -> str:
"""
Enhance a prompt with context information.
Args:
base_prompt: The original prompt
document_content: Extracted document content
dependency_context: Context from dependent activities
previous_result: Result from previous agent in sequence
is_continuation: Flag indicating if this is a continuation
Returns:
Enhanced prompt
"""
enhanced_prompt = base_prompt
# Add continuation context if this is a continuation
if is_continuation and previous_result:
enhanced_prompt = f"""
{enhanced_prompt}
=== PREVIOUS AGENT OUTPUT ===
{previous_result}
"""
# Add document content if available
if document_content:
enhanced_prompt += f"\n\n{document_content}"
# Add dependency context if available
if dependency_context:
dependency_section = "\n\n=== OUTPUTS FROM PREVIOUS ACTIVITIES ===\n\n"
for name, value in dependency_context.items():
if isinstance(value, dict) and "content" in value:
# Extract content if it's in the standard format
dependency_section += f"--- {name} ---\n{value['content']}\n\n"
else:
# Use the value directly
dependency_section += f"--- {name} ---\n{str(value)}\n\n"
enhanced_prompt += dependency_section
return enhanced_prompt