228 lines
9.5 KiB
Python
228 lines
9.5 KiB
Python
"""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)
|
|
)
|