import logging import importlib import pkgutil import inspect from typing import Dict, Any, Optional, List, Type, Callable, Awaitable, Union from datetime import datetime, UTC import json import base64 import uuid from modules.interfaces.serviceAppClass import User from modules.methods.methodBase import MethodBase, AuthSource, MethodResult from modules.workflow.serviceContainer import ServiceContainer from modules.interfaces.serviceChatModel import ( TaskStatus, UserInputRequest, ContentMetadata, ContentItem, ChatDocument, TaskDocument, ExtractedContent, TaskItem, TaskResult, ChatStat, ChatLog, ChatMessage, ChatWorkflow ) from modules.workflow.processorDocument import DocumentProcessor logger = logging.getLogger(__name__) class ChatManager: """Chat manager with improved AI integration and method handling""" def __init__(self, currentUser: User): self._discoverMethods() self.workflow: Optional[ChatWorkflow] = None self.currentTask: Optional[TaskItem] = None self.workflowHistory: List[ChatMessage] = [] self.documentProcessor = DocumentProcessor() self.userLanguage = None self.currentUser = currentUser # ===== Initialization and Setup ===== 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 setUserLanguage(self, languageCode: str): """Set the user's preferred language""" self.userLanguage = languageCode logger.debug(f"User language set to: {languageCode}") 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)}") # ===== Task Creation and Management ===== async def createInitialTask(self, workflow: ChatWorkflow, initialMessage: ChatMessage) -> Optional[TaskItem]: """Create the initial task from the first message""" try: # Get available methods and their actions methodCatalog = self.service.getAvailableMethods() # Process user input with AI processedInput = await self._processUserInput(initialMessage.message, methodCatalog) # Create actions from processed input actions = await self._createActions(processedInput['actions']) # Create task task = TaskItem( id=f"task_{datetime.now(UTC).timestamp()}", workflowId=workflow.id, userInput=processedInput['objective'], dataList=initialMessage.documents, actionList=actions, status=TaskStatus.PENDING, createdAt=datetime.now(UTC), updatedAt=datetime.now(UTC) ) # Add task to workflow workflow.tasks.append(task) return task except Exception as e: logger.error(f"Error creating initial task: {str(e)}") return None async def createNextTask(self, workflow: ChatWorkflow, previousResult: TaskResult) -> Optional[TaskItem]: """Create next task based on previous result""" try: # Check if previous result was successful if not previousResult.success: logger.error(f"Previous task failed: {previousResult.error}") return None # Extract task data from previous result taskData = previousResult.data if not taskData: logger.error("No task data in previous result") return None # Create next task nextTask = TaskItem( id=f"task_{datetime.now(UTC).timestamp()}", workflowId=workflow.id, userInput=taskData.get('objective', ''), actionList=await self._createActions(taskData.get('actions', [])), status=TaskStatus.PENDING, createdAt=datetime.now(UTC), updatedAt=datetime.now(UTC) ) # Add task to workflow workflow.tasks.append(nextTask) return nextTask except Exception as e: logger.error(f"Error creating next task: {str(e)}") return None async def identifyNextTask(self, workflow: ChatWorkflow) -> TaskResult: """Identify the next task based on workflow state""" try: # Get workflow summary summary = await self._summarizeWorkflow() # Generate prompt for next task prompt = f"""Based on the workflow summary: {summary} Determine what the next task should be. Return a JSON object with: - objective: The main goal or task to accomplish - actions: List of required actions with method and parameters """ # Get AI response response = await self._callAiBasic(prompt) # Parse response try: result = json.loads(response) return TaskResult( taskId=f"analysis_{datetime.now(UTC).timestamp()}", status=TaskStatus.COMPLETED, success=True, timestamp=datetime.now(UTC), data=result ) except json.JSONDecodeError as e: logger.error(f"Error parsing AI response: {str(e)}") return TaskResult( taskId=f"analysis_{datetime.now(UTC).timestamp()}", status=TaskStatus.FAILED, success=False, timestamp=datetime.now(UTC), error=f"Error parsing AI response: {str(e)}" ) except Exception as e: logger.error(f"Error identifying next task: {str(e)}") return TaskResult( taskId=f"analysis_{datetime.now(UTC).timestamp()}", status=TaskStatus.FAILED, success=False, timestamp=datetime.now(UTC), error=f"Error identifying next task: {str(e)}" ) async def callAi(self, messages: List[Dict[str, str]], produceUserAnswer: bool = False, temperature: float = None) -> str: """Enhanced AI service call with language support.""" if not self.service or not self.service.base: logger.error("AI service not set in ChatManager") return "Error: AI service not available" # Add language instruction for user-facing responses if produceUserAnswer and self.userLanguage: ltext = f"Please respond in '{self.userLanguage}' language." if messages and messages[0]["role"] == "system": if "language" not in messages[0]["content"].lower(): messages[0]["content"] = f"{ltext} {messages[0]['content']}" else: # Insert a system message with language instruction messages.insert(0, { "role": "system", "content": ltext }) # Call the AI service if temperature is not None: return await self.service.base.callAi(messages, temperature=temperature) else: return await self.service.base.callAi(messages) async def callAi4Image(self, imageData: Union[str, bytes], mimeType: str = None, prompt: str = "Describe this image") -> str: """Enhanced AI service call with language support.""" if not self.service or not self.service.base: logger.error("AI service not set in ChatManager") return "Error: AI service not available" return await self.service.base.analyzeImage(imageData, mimeType, prompt) 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.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' self.setUserLanguage(userLanguage) # 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.callAi(messages, produceUserAnswer=True, temperature=0.7) return feedback except Exception as e: logger.error(f"Error generating workflow feedback: {str(e)}") return "Workflow completed successfully." 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 "" # ===== Task Execution and Processing ===== async def executeTask(self, task: TaskItem) -> TaskResult: """Execute a task and return its result""" try: # Create result object result = TaskResult( taskId=task.id, status=task.status, success=True, timestamp=datetime.now(UTC) ) # Start timing startTime = datetime.now(UTC) # Execute each action for action in task.actionList: try: # Execute action actionResult = await action.execute() # Update action status action.status = actionResult.status if actionResult.error: action.error = actionResult.error except Exception as e: logger.error(f"Action execution error: {str(e)}") action.status = TaskStatus.FAILED action.error = str(e) # Calculate processing time endTime = datetime.now(UTC) result.processingTime = (endTime - startTime).total_seconds() # Update task status if all(action.status == TaskStatus.COMPLETED for action in task.actionList): result.status = TaskStatus.COMPLETED result.success = True else: result.status = TaskStatus.FAILED result.success = False result.error = "One or more actions failed" # Generate feedback and documents if task completed successfully if result.status == TaskStatus.COMPLETED: # Generate feedback using AI result.feedback = await self._processTaskResults(task) # Create output documents result.documents = await self._createOutputDocuments(task) else: result.feedback = f"Task failed: {result.error}" return result except Exception as e: logger.error(f"Task execution error: {str(e)}") raise async def parseTaskResult(self, workflow: ChatWorkflow, result: TaskResult) -> None: """Process and store task result in workflow""" try: # Find task in workflow task = next((t for t in workflow.tasks if t.id == result.taskId), None) if not task: logger.error(f"Task {result.taskId} not found in workflow") return # Update task status task.status = result.status if result.error: task.error = result.error # Create feedback message if available if result.feedback: message = ChatMessage( id=str(uuid.uuid4()), workflowId=workflow.id, role="assistant", message=result.feedback, status="step", documents=result.documents ) workflow.messages.append(message) # Update workflow stats if result.processingTime: if not workflow.stats: workflow.stats = ChatStat() workflow.stats.processingTime = (workflow.stats.processingTime or 0) + result.processingTime except Exception as e: logger.error(f"Error parsing task result: {str(e)}") raise async def shouldContinue(self, workflow: ChatWorkflow) -> bool: """Determine if workflow should continue""" try: # Check if workflow is in a terminal state if workflow.status in ["completed", "failed", "stopped"]: return False # Check if there are any pending tasks hasPendingTasks = any(t.status == TaskStatus.PENDING for t in workflow.tasks) if not hasPendingTasks: return False # Check if any task is currently running hasRunningTasks = any(t.status == TaskStatus.RUNNING for t in workflow.tasks) if hasRunningTasks: return True return False except Exception as e: logger.error(f"Error checking workflow continuation: {str(e)}") return False 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: TaskItem) -> Dict[str, Any]: """Analyze task results to determine next steps""" # Get workflow summary summary = await self._summarizeWorkflow() # Generate prompt for analysis prompt = f"""Based on the workflow summary and task results: {summary} Task: {task.userInput} Status: {task.status} Error: {task.error if task.error else 'None'} Determine if the workflow is complete and what the next steps should be. Return a JSON object with: - isComplete: boolean - objective: string - nextActions: array of action objects """ # Get AI response response = await self._callAiBasic(prompt) # Parse response return json.loads(response) def _promptInstructions(self, methodCatalog: Dict[str, Any], isInitialTask: bool = False) -> str: """Generate common prompt instructions for task analysis""" instructions = f"""Available Methods and Actions: {json.dumps(methodCatalog, indent=2)} """ if isInitialTask: instructions += """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" } } ] }""" else: instructions += """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"] } Note: Only use methods and actions that are available in the method catalog above.""" return instructions 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', '')} {self._promptInstructions(methodCatalog, isInitialTask=True)}""" # Call AI service response = await self._callAiBasic(prompt) return json.loads(response) async def _createActions(self, actionsData: List[Dict[str, Any]]) -> List[TaskItem]: """Create action objects from processed input""" actions = [] for actionData in actionsData: try: # Validate required fields if not all(k in actionData for k in ['method', 'action']): logger.warning(f"Skipping invalid action data: {actionData}") continue action = TaskItem( id=f"action_{datetime.now(UTC).timestamp()}", method=actionData['method'], action=actionData['action'], parameters=actionData.get('parameters', {}), status=TaskStatus.PENDING, retryCount=0, retryMax=actionData.get('retryMax', 3) ) actions.append(action) except Exception as e: logger.error(f"Error creating action: {str(e)}") continue return actions async def _processTaskResults(self, task: TaskItem) -> 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: TaskItem) -> 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.processFileIds(fileIds) return [] except Exception as e: logger.error(f"Error creating output documents: {str(e)}") return [] async def processFileIds(self, fileIds: List[str]) -> List[ChatDocument]: """Process multiple files and extract their contents.""" documents = [] for fileId in fileIds: # Get file metadata fileMetadata = self.service.functions.getFile(fileId) if not fileMetadata: logger.warning(f"File metadata not found for {fileId}") continue # Create ChatDocument document = ChatDocument( id=str(uuid.uuid4()), fileId=fileId, filename=fileMetadata.get("name", "Unknown"), fileSize=fileMetadata.get("size", 0), mimeType=fileMetadata.get("mimeType", "text/plain") ) documents.append(document) return documents async def addTaskResult(self, workflow: ChatWorkflow, result: TaskResult) -> None: """Add task result to workflow and update status""" try: # Find task in workflow task = next((t for t in workflow.tasks if t.id == result.taskId), None) if not task: logger.error(f"Task {result.taskId} not found in workflow") return # Update task status task.status = result.status if result.error: task.error = result.error # Create feedback message if available if result.feedback: message = ChatMessage( id=str(uuid.uuid4()), workflowId=workflow.id, role="assistant", message=result.feedback, status="step", documents=result.documents ) workflow.messages.append(message) # Update workflow stats if result.processingTime: if not workflow.stats: workflow.stats = ChatStat() workflow.stats.processingTime = (workflow.stats.processingTime or 0) + result.processingTime except Exception as e: logger.error(f"Error adding task result: {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. """