288 lines
15 KiB
Python
288 lines
15 KiB
Python
"""
|
|
AI processing method module.
|
|
Handles direct AI calls for any type of task.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Dict, Any, List, Optional
|
|
from datetime import datetime, UTC
|
|
|
|
from modules.workflows.methods.methodBase import MethodBase, action
|
|
from modules.datamodels.datamodelWorkflow import ActionResult
|
|
from modules.shared.timezoneUtils import get_utc_timestamp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class MethodAi(MethodBase):
|
|
"""AI processing methods."""
|
|
|
|
def __init__(self, service):
|
|
super().__init__(service)
|
|
self.name = "ai"
|
|
self.description = "AI processing methods"
|
|
# Centralized services interface (for AI)
|
|
from modules.services import getInterface as getServices
|
|
self.services = getServices(self.service.user, self.service.workflow)
|
|
|
|
def _format_timestamp_for_filename(self) -> str:
|
|
"""Format current timestamp as YYYYMMDD-hhmmss for filenames."""
|
|
return datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
|
|
|
|
@action
|
|
async def process(self, parameters: Dict[str, Any]) -> ActionResult:
|
|
"""
|
|
Perform an AI call for any type of task with optional document references
|
|
|
|
Parameters:
|
|
aiPrompt (str): The AI prompt for processing
|
|
documentList (list, optional): List of document references to include in context
|
|
expectedDocumentFormats (list, optional): Expected output formats with extension, mimeType, description
|
|
processingMode (str, optional): Processing mode ('basic', 'advanced', 'detailed') - defaults to 'basic'
|
|
includeMetadata (bool, optional): Whether to include metadata (default: True)
|
|
customInstructions (str, optional): Additional custom instructions for the AI
|
|
"""
|
|
try:
|
|
aiPrompt = parameters.get("aiPrompt")
|
|
documentList = parameters.get("documentList", [])
|
|
if isinstance(documentList, str):
|
|
documentList = [documentList]
|
|
expectedDocumentFormats = parameters.get("expectedDocumentFormats", [])
|
|
processingMode = parameters.get("processingMode", "basic")
|
|
includeMetadata = parameters.get("includeMetadata", True)
|
|
customInstructions = parameters.get("customInstructions", "")
|
|
|
|
if not aiPrompt:
|
|
return ActionResult.isFailure(
|
|
error="AI prompt is required"
|
|
)
|
|
|
|
# Determine output format first (needed for context building)
|
|
output_extension = ".txt" # Default
|
|
output_mime_type = "text/plain" # Default
|
|
|
|
if expectedDocumentFormats and len(expectedDocumentFormats) > 0:
|
|
expected_format = expectedDocumentFormats[0]
|
|
output_extension = expected_format.get("extension", ".txt")
|
|
output_mime_type = expected_format.get("mimeType", "text/plain")
|
|
logger.info(f"Using expected format: {output_extension} ({output_mime_type})")
|
|
|
|
# Build context from documents if provided
|
|
context = ""
|
|
if documentList:
|
|
chatDocuments = self.service.getChatDocumentsFromDocumentList(documentList)
|
|
if chatDocuments:
|
|
context_parts = []
|
|
for doc in chatDocuments:
|
|
file_info = self.service.getFileInfo(doc.fileId)
|
|
|
|
try:
|
|
# Use the document content extraction service with the specific AI prompt context
|
|
# This tells the extraction engine exactly what and how to extract
|
|
extraction_prompt = f"""
|
|
Extract content from this document for AI processing context.
|
|
|
|
AI Task: {aiPrompt}
|
|
Processing Mode: {processingMode}
|
|
Expected Output: {output_extension.upper()} format
|
|
|
|
Requirements:
|
|
1. Extract the most relevant text content that would be useful for the AI task
|
|
2. Focus on content that directly relates to: {aiPrompt}
|
|
3. Include key information, data, and insights that the AI needs
|
|
4. Provide clean, readable text without formatting artifacts
|
|
|
|
Document: {doc.fileName}
|
|
"""
|
|
|
|
logger.debug(f"Extracting content from {doc.fileName} with task-specific prompt: {extraction_prompt[:100]}...")
|
|
|
|
extracted_content = await self.service.extractContentFromDocument(
|
|
prompt=extraction_prompt.strip(),
|
|
document=doc
|
|
)
|
|
|
|
if extracted_content and extracted_content.contents:
|
|
# Get the first content item's data
|
|
content = ""
|
|
for content_item in extracted_content.contents:
|
|
if hasattr(content_item, 'data') and content_item.data:
|
|
content += content_item.data + " "
|
|
|
|
|
|
if content.strip():
|
|
metadata_info = ""
|
|
if file_info and includeMetadata:
|
|
metadata_info = f" (Size: {file_info.get('fileSize', 'unknown')}, Type: {file_info.get('mimeType', 'unknown')})"
|
|
|
|
# Adjust context length based on processing mode and AI task relevance
|
|
base_length = 5000 if processingMode == "detailed" else 3000 if processingMode == "advanced" else 2000
|
|
|
|
# For detailed mode, include more context
|
|
if processingMode == "detailed":
|
|
context_parts.append(f"Document: {doc.fileName}{metadata_info}\nRelevance to AI Task: This document contains content directly related to '{aiPrompt[:100]}...'\nContent:\n{content[:base_length]}...")
|
|
else:
|
|
context_parts.append(f"Document: {doc.fileName}{metadata_info}\nContent:\n{content[:base_length]}...")
|
|
else:
|
|
context_parts.append(f"Document: {doc.fileName} [No readable text content - binary file]")
|
|
else:
|
|
context_parts.append(f"Document: {doc.fileName} [No readable text content - binary file]")
|
|
|
|
except Exception as extract_error:
|
|
context_parts.append(f"Document: {doc.fileName} [Could not extract content - binary file]")
|
|
|
|
if context_parts:
|
|
# Add a summary header to help the AI understand the context
|
|
context_header = f"""
|
|
=== DOCUMENT CONTEXT FOR AI PROCESSING ===
|
|
AI Task: {aiPrompt[:100]}...
|
|
Processing Mode: {processingMode}
|
|
Expected Output Format: {output_extension.upper()}
|
|
Total Documents: {len(chatDocuments)}
|
|
|
|
The following documents contain content relevant to your task.
|
|
Use this information to provide the most accurate and helpful response.
|
|
================================================
|
|
"""
|
|
|
|
context = context_header + "\n\n" + "\n\n".join(context_parts)
|
|
logger.info(f"Included {len(chatDocuments)} documents in AI context with task-specific extraction")
|
|
|
|
# Build enhanced prompt
|
|
enhanced_prompt = aiPrompt
|
|
|
|
# Add processing mode instructions if specified (generic, not analysis-specific)
|
|
if processingMode == "detailed":
|
|
enhanced_prompt += "\n\nPlease provide a detailed response with comprehensive information."
|
|
elif processingMode == "advanced":
|
|
enhanced_prompt += "\n\nPlease provide an advanced response with deep insights."
|
|
|
|
# Add custom instructions if provided
|
|
if customInstructions:
|
|
enhanced_prompt += f"\n\nAdditional Instructions: {customInstructions}"
|
|
|
|
# Add format-specific instructions only if non-text format is requested
|
|
if output_extension != ".txt":
|
|
if output_extension == ".csv":
|
|
enhanced_prompt += f"\n\nCRITICAL: Deliver the result as pure CSV data without any markdown formatting, code blocks, or additional text. Output only the CSV content with proper headers and data rows."
|
|
elif output_extension == ".json":
|
|
enhanced_prompt += f"\n\nCRITICAL: Deliver the result as pure JSON data without any markdown formatting, code blocks, or additional text. Output only the JSON content."
|
|
elif output_extension == ".xml":
|
|
enhanced_prompt += f"\n\nCRITICAL: Deliver the result as pure XML data without any markdown formatting, code blocks, or additional text. Output only the XML content."
|
|
else:
|
|
enhanced_prompt += f"\n\nCRITICAL: Deliver the result as pure {output_extension.upper()} data without any markdown formatting, code blocks, or additional text."
|
|
|
|
# Call appropriate AI service based on processing mode
|
|
logger.info(f"Executing AI call with mode: {processingMode}, prompt length: {len(enhanced_prompt)}")
|
|
if context:
|
|
logger.info(f"Including context from {len(documentList)} documents")
|
|
|
|
# Encourage longer, structured outputs with a min-length hint
|
|
min_tokens_hint = "\n\nPlease ensure the response is substantial and complete."
|
|
call_prompt = enhanced_prompt + min_tokens_hint
|
|
|
|
# Centralized AI call with optional document context
|
|
documents = []
|
|
try:
|
|
if documentList:
|
|
for d in (chatDocuments or []):
|
|
try:
|
|
file_data = self.service.getFileData(d.fileId)
|
|
documents.append(
|
|
ChatDocument(
|
|
fileData=file_data,
|
|
fileName=d.fileName,
|
|
mimeType=d.mimeType
|
|
)
|
|
)
|
|
except Exception:
|
|
continue
|
|
except Exception:
|
|
documents = None
|
|
|
|
output_format = output_extension.replace('.', '') or 'txt'
|
|
result = await self.services.ai.callAi(
|
|
prompt=call_prompt,
|
|
documents=documents or None,
|
|
options={
|
|
"process_type": "text",
|
|
"operation_type": "generate_content",
|
|
"priority": "quality" if processingMode in ["advanced", "detailed"] else "speed",
|
|
"compress_prompt": processingMode != "detailed",
|
|
"compress_documents": True,
|
|
"process_documents_individually": True,
|
|
"processing_mode": processingMode,
|
|
"result_format_requested": output_format,
|
|
"include_metadata": includeMetadata,
|
|
"max_cost": 0.05 if processingMode in ["advanced", "detailed"] else 0.02,
|
|
"max_processing_time": 45 if processingMode in ["advanced", "detailed"] else 20
|
|
}
|
|
)
|
|
|
|
# If expected JSON and too short/not JSON, retry with stricter JSON guardrails
|
|
if output_extension == ".json":
|
|
import json
|
|
cleaned = (result or "").strip()
|
|
if cleaned.startswith('```json'):
|
|
cleaned = cleaned[7:]
|
|
if cleaned.endswith('```'):
|
|
cleaned = cleaned[:-3]
|
|
cleaned = cleaned.strip()
|
|
needs_retry = False
|
|
try:
|
|
parsed = json.loads(cleaned)
|
|
# Heuristic: small dict -> possibly underfilled
|
|
if isinstance(parsed, dict) and len(parsed.keys()) <= 2:
|
|
needs_retry = True
|
|
except Exception:
|
|
needs_retry = True
|
|
|
|
if needs_retry:
|
|
guardrail_prompt = (
|
|
enhanced_prompt
|
|
+ "\n\nCRITICAL: Return ONLY valid JSON, no markdown, no code fences. "
|
|
"Include all requested fields with detailed content."
|
|
)
|
|
try:
|
|
result = await self.services.ai.callAi(
|
|
prompt=guardrail_prompt,
|
|
documents=context or None,
|
|
options={
|
|
"process_type": "text",
|
|
"operation_type": "generate_content",
|
|
"priority": "quality",
|
|
"compress_prompt": False,
|
|
"compress_documents": True,
|
|
"process_documents_individually": True,
|
|
"processing_mode": "detailed",
|
|
"result_format_requested": "json",
|
|
"include_metadata": False,
|
|
"max_cost": 0.03,
|
|
"max_processing_time": 30
|
|
}
|
|
)
|
|
except Exception:
|
|
result = cleaned # fallback to first attempt
|
|
|
|
# Create result document
|
|
fileName = f"ai_{processingMode}_{self._format_timestamp_for_filename()}{output_extension}"
|
|
|
|
|
|
|
|
# Return result in the standard ActionResult format
|
|
return ActionResult.isSuccess(
|
|
documents=[{
|
|
"documentName": fileName,
|
|
"documentData": {
|
|
"result": result,
|
|
"fileName": fileName,
|
|
"processedDocuments": len(documentList) if documentList else 0
|
|
},
|
|
"mimeType": output_mime_type
|
|
}]
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in AI processing: {str(e)}")
|
|
return ActionResult.isFailure(
|
|
error=str(e)
|
|
)
|