395 lines
No EOL
14 KiB
Python
395 lines
No EOL
14 KiB
Python
import logging
|
|
import importlib
|
|
import pkgutil
|
|
import inspect
|
|
from typing import Dict, Any, Optional, List, Type
|
|
from datetime import datetime, UTC
|
|
import json
|
|
import asyncio
|
|
|
|
from modules.methods.methodBase import MethodBase, AuthSource, MethodResult
|
|
from modules.workflow.serviceContainer import ServiceContainer
|
|
from modules.interfaces.serviceChatModel import AgentTask, AgentAction, AgentResult, Action, TaskStatus
|
|
from modules.workflow.managerPrompt import AIPromptManager
|
|
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._discover_methods()
|
|
self.workflow = None
|
|
self.current_task = None
|
|
self.workflow_history = []
|
|
|
|
def _discover_methods(self):
|
|
"""Dynamically discover all method classes in modules.methods package"""
|
|
try:
|
|
# Import the methods package
|
|
methods_package = importlib.import_module('modules.methods')
|
|
|
|
# Discover all modules in the package
|
|
for _, name, is_pkg in pkgutil.iter_modules(methods_package.__path__):
|
|
if not is_pkg 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 item_name, item in inspect.getmembers(module):
|
|
if (inspect.isclass(item) and
|
|
issubclass(item, MethodBase) and
|
|
item != MethodBase):
|
|
# Instantiate the method and add to service
|
|
method_instance = item()
|
|
self.service.methods[method_instance.name] = method_instance
|
|
logger.info(f"Discovered method: {method_instance.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: Any, context: Dict[str, Any]) -> None:
|
|
"""Initialize chat manager with workflow and context"""
|
|
self.service.workflow = workflow
|
|
self.service.context = context
|
|
|
|
# Initialize AI model
|
|
self.service.model = {
|
|
'callAiBasic': self._call_ai_basic,
|
|
'callAiAdvanced': self._call_ai_advanced
|
|
}
|
|
|
|
# Initialize document processor
|
|
self.service.document_processor.initialize(context)
|
|
|
|
async def create_initial_task(self, user_input: Dict[str, Any]) -> AgentTask:
|
|
"""Create initial task from user input"""
|
|
# Get available methods and their actions
|
|
method_catalog = self.service.get_available_methods()
|
|
|
|
# Process user input with AI
|
|
processed_input = await self._process_user_input(user_input, method_catalog)
|
|
|
|
# Create actions from processed input
|
|
actions = await self._create_actions(processed_input['actions'])
|
|
|
|
# Create task
|
|
task = AgentTask(
|
|
id=f"task_{datetime.now(UTC).timestamp()}",
|
|
workflowId=self.workflow.id,
|
|
userInput=processed_input['objective'],
|
|
dataList=user_input.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 execute_current_task(self) -> None:
|
|
"""Execute current task"""
|
|
task = self.service.tasks.get('current')
|
|
if not task:
|
|
raise ValueError("No current task to execute")
|
|
|
|
await self.service.execute_task(task)
|
|
|
|
async def define_next_task(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._analyze_task_results(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._create_next_task(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 _process_user_input(self, user_input: Dict[str, Any], method_catalog: 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: {user_input.get('message', '')}
|
|
|
|
Available Methods and Actions:
|
|
{json.dumps(method_catalog, 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.service.model['callAiBasic'](prompt)
|
|
return json.loads(response)
|
|
|
|
async def _create_actions(self, actions_data: List[Dict[str, Any]]) -> List[AgentAction]:
|
|
"""Create action objects from processed input"""
|
|
actions = []
|
|
for action_data in actions_data:
|
|
method = self.service.get_method(action_data['method'])
|
|
if not method:
|
|
continue
|
|
|
|
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)
|
|
|
|
return actions
|
|
|
|
async def _summarize_workflow(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.service.model['callAiBasic'](prompt)
|
|
|
|
async def _analyze_task_results(self, task: AgentTask) -> Dict[str, Any]:
|
|
"""Analyze task results to determine next steps"""
|
|
# Get workflow summary
|
|
summary = await self._summarize_workflow()
|
|
|
|
# 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.service.model['callAiBasic'](prompt)
|
|
return json.loads(response)
|
|
|
|
def _create_next_task(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 process_task(self, task: Any) -> Dict[str, Any]:
|
|
"""Process a task with improved error handling and AI integration"""
|
|
try:
|
|
# Execute task
|
|
await self.service.execute_task(task)
|
|
|
|
# Process results
|
|
if task.status == 'success':
|
|
# Generate feedback using AI
|
|
feedback = await self._process_task_results(task)
|
|
task.thisTaskFeedback = feedback
|
|
|
|
# Create output documents
|
|
documents = await self._create_output_documents(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)}"
|
|
}
|
|
|
|
async def _process_task_results(self, task: Any) -> str:
|
|
"""Process task results and generate feedback using AI"""
|
|
try:
|
|
# Create context for AI
|
|
context = {
|
|
"task": "Process task results",
|
|
"document": {"name": "Task Results", "type": "json"}
|
|
}
|
|
|
|
# Generate prompt
|
|
prompt = self.service.prompt_manager.generate_prompt(
|
|
context,
|
|
[
|
|
{"input": "Task results", "output": "Generate summary"}
|
|
]
|
|
)
|
|
|
|
# Call AI
|
|
response = await self.service.model['callAiBasic'](
|
|
f"""Process task results and generate feedback:
|
|
Task Input: {task.userInput}
|
|
Method Results: {task.result}
|
|
Generated Documents: {task.documentsOutput}
|
|
|
|
{prompt}
|
|
|
|
Please provide:
|
|
1. Summary of completed actions
|
|
2. Generated document descriptions
|
|
3. Next steps or completion status
|
|
|
|
Format your response as JSON:
|
|
{{
|
|
"summary": "string",
|
|
"documents": ["string"],
|
|
"nextSteps": ["string"]
|
|
}}
|
|
"""
|
|
)
|
|
|
|
# Parse and validate response
|
|
try:
|
|
result = json.loads(response)
|
|
return result.get("summary", "Task completed successfully")
|
|
except json.JSONDecodeError:
|
|
return response.strip()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing task results: {str(e)}")
|
|
return f"Error processing results: {str(e)}"
|
|
|
|
async def _create_output_documents(self, task: Any) -> List[Dict[str, Any]]:
|
|
"""Create output documents from task results"""
|
|
try:
|
|
documents = []
|
|
|
|
# Process each document
|
|
for doc in task.documentsOutput:
|
|
processed = self.service.document_processor.process_with_context(
|
|
doc,
|
|
{
|
|
"id": doc.get("id", ""),
|
|
"extractionHistory": doc.get("extractionHistory", []),
|
|
"relevantSections": doc.get("relevantSections", []),
|
|
"processingStatus": doc.get("processingStatus", {})
|
|
}
|
|
)
|
|
|
|
if processed:
|
|
documents.append(processed)
|
|
|
|
return documents
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating output documents: {str(e)}")
|
|
return []
|
|
|
|
async def _call_ai_basic(self, prompt: str) -> str:
|
|
"""Call basic AI model"""
|
|
# TODO: Implement actual AI call
|
|
return "AI response placeholder"
|
|
|
|
async def _call_ai_advanced(self, prompt: str, context: Dict[str, Any]) -> str:
|
|
"""Call advanced AI model with context"""
|
|
# TODO: Implement actual AI call
|
|
return "AI response placeholder" |