199 lines
11 KiB
Python
199 lines
11 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.chat.methodBase import MethodBase, action
|
|
from modules.interfaces.interfaceChatModel import ActionResult
|
|
from modules.shared.timezoneUtils import get_utc_timestamp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class MethodAi(MethodBase):
|
|
"""AI method implementation for direct AI processing"""
|
|
|
|
def __init__(self, serviceCenter: Any):
|
|
"""Initialize the AI method"""
|
|
super().__init__(serviceCenter)
|
|
self.name = "ai"
|
|
self.description = "Handle direct AI processing for any type of task"
|
|
|
|
@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", [])
|
|
expectedDocumentFormats = parameters.get("expectedDocumentFormats", [])
|
|
processingMode = parameters.get("processingMode", "basic")
|
|
includeMetadata = parameters.get("includeMetadata", True)
|
|
customInstructions = parameters.get("customInstructions", "")
|
|
|
|
if not aiPrompt:
|
|
return ActionResult.failure(
|
|
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")
|
|
|
|
if processingMode in ["advanced", "detailed"]:
|
|
result = await self.service.callAiTextAdvanced(enhanced_prompt, context)
|
|
else:
|
|
result = await self.service.callAiTextBasic(enhanced_prompt, context)
|
|
|
|
# Create result document
|
|
timestamp = int(get_utc_timestamp())
|
|
fileName = f"ai_{processingMode}_{timestamp}{output_extension}"
|
|
|
|
|
|
|
|
# Return result in the standard ActionResult format
|
|
return ActionResult.success(
|
|
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.failure(
|
|
error=str(e)
|
|
)
|