"""Operator method implementation for handling collections and AI operations""" from typing import Dict, List, Any, Optional from datetime import datetime, UTC import logging from modules.workflow.methodBase import MethodBase, ActionResult, action logger = logging.getLogger(__name__) class OperatorService: """Service for operator operations like forEach and AI calls""" def __init__(self, serviceContainer: Any): self.serviceContainer = serviceContainer async def executeForEach(self, items: List[Any], action: Dict[str, Any]) -> List[Any]: """Execute an action for each item in a list""" try: results = [] for i, item in enumerate(items): logger.info(f"Executing forEach action {i+1}/{len(items)}") # Create context with current item context = { "item": item, "index": i, "total": len(items), "isFirst": i == 0, "isLast": i == len(items) - 1 } # Execute the action using the service container if "method" in action and "action" in action: methodName = action["method"] actionName = action["action"] parameters = action.get("parameters", {}) # Add context to parameters parameters["context"] = context parameters["currentItem"] = item # Execute the method action result = await self.serviceContainer.executeAction( methodName=methodName, actionName=actionName, parameters=parameters ) # Return the exact result data, not wrapped if result.success: results.append(result.data) else: results.append({"error": result.error}) else: # Simple action without method call results.append({"error": "No method specified"}) return results except Exception as e: logger.error(f"Error executing forEach: {str(e)}") return [{"error": str(e)}] * len(items) if items else [] async def executeAiCall(self, prompt: str, documents: List[Dict[str, Any]] = None) -> Dict[str, Any]: """Call AI service with document content""" try: # Prepare context from documents context = "" extractedDocuments = [] if documents: for i, doc in enumerate(documents): documentReference = doc.get('documentReference') contentExtractionPrompt = doc.get('contentExtractionPrompt', 'Extract the main content from this document') if documentReference: # Get documents from reference chatDocuments = self.serviceContainer.getChatDocumentsFromDocumentReference(documentReference) if chatDocuments: # Extract content from each document for j, chatDoc in enumerate(chatDocuments): try: # Extract content using the document manager extractedContent = await self.serviceContainer.documentManager.extractContentFromChatDocument( chatDocument=chatDoc, extractionPrompt=contentExtractionPrompt ) extractedDocuments.append({ "documentReference": documentReference, "documentId": chatDoc.id, "extractionPrompt": contentExtractionPrompt, "extractedContent": extractedContent }) # Add to context context += f"\n\nDocument {len(extractedDocuments)} (from {documentReference}):\n{extractedContent}" except Exception as e: logger.warning(f"Error extracting content from document {chatDoc.id}: {str(e)}") extractedDocuments.append({ "documentReference": documentReference, "documentId": chatDoc.id, "extractionPrompt": contentExtractionPrompt, "extractedContent": f"Error extracting content: {str(e)}" }) else: logger.warning(f"No documents found for reference: {documentReference}") extractedDocuments.append({ "documentReference": documentReference, "extractionPrompt": contentExtractionPrompt, "extractedContent": f"No documents found for reference: {documentReference}" }) # Create full prompt with context fullPrompt = f"{prompt}\n\nContext:\n{context}" if context else prompt # Call AI service aiResponse = await self.serviceContainer.interfaceAiCalls.callAiTextAdvanced(fullPrompt) return { "prompt": prompt, "documentsProcessed": len(extractedDocuments), "extractedDocuments": extractedDocuments, "response": aiResponse, "timestamp": datetime.now(UTC).isoformat() } except Exception as e: logger.error(f"Error executing AI call: {str(e)}") return { "error": str(e), "prompt": prompt, "documentsProcessed": 0, "extractedDocuments": [], "response": None } class MethodOperator(MethodBase): """Operator method implementation for handling collections and AI operations""" def __init__(self, serviceContainer: Any): super().__init__(serviceContainer) self.name = "operator" self.description = "Handle operations like forEach and AI calls" self.operatorService = OperatorService(serviceContainer) @action async def forEach(self, parameters: Dict[str, Any]) -> ActionResult: """ Execute an action for each item in a list Parameters: items (List[Any]): List of items to process action (Dict[str, Any]): Action to execute for each item (contains method, action, parameters) """ try: items = parameters.get("items", []) action = parameters.get("action", {}) if not items or not action: return self._createResult( success=False, data={}, error="Items and action are required" ) # Execute forEach operation results = await self.operatorService.executeForEach( items=items, action=action ) return self._createResult( success=True, data=results ) except Exception as e: logger.error(f"Error in forEach execution: {str(e)}") return self._createResult( success=False, data={}, error=str(e) ) @action async def aiCall(self, parameters: Dict[str, Any]) -> ActionResult: """ Call AI service with document content Parameters: prompt (str): The prompt to send to the AI service documents (List[Dict[str, Any]], optional): List of documents to include in context Each document should have: documentReference (str), contentExtractionPrompt (str, optional) """ try: prompt = parameters.get("prompt") documents = parameters.get("documents", []) # List of {documentReference, contentExtractionPrompt} if not prompt: return self._createResult( success=False, data={}, error="Prompt is required" ) # Execute AI call result = await self.operatorService.executeAiCall( prompt=prompt, documents=documents ) return self._createResult( success=True, data=result ) except Exception as e: logger.error(f"Error in AI call execution: {str(e)}") return self._createResult( success=False, data={}, error=str(e) )