499 lines
No EOL
19 KiB
Python
499 lines
No EOL
19 KiB
Python
import logging
|
|
import importlib
|
|
import pkgutil
|
|
import inspect
|
|
from typing import Dict, Any, Optional, List, Type, Callable, Awaitable
|
|
from datetime import datetime, UTC
|
|
import json
|
|
import asyncio
|
|
import base64
|
|
|
|
from modules.methods.methodBase import MethodBase, AuthSource, MethodResult
|
|
from modules.workflow.serviceContainer import ServiceContainer
|
|
from modules.interfaces.serviceChatModel import (
|
|
AgentTask, AgentAction, AgentResult, Action, TaskStatus, ChatWorkflow,
|
|
ChatMessage, ChatDocument, ChatStat, ExtractedContent, ContentItem
|
|
)
|
|
from modules.workflow.processorDocument import DocumentProcessor
|
|
from modules.shared.configuration import APP_CONFIG
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ChatManager:
|
|
"""Chat manager with improved AI integration and method handling"""
|
|
|
|
def __init__(self):
|
|
self.service = ServiceContainer()
|
|
self._discoverMethods()
|
|
self.workflow: Optional[ChatWorkflow] = None
|
|
self.currentTask: Optional[AgentTask] = None
|
|
self.workflowHistory: List[ChatMessage] = []
|
|
self.documentProcessor = DocumentProcessor()
|
|
|
|
def _discoverMethods(self):
|
|
"""Dynamically discover all method classes in modules.methods package"""
|
|
try:
|
|
# Import the methods package
|
|
methodsPackage = importlib.import_module('modules.methods')
|
|
|
|
# Discover all modules in the package
|
|
for _, name, isPkg in pkgutil.iter_modules(methodsPackage.__path__):
|
|
if not isPkg and name.startswith('method'):
|
|
try:
|
|
# Import the module
|
|
module = importlib.import_module(f'modules.methods.{name}')
|
|
|
|
# Find all classes in the module that inherit from MethodBase
|
|
for itemName, item in inspect.getmembers(module):
|
|
if (inspect.isclass(item) and
|
|
issubclass(item, MethodBase) and
|
|
item != MethodBase):
|
|
# Instantiate the method and add to service
|
|
methodInstance = item()
|
|
self.service.methods[methodInstance.name] = methodInstance
|
|
logger.info(f"Discovered method: {methodInstance.name}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error loading method module {name}: {str(e)}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error discovering methods: {str(e)}")
|
|
|
|
async def initialize(self, workflow: ChatWorkflow) -> None:
|
|
"""Initialize chat manager with workflow"""
|
|
self.service.workflow = workflow
|
|
|
|
# Initialize AI model
|
|
self.service.model = {
|
|
'callAiBasic': self._callAiBasic,
|
|
'callAiAdvanced': self._callAiAdvanced
|
|
}
|
|
|
|
# Initialize document processor
|
|
self.service.documentProcessor.initialize()
|
|
|
|
def _generatePrompt(self, task: str, document: ChatDocument, examples: List[Dict[str, str]] = None) -> str:
|
|
"""Generate a prompt based on task and document"""
|
|
try:
|
|
# Create base prompt
|
|
prompt = f"""Task: {task}
|
|
Document: {document.filename} ({document.mimeType})
|
|
|
|
"""
|
|
|
|
# Add examples if provided
|
|
if examples:
|
|
prompt += "\nExamples:\n"
|
|
for example in examples:
|
|
prompt += f"Input: {example.get('input', '')}\n"
|
|
prompt += f"Output: {example.get('output', '')}\n\n"
|
|
|
|
return prompt
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating prompt: {str(e)}")
|
|
return ""
|
|
|
|
async def createInitialTask(self, userInput: Dict[str, Any]) -> AgentTask:
|
|
"""Create initial task from user input"""
|
|
# Get available methods and their actions
|
|
methodCatalog = self.service.getAvailableMethods()
|
|
|
|
# Process user input with AI
|
|
processedInput = await self._processUserInput(userInput, methodCatalog)
|
|
|
|
# Create actions from processed input
|
|
actions = await self._createActions(processedInput['actions'])
|
|
|
|
# Create task
|
|
task = AgentTask(
|
|
id=f"task_{datetime.now(UTC).timestamp()}",
|
|
workflowId=self.workflow.id,
|
|
userInput=processedInput['objective'],
|
|
dataList=userInput.get('connections', []),
|
|
actionList=actions,
|
|
status=TaskStatus.PENDING,
|
|
createdAt=datetime.now(UTC),
|
|
updatedAt=datetime.now(UTC)
|
|
)
|
|
|
|
# Store in service
|
|
self.service.tasks['current'] = task
|
|
return task
|
|
|
|
async def executeCurrentTask(self) -> None:
|
|
"""Execute current task"""
|
|
task = self.service.tasks.get('current')
|
|
if not task:
|
|
raise ValueError("No current task to execute")
|
|
|
|
await self.service.executeTask(task)
|
|
|
|
async def defineNextTask(self) -> Optional[AgentTask]:
|
|
"""Define next task based on current task results"""
|
|
current_task = self.service.tasks.get('current')
|
|
if not current_task:
|
|
return None
|
|
|
|
try:
|
|
# Analyze task results
|
|
analysis = await self._analyzeTaskResults(current_task)
|
|
|
|
# If workflow is complete, update task status
|
|
if analysis['isComplete']:
|
|
current_task.status = TaskStatus.COMPLETED
|
|
current_task.updatedAt = datetime.now(UTC)
|
|
return None
|
|
|
|
# If more actions needed, create next task
|
|
if not analysis['isComplete']:
|
|
next_task = self._createNextTask(current_task, analysis)
|
|
self.service.tasks['previous'] = current_task
|
|
self.service.tasks['current'] = next_task
|
|
return next_task
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error defining next task: {e}")
|
|
current_task.status = TaskStatus.FAILED
|
|
current_task.updatedAt = datetime.now(UTC)
|
|
return None
|
|
|
|
async def _processUserInput(self, userInput: Dict[str, Any], methodCatalog: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Process user input with AI to extract objectives and actions"""
|
|
# Create prompt with available methods and actions
|
|
prompt = f"""Given the following user input and available methods/actions, extract the objective and required actions:
|
|
|
|
User Input: {userInput.get('message', '')}
|
|
|
|
Available Methods and Actions:
|
|
{json.dumps(methodCatalog, indent=2)}
|
|
|
|
Please provide a JSON response with:
|
|
1. objective: The main goal or task to accomplish
|
|
2. actions: List of required actions with method and parameters
|
|
|
|
Example format:
|
|
{{
|
|
"objective": "Search for documents about project X",
|
|
"actions": [
|
|
{{
|
|
"method": "sharepoint",
|
|
"action": "search",
|
|
"parameters": {{
|
|
"query": "project X",
|
|
"site": "projects"
|
|
}}
|
|
}}
|
|
]
|
|
}}
|
|
"""
|
|
|
|
# Call AI service
|
|
response = await self._callAiBasic(prompt)
|
|
return json.loads(response)
|
|
|
|
async def _createActions(self, actionsData: List[Dict[str, Any]]) -> List[AgentAction]:
|
|
"""Create action objects from processed input"""
|
|
actions = []
|
|
for actionData in actionsData:
|
|
method = self.service.getMethod(actionData['method'])
|
|
if not method:
|
|
continue
|
|
|
|
action = AgentAction(
|
|
id=f"action_{datetime.now(UTC).timestamp()}",
|
|
method=actionData['method'],
|
|
action=actionData['action'],
|
|
parameters=actionData.get('parameters', {}),
|
|
status=TaskStatus.PENDING,
|
|
createdAt=datetime.now(UTC),
|
|
updatedAt=datetime.now(UTC)
|
|
)
|
|
actions.append(action)
|
|
|
|
return actions
|
|
|
|
async def _summarizeWorkflow(self) -> str:
|
|
"""Summarize workflow history"""
|
|
if not self.workflow.messages:
|
|
return ""
|
|
|
|
prompt = f"""Summarize the following chat history:
|
|
{json.dumps([m.dict() for m in self.workflow.messages], indent=2)}
|
|
|
|
Please provide a concise summary focusing on:
|
|
1. Main objectives
|
|
2. Key actions taken
|
|
3. Current status
|
|
4. Any issues or blockers
|
|
"""
|
|
|
|
return await self._callAiBasic(prompt)
|
|
|
|
async def _analyzeTaskResults(self, task: AgentTask) -> Dict[str, Any]:
|
|
"""Analyze task results to determine next steps"""
|
|
# Get workflow summary
|
|
summary = await self._summarizeWorkflow()
|
|
|
|
# Create prompt for analysis
|
|
prompt = f"""Analyze the following task results and workflow history to determine next steps:
|
|
|
|
Task Results:
|
|
{json.dumps([a.dict() for a in task.actionList], indent=2)}
|
|
|
|
Workflow Summary:
|
|
{summary}
|
|
|
|
Please provide a JSON response with:
|
|
1. isComplete: Whether the workflow is complete
|
|
2. nextActions: List of next actions needed (if any)
|
|
3. issues: Any issues or blockers identified
|
|
|
|
Example format:
|
|
{{
|
|
"isComplete": false,
|
|
"nextActions": [
|
|
{{
|
|
"method": "sharepoint",
|
|
"action": "read",
|
|
"parameters": {{
|
|
"documentId": "doc123"
|
|
}}
|
|
}}
|
|
],
|
|
"issues": ["Need authentication for SharePoint"]
|
|
}}
|
|
"""
|
|
|
|
response = await self._callAiBasic(prompt)
|
|
return json.loads(response)
|
|
|
|
def _createNextTask(self, current_task: AgentTask, analysis: Dict[str, Any]) -> AgentTask:
|
|
"""Create next task based on analysis"""
|
|
# Create actions for next task
|
|
actions = []
|
|
for action_data in analysis.get('nextActions', []):
|
|
action = AgentAction(
|
|
id=f"action_{datetime.now(UTC).timestamp()}",
|
|
method=action_data['method'],
|
|
action=action_data['action'],
|
|
parameters=action_data.get('parameters', {}),
|
|
status=TaskStatus.PENDING,
|
|
createdAt=datetime.now(UTC),
|
|
updatedAt=datetime.now(UTC)
|
|
)
|
|
actions.append(action)
|
|
|
|
# Create and return next task
|
|
return AgentTask(
|
|
id=f"task_{datetime.now(UTC).timestamp()}",
|
|
workflowId=self.workflow.id,
|
|
userInput=current_task.userInput,
|
|
dataList=current_task.dataList,
|
|
actionList=actions,
|
|
status=TaskStatus.PENDING,
|
|
createdAt=datetime.now(UTC),
|
|
updatedAt=datetime.now(UTC)
|
|
)
|
|
|
|
async def processTask(self, task: AgentTask) -> Dict[str, Any]:
|
|
"""Process a task with improved error handling and AI integration"""
|
|
try:
|
|
# Execute task
|
|
await self.service.executeTask(task)
|
|
|
|
# Process results
|
|
if task.status == TaskStatus.COMPLETED:
|
|
# Generate feedback using AI
|
|
feedback = await self._processTaskResults(task)
|
|
task.thisTaskFeedback = feedback
|
|
|
|
# Create output documents
|
|
documents = await self._createOutputDocuments(task)
|
|
task.documentsOutput = documents
|
|
|
|
return {
|
|
"status": "success",
|
|
"feedback": feedback,
|
|
"documents": documents
|
|
}
|
|
else:
|
|
return {
|
|
"status": task.status,
|
|
"error": task.error,
|
|
"feedback": f"Task failed: {task.error}"
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing task: {str(e)}")
|
|
return {
|
|
"status": "error",
|
|
"error": str(e),
|
|
"feedback": f"Error processing task: {str(e)}"
|
|
}
|
|
|
|
def _generateDocumentPrompt(self, task: str) -> str:
|
|
"""Generate a prompt for document generation"""
|
|
return f"""Generate output documents for the following task:
|
|
|
|
Task: {task}
|
|
|
|
For each document you need to generate, provide a TaskDocument object with the following structure:
|
|
{{
|
|
"filename": "string", # Filename with extension
|
|
"mimeType": "string", # MIME type of the file
|
|
"data": "string", # File content as text or base64
|
|
"base64Encoded": boolean # True if data is base64 encoded
|
|
}}
|
|
|
|
Rules:
|
|
1. For text files (txt, json, xml, etc.), provide content directly in the data field
|
|
2. For binary files (images, videos, etc.), encode content in base64 and set base64Encoded to true
|
|
3. Use appropriate MIME types (e.g., text/plain, image/jpeg, application/pdf)
|
|
4. Include file extensions in filenames
|
|
|
|
Return a JSON array of TaskDocument objects.
|
|
"""
|
|
|
|
async def _processTaskResults(self, task: AgentTask) -> str:
|
|
"""Process task results and generate feedback"""
|
|
try:
|
|
# Generate document prompt
|
|
docPrompt = self._generateDocumentPrompt(task.userInput)
|
|
|
|
# Get AI response for document generation
|
|
docResponse = await self._callAiBasic(docPrompt)
|
|
|
|
# Parse response into TaskDocument objects
|
|
try:
|
|
taskDocs = json.loads(docResponse)
|
|
task.documentsOutput = taskDocs
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Error parsing document response: {str(e)}")
|
|
return f"Error processing results: {str(e)}"
|
|
|
|
# Generate feedback
|
|
feedback = await self._callAiBasic(
|
|
f"""Generate feedback for the completed task:
|
|
Task: {task.userInput}
|
|
Generated Documents: {len(task.documentsOutput)} files
|
|
|
|
Provide a concise summary of what was accomplished.
|
|
"""
|
|
)
|
|
|
|
return feedback
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing task results: {str(e)}")
|
|
return f"Error processing results: {str(e)}"
|
|
|
|
async def _createOutputDocuments(self, task: AgentTask) -> List[ChatDocument]:
|
|
"""Create output documents from task results"""
|
|
try:
|
|
fileIds = []
|
|
|
|
# Process each TaskDocument from AI output
|
|
for taskDoc in task.documentsOutput:
|
|
# Store file in database
|
|
fileItem = self.service.functions.createFile(
|
|
name=taskDoc.filename,
|
|
mimeType=taskDoc.mimeType
|
|
)
|
|
|
|
# Store file content
|
|
if taskDoc.base64Encoded:
|
|
# Decode base64 content
|
|
content = base64.b64decode(taskDoc.data)
|
|
else:
|
|
# Use text content directly
|
|
content = taskDoc.data.encode('utf-8')
|
|
|
|
# Store file data
|
|
self.service.functions.createFileData(fileItem.id, content)
|
|
fileIds.append(fileItem.id)
|
|
|
|
# Convert all files to ChatDocuments in one call
|
|
if fileIds:
|
|
return await self.service.chat.processFileIds(fileIds)
|
|
return []
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating output documents: {str(e)}")
|
|
return []
|
|
|
|
async def _callAiBasic(self, prompt: str) -> str:
|
|
"""Call basic AI service"""
|
|
try:
|
|
if not self.service or not self.service.base:
|
|
raise ValueError("Service or base interface not initialized")
|
|
return await self.service.base.callAi([
|
|
{"role": "system", "content": "You are an AI assistant that helps process user requests."},
|
|
{"role": "user", "content": prompt}
|
|
])
|
|
except Exception as e:
|
|
logger.error(f"Error calling AI service: {str(e)}")
|
|
raise
|
|
|
|
async def _callAiAdvanced(self, prompt: str, context: Dict[str, Any]) -> str:
|
|
"""Call advanced AI model with context"""
|
|
# TODO: Implement actual AI call
|
|
return "AI response placeholder"
|
|
|
|
async def generateWorkflowFeedback(self, workflow: ChatWorkflow) -> str:
|
|
"""
|
|
Generates a final feedback message for the workflow in the user's language.
|
|
|
|
Args:
|
|
workflow: The completed workflow to generate feedback for
|
|
|
|
Returns:
|
|
str: The generated feedback message
|
|
"""
|
|
try:
|
|
# Get workflow summary
|
|
workflowSummary = {
|
|
"status": workflow.status,
|
|
"totalMessages": len(workflow.messages),
|
|
"totalDocuments": sum(len(msg.documents) for msg in workflow.messages),
|
|
"duration": (datetime.now(UTC) - datetime.fromisoformat(workflow.startedAt)).total_seconds()
|
|
}
|
|
|
|
# Get user language from workflow mandate
|
|
userLanguage = workflow.mandateId.split('_')[0] if workflow.mandateId else 'en'
|
|
|
|
# Prepare messages for AI context
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": f"You are an AI assistant providing a summary of a completed workflow. "
|
|
f"Please respond in '{userLanguage}' language. "
|
|
f"Summarize the workflow's activities, outcomes, and any important points. "
|
|
f"Be concise but informative. Use a professional but friendly tone."
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": f"Please provide a summary of this workflow:\n"
|
|
f"Status: {workflowSummary['status']}\n"
|
|
f"Total Messages: {workflowSummary['totalMessages']}\n"
|
|
f"Total Documents: {workflowSummary['totalDocuments']}\n"
|
|
f"Duration: {workflowSummary['duration']:.1f} seconds"
|
|
}
|
|
]
|
|
|
|
# Add relevant workflow messages for context
|
|
for msg in workflow.messages:
|
|
if msg.role == "user" or msg.status in ["first", "last"]:
|
|
messages.append({
|
|
"role": msg.role,
|
|
"content": msg.message
|
|
})
|
|
|
|
# Generate feedback using AI
|
|
feedback = await self.service.aiService.callApi(messages, temperature=0.7)
|
|
|
|
return feedback
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating workflow feedback: {str(e)}")
|
|
return "Workflow completed successfully." |