""" 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.methods.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 processing methods.""" def __init__(self, service): super().__init__(service) self.name = "ai" self.description = "AI processing methods" 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 if processingMode in ["advanced", "detailed"]: result = await self.service.callAiTextAdvanced(call_prompt, context) else: result = await self.service.callAiTextBasic(call_prompt, context) # 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.service.callAiTextAdvanced(guardrail_prompt, context) 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) )