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"