178 lines
6.9 KiB
Python
178 lines
6.9 KiB
Python
from typing import Dict, List, Any, Optional
|
|
from datetime import datetime, UTC
|
|
import logging
|
|
from .methodBase import MethodBase
|
|
from modules.interfaces.interfaceChatModel import MethodResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class MethodOperator(MethodBase):
|
|
"""Operator methods for handling collections and AI operations"""
|
|
|
|
def __init__(self, serviceContainer: Any):
|
|
super().__init__(serviceContainer)
|
|
self.name = "operator"
|
|
self.description = "Operator methods for handling collections and AI operations"
|
|
|
|
@property
|
|
def actions(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Available actions and their parameters"""
|
|
return {
|
|
"forEach": {
|
|
"description": "Execute an action for each item in a list",
|
|
"parameters": {
|
|
"items": {
|
|
"type": "List[Any]",
|
|
"description": "List of items to process",
|
|
"required": True
|
|
},
|
|
"action": {
|
|
"type": "Dict[str, Any]",
|
|
"description": "Action to execute for each item",
|
|
"required": True,
|
|
"properties": {
|
|
"method": {"type": "str", "required": True},
|
|
"action": {"type": "str", "required": True},
|
|
"parameters": {"type": "Dict[str, Any]", "required": False}
|
|
}
|
|
}
|
|
}
|
|
},
|
|
"aiCall": {
|
|
"description": "Call AI service with document content",
|
|
"parameters": {
|
|
"prompt": {
|
|
"type": "str",
|
|
"description": "Prompt for AI processing",
|
|
"required": True
|
|
},
|
|
"extractedDocumentContent": {
|
|
"type": "List[Dict[str, str]]",
|
|
"description": "List of documents and their extraction prompts",
|
|
"required": True,
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"document": {"type": "str", "required": True},
|
|
"promptForContentExtraction": {"type": "str", "required": True}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async def _executeAction(self, action: str, parameters: Dict[str, Any], authData: Optional[Dict[str, Any]] = None) -> MethodResult:
|
|
"""Execute operator action"""
|
|
if action == "forEach":
|
|
return await self._executeForEach(parameters)
|
|
elif action == "aiCall":
|
|
return await self._executeAiCall(parameters)
|
|
else:
|
|
raise ValueError(f"Unsupported action: {action}")
|
|
|
|
async def _executeForEach(self, parameters: Dict[str, Any]) -> MethodResult:
|
|
"""Execute forEach operation"""
|
|
try:
|
|
items = parameters.get("items", [])
|
|
action = parameters.get("action", {})
|
|
|
|
if not items or not action:
|
|
return self._createResult(
|
|
success=False,
|
|
data={},
|
|
error="Missing required parameters"
|
|
)
|
|
|
|
results = []
|
|
for item in items:
|
|
try:
|
|
# Execute action for each item
|
|
method = action.get("method")
|
|
action_name = action.get("action")
|
|
action_params = action.get("parameters", {})
|
|
|
|
# Add current item to parameters
|
|
action_params["item"] = item
|
|
|
|
# Execute method action
|
|
method_result = await self.service.methods[method][action_name](action_params)
|
|
results.append(method_result)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing item: {str(e)}")
|
|
results.append({
|
|
"success": False,
|
|
"error": str(e)
|
|
})
|
|
|
|
return self._createResult(
|
|
success=True,
|
|
data={"results": results}
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in forEach execution: {str(e)}")
|
|
return self._createResult(
|
|
success=False,
|
|
data={},
|
|
error=str(e)
|
|
)
|
|
|
|
async def _executeAiCall(self, parameters: Dict[str, Any]) -> MethodResult:
|
|
"""Execute AI call with document content"""
|
|
try:
|
|
prompt = parameters.get("prompt")
|
|
documents = parameters.get("extractedDocumentContent", [])
|
|
|
|
if not prompt:
|
|
return self._createResult(
|
|
success=False,
|
|
data={},
|
|
error="Missing prompt parameter"
|
|
)
|
|
|
|
# Extract content from documents
|
|
extracted_content = []
|
|
for doc in documents:
|
|
try:
|
|
doc_ref = doc.get("document")
|
|
doc_prompt = doc.get("promptForContentExtraction")
|
|
|
|
if not doc_ref or not doc_prompt:
|
|
continue
|
|
|
|
# Extract content using document manager
|
|
content = self.service.extractContent(doc_prompt, doc_ref)
|
|
extracted_content.append({
|
|
"document": doc_ref,
|
|
"content": content
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error extracting document content: {str(e)}")
|
|
continue
|
|
|
|
# Prepare AI prompt with extracted content
|
|
full_prompt = f"{prompt}\n\nExtracted Content:\n"
|
|
for content in extracted_content:
|
|
full_prompt += f"\nDocument: {content['document']}\n{content['content']}\n"
|
|
|
|
# Call AI service
|
|
response = await self.service.callAiTextBasic(full_prompt)
|
|
|
|
return self._createResult(
|
|
success=True,
|
|
data={
|
|
"response": response,
|
|
"processedDocuments": len(extracted_content)
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in AI call execution: {str(e)}")
|
|
return self._createResult(
|
|
success=False,
|
|
data={},
|
|
error=str(e)
|
|
)
|