import json import logging import time from typing import Dict, Any, List, Optional, Tuple, Union from modules.datamodels.datamodelChat import PromptPlaceholder, ChatDocument from modules.services.serviceExtraction.mainServiceExtraction import ExtractionService from modules.datamodels.datamodelAi import AiCallRequest, AiCallOptions, OperationTypeEnum, PriorityEnum, ProcessingModeEnum from modules.interfaces.interfaceAiObjects import AiObjects from modules.shared.jsonUtils import ( extractJsonString, repairBrokenJson, extractSectionsFromDocument, buildContinuationContext ) logger = logging.getLogger(__name__) # Rebuild the model to resolve forward references AiCallRequest.model_rebuild() class AiService: """AI service with core operations integrated.""" def __init__(self, serviceCenter=None) -> None: """Initialize AI service with service center access. Args: serviceCenter: Service center instance for accessing other services """ self.services = serviceCenter # Only depend on interfaces self.aiObjects = None # Will be initialized in create() or _ensureAiObjectsInitialized() # Submodules initialized as None - will be set in _initializeSubmodules() after aiObjects is ready self.extractionService = None def _initializeSubmodules(self): """Initialize all submodules after aiObjects is ready.""" if self.aiObjects is None: raise RuntimeError("aiObjects must be initialized before initializing submodules") if self.extractionService is None: logger.info("Initializing ExtractionService...") self.extractionService = ExtractionService(self.services) async def _ensureAiObjectsInitialized(self): """Ensure aiObjects is initialized and submodules are ready.""" if self.aiObjects is None: logger.info("Lazy initializing AiObjects...") self.aiObjects = await AiObjects.create() logger.info("AiObjects initialization completed") # Initialize submodules after aiObjects is ready self._initializeSubmodules() @classmethod async def create(cls, serviceCenter=None) -> "AiService": """Create AiService instance with all connectors and submodules initialized.""" logger.info("AiService.create() called") instance = cls(serviceCenter) logger.info("AiService created, about to call AiObjects.create()...") instance.aiObjects = await AiObjects.create() logger.info("AiObjects.create() completed") # Initialize all submodules after aiObjects is ready instance._initializeSubmodules() logger.info("AiService submodules initialized") return instance # Helper methods def _buildPromptWithPlaceholders(self, prompt: str, placeholders: Optional[Dict[str, str]]) -> str: """ Build full prompt by replacing placeholders with their content. Uses the new {{KEY:placeholder}} format. Args: prompt: The base prompt template placeholders: Dictionary of placeholder key-value pairs Returns: Prompt with placeholders replaced """ if not placeholders: return prompt full_prompt = prompt for placeholder, content in placeholders.items(): # Skip if content is None or empty if content is None: continue # Replace {{KEY:placeholder}} full_prompt = full_prompt.replace(f"{{{{KEY:{placeholder}}}}}", str(content)) return full_prompt async def _analyzePromptAndCreateOptions(self, prompt: str) -> AiCallOptions: """Analyze prompt to determine appropriate AiCallOptions parameters.""" try: # Get dynamic enum values from Pydantic models operationTypes = [e.value for e in OperationTypeEnum] priorities = [e.value for e in PriorityEnum] processingModes = [e.value for e in ProcessingModeEnum] # Create analysis prompt for AI to determine operation type and parameters analysisPrompt = f""" You are an AI operation analyzer. Analyze the following prompt and determine the most appropriate operation type and parameters. PROMPT TO ANALYZE: {self.services.utils.sanitizePromptContent(prompt, 'userinput')} Based on the prompt content, determine: 1. operationType: Choose the most appropriate from: {', '.join(operationTypes)} 2. priority: Choose from: {', '.join(priorities)} 3. processingMode: Choose from: {', '.join(processingModes)} 4. compressPrompt: true/false (true for story-like prompts, false for structured prompts with JSON/schemas) 5. compressContext: true/false (true to summarize context, false to process fully) Respond with ONLY a JSON object in this exact format: {{ "operationType": "dataAnalyse", "priority": "balanced", "processingMode": "basic", "compressPrompt": true, "compressContext": true }} """ # Use AI to analyze the prompt request = AiCallRequest( prompt=analysisPrompt, options=AiCallOptions( operationType=OperationTypeEnum.DATA_ANALYSE, priority=PriorityEnum.SPEED, processingMode=ProcessingModeEnum.BASIC, compressPrompt=True, compressContext=False ) ) response = await self.aiObjects.call(request) # Parse AI response try: jsonStart = response.content.find('{') jsonEnd = response.content.rfind('}') + 1 if jsonStart != -1 and jsonEnd > jsonStart: analysis = json.loads(response.content[jsonStart:jsonEnd]) # Map string values to enums operationType = OperationTypeEnum(analysis.get('operationType', 'dataAnalyse')) priority = PriorityEnum(analysis.get('priority', 'balanced')) processingMode = ProcessingModeEnum(analysis.get('processingMode', 'basic')) return AiCallOptions( operationType=operationType, priority=priority, processingMode=processingMode, compressPrompt=analysis.get('compressPrompt', True), compressContext=analysis.get('compressContext', True) ) except Exception as e: logger.warning(f"Failed to parse AI analysis response: {e}") except Exception as e: logger.warning(f"Prompt analysis failed: {e}") # Fallback to default options return AiCallOptions( operationType=OperationTypeEnum.DATA_ANALYSE, priority=PriorityEnum.BALANCED, processingMode=ProcessingModeEnum.BASIC ) async def _callAiWithLooping( self, prompt: str, options: AiCallOptions, debugPrefix: str = "ai_call", promptBuilder: Optional[callable] = None, promptArgs: Optional[Dict[str, Any]] = None, operationId: Optional[str] = None ) -> str: """ Shared core function for AI calls with repair-based looping system. Automatically repairs broken JSON and continues generation seamlessly. Args: prompt: The prompt to send to AI options: AI call configuration options debugPrefix: Prefix for debug file names promptBuilder: Optional function to rebuild prompts for continuation promptArgs: Optional arguments for prompt builder operationId: Optional operation ID for progress tracking Returns: Complete AI response after all iterations """ maxIterations = 50 # Prevent infinite loops iteration = 0 allSections = [] # Accumulate all sections across iterations lastRawResponse = None # Store last raw JSON response for continuation documentMetadata = None # Store document metadata (title, filename) from first iteration while iteration < maxIterations: iteration += 1 # Update progress for iteration start if operationId: if iteration == 1: self.services.workflow.progressLogUpdate(operationId, 0.5, f"Starting AI call iteration {iteration}") else: # For continuation iterations, show progress incrementally baseProgress = 0.5 + (min(iteration - 1, maxIterations) / maxIterations * 0.4) # Progress from 0.5 to 0.9 over maxIterations iterations self.services.workflow.progressLogUpdate(operationId, baseProgress, f"Continuing generation (iteration {iteration})") # Build iteration prompt if len(allSections) > 0 and promptBuilder and promptArgs: # This is a continuation - build continuation context with raw JSON and rebuild prompt continuationContext = buildContinuationContext(allSections, lastRawResponse) if not lastRawResponse: logger.warning(f"Iteration {iteration}: No previous response available for continuation!") # Rebuild prompt with continuation context using the provided prompt builder iterationPrompt = await promptBuilder(**promptArgs, continuationContext=continuationContext) else: # First iteration - use original prompt iterationPrompt = prompt # Make AI call try: if operationId and iteration == 1: self.services.workflow.progressLogUpdate(operationId, 0.51, "Calling AI model") request = AiCallRequest( prompt=iterationPrompt, context="", options=options ) # Write the ACTUAL prompt sent to AI if iteration == 1: self.services.utils.writeDebugFile(iterationPrompt, f"{debugPrefix}_prompt") else: self.services.utils.writeDebugFile(iterationPrompt, f"{debugPrefix}_prompt_iteration_{iteration}") response = await self.aiObjects.call(request) result = response.content # Update progress after AI call if operationId: if iteration == 1: self.services.workflow.progressLogUpdate(operationId, 0.6, f"AI response received (iteration {iteration})") else: progress = 0.6 + (min(iteration - 1, 10) * 0.03) self.services.workflow.progressLogUpdate(operationId, progress, f"Processing response (iteration {iteration})") # Write raw AI response to debug file if iteration == 1: self.services.utils.writeDebugFile(result, f"{debugPrefix}_response") else: self.services.utils.writeDebugFile(result, f"{debugPrefix}_response_iteration_{iteration}") # Emit stats for this iteration self.services.workflow.storeWorkflowStat( self.services.currentWorkflow, response, f"ai.call.{debugPrefix}.iteration_{iteration}" ) if not result or not result.strip(): logger.warning(f"Iteration {iteration}: Empty response, stopping") break # Store raw response for continuation (even if broken) lastRawResponse = result # Check for complete_response flag in raw response (before parsing) import re if re.search(r'"complete_response"\s*:\s*true', result, re.IGNORECASE): pass # Flag detected, will stop in _shouldContinueGeneration # Extract sections from response (handles both valid and broken JSON) extractedSections, wasJsonComplete, parsedResult = self._extractSectionsFromResponse(result, iteration, debugPrefix) # Extract document metadata from first iteration if available if iteration == 1 and parsedResult and not documentMetadata: documentMetadata = self._extractDocumentMetadata(parsedResult) # Update progress after parsing if operationId: if extractedSections: self.services.workflow.progressLogUpdate(operationId, 0.65 + (min(iteration - 1, 10) * 0.025), f"Extracted {len(extractedSections)} sections (iteration {iteration})") if not extractedSections: # If we're in continuation mode and JSON was incomplete, don't stop - continue to allow retry if iteration > 1 and not wasJsonComplete: logger.warning(f"Iteration {iteration}: No sections extracted from continuation fragment, continuing for another attempt") continue # Otherwise, stop if no sections logger.warning(f"Iteration {iteration}: No sections extracted, stopping") break # Add new sections to accumulator allSections.extend(extractedSections) # Check if we should continue (completion detection) if self._shouldContinueGeneration(allSections, iteration, wasJsonComplete, result): continue else: # Done - build final result if operationId: self.services.workflow.progressLogUpdate(operationId, 0.95, f"Generation complete ({iteration} iterations, {len(allSections)} sections)") logger.info(f"Generation complete after {iteration} iterations: {len(allSections)} sections") break except Exception as e: logger.error(f"Error in AI call iteration {iteration}: {str(e)}") break if iteration >= maxIterations: logger.warning(f"AI call stopped after maximum iterations ({maxIterations})") # Build final result from accumulated sections final_result = self._buildFinalResultFromSections(allSections, documentMetadata) # Write final result to debug file self.services.utils.writeDebugFile(final_result, f"{debugPrefix}_final_result") return final_result def _extractSectionsFromResponse( self, result: str, iteration: int, debugPrefix: str ) -> Tuple[List[Dict[str, Any]], bool, Optional[Dict[str, Any]]]: """ Extract sections from AI response, handling both valid and broken JSON. Uses repair mechanism for broken JSON. Checks for "complete_response": true flag to determine completion. Returns (sections, wasJsonComplete, parsedResult) """ # First, try to parse as valid JSON try: extracted = extractJsonString(result) parsed_result = json.loads(extracted) # Check if AI marked response as complete isComplete = parsed_result.get("complete_response", False) == True # Extract sections from parsed JSON sections = extractSectionsFromDocument(parsed_result) # If AI marked as complete, always return as complete if isComplete: return sections, True, parsed_result # If in continuation mode (iteration > 1), continuation responses are expected to be fragments # A fragment with 0 extractable sections means JSON is incomplete - need another iteration if len(sections) == 0 and iteration > 1: return sections, False, parsed_result # Mark as incomplete so loop continues # First iteration with 0 sections means empty response - stop if len(sections) == 0: return sections, True, parsed_result # Complete but empty return sections, True, parsed_result # JSON was complete with sections except json.JSONDecodeError as e: # Broken JSON - try repair mechanism (normal in iterative generation) self.services.utils.writeDebugFile(result, f"{debugPrefix}_broken_json_iteration_{iteration}") # Try to repair repaired_json = repairBrokenJson(result) if repaired_json: # Extract sections from repaired JSON sections = extractSectionsFromDocument(repaired_json) return sections, False, repaired_json # JSON was broken but repaired else: # Repair failed - log error logger.error(f"Iteration {iteration}: All repair strategies failed") return [], False, None except Exception as e: logger.error(f"Iteration {iteration}: Unexpected error during parsing: {str(e)}") return [], False, None def _shouldContinueGeneration( self, allSections: List[Dict[str, Any]], iteration: int, wasJsonComplete: bool, rawResponse: str = None ) -> bool: """ Determine if generation should continue based on JSON completeness, complete_response flag, and task completion. Returns True if we should continue, False if done. """ if len(allSections) == 0: return True # No sections yet, continue # Check for complete_response flag in raw response if rawResponse: import re if re.search(r'"complete_response"\s*:\s*true', rawResponse, re.IGNORECASE): logger.info(f"Iteration {iteration}: AI marked response as complete (complete_response flag detected)") return False # If JSON was complete, stop (AI should have set complete_response if task is done) # For continuation iterations (iteration > 1), if JSON is complete but no flag was set, # stop to prevent infinite loops - AI had a chance to set the flag if wasJsonComplete: if iteration > 1: # Continuation mode: JSON complete without flag means we're likely done # Stop to prevent infinite loops logger.info(f"Iteration {iteration}: JSON complete without complete_response flag - stopping") return False # First iteration with complete JSON - done return False else: # JSON was incomplete/broken - continue return True def _extractDocumentMetadata( self, parsedResult: Dict[str, Any] ) -> Optional[Dict[str, Any]]: """ Extract document metadata (title, filename) from parsed AI response. Returns dict with 'title' and 'filename' keys if found, None otherwise. """ if not isinstance(parsedResult, dict): return None # Try to get from documents array (preferred structure) if "documents" in parsedResult and isinstance(parsedResult["documents"], list) and len(parsedResult["documents"]) > 0: firstDoc = parsedResult["documents"][0] if isinstance(firstDoc, dict): title = firstDoc.get("title") filename = firstDoc.get("filename") if title or filename: return { "title": title, "filename": filename } return None def _buildFinalResultFromSections( self, allSections: List[Dict[str, Any]], documentMetadata: Optional[Dict[str, Any]] = None ) -> str: """ Build final JSON result from accumulated sections. Uses AI-provided metadata (title, filename) if available. """ if not allSections: return "" # Extract metadata from AI response if available title = "Generated Document" filename = "document.json" if documentMetadata: if documentMetadata.get("title"): title = documentMetadata["title"] if documentMetadata.get("filename"): filename = documentMetadata["filename"] # Build documents structure # Assuming single document for now documents = [{ "id": "doc_1", "title": title, "filename": filename, "sections": allSections }] result = { "metadata": { "split_strategy": "single_document", "source_documents": [], "extraction_method": "ai_generation" }, "documents": documents } return json.dumps(result, indent=2) # Public API Methods # Planning AI Call async def callAiPlanning( self, prompt: str, placeholders: Optional[List[PromptPlaceholder]] = None ) -> str: """ Planning AI call for task planning, action planning, action selection, etc. Always uses static parameters optimized for planning tasks. Args: prompt: The planning prompt placeholders: Optional list of placeholder replacements Returns: Planning JSON response """ await self._ensureAiObjectsInitialized() # Planning calls always use static parameters options = AiCallOptions( operationType=OperationTypeEnum.PLAN, priority=PriorityEnum.QUALITY, processingMode=ProcessingModeEnum.DETAILED, compressPrompt=False, compressContext=False ) # Build full prompt with placeholders if placeholders: placeholdersDict = {p.label: p.content for p in placeholders} fullPrompt = self._buildPromptWithPlaceholders(prompt, placeholdersDict) else: fullPrompt = prompt # Root-cause fix: planning must return raw single-shot JSON, not section-based output request = AiCallRequest( prompt=fullPrompt, context="", options=options ) # Debug: persist prompt/response for analysis self.services.utils.writeDebugFile(fullPrompt, "plan_prompt") response = await self.aiObjects.call(request) result = response.content or "" self.services.utils.writeDebugFile(result, "plan_response") return result # Document Generation AI Call async def callAiDocuments( self, prompt: str, documents: Optional[List[ChatDocument]] = None, options: Optional[AiCallOptions] = None, outputFormat: Optional[str] = None, title: Optional[str] = None ) -> Union[str, Dict[str, Any]]: """ Document generation AI call for all non-planning calls. Uses the current unified path with extraction and generation. Args: prompt: The main prompt for the AI call documents: Optional list of documents to process options: AI call configuration options outputFormat: Optional output format for document generation title: Optional title for generated documents Returns: AI response as string, or dict with documents if outputFormat is specified """ await self._ensureAiObjectsInitialized() # Create separate operationId for detailed progress tracking workflowId = self.services.currentWorkflow.id if self.services.currentWorkflow else f"no-workflow-{int(time.time())}" aiOperationId = f"ai_documents_{workflowId}_{int(time.time())}" # Start progress tracking for this operation self.services.workflow.progressLogStart( aiOperationId, "AI call with documents", "Document Generation", f"Format: {outputFormat or 'text'}" ) try: if options is None or (hasattr(options, 'operationType') and options.operationType is None): # Use AI to determine parameters ONLY when truly needed (options=None OR operationType=None) self.services.workflow.progressLogUpdate(aiOperationId, 0.1, "Analyzing prompt parameters") options = await self._analyzePromptAndCreateOptions(prompt) # Handle image generation requests directly via generic path opType = getattr(options, "operationType", None) isImageRequest = (opType == OperationTypeEnum.IMAGE_GENERATE) if isImageRequest: # Image generation uses generic call path but bypasses document generation pipeline self.services.workflow.progressLogUpdate(aiOperationId, 0.4, "Calling AI for image generation") # Call via generic path (no looping for images) request = AiCallRequest( prompt=prompt, context="", options=options ) response = await self.aiObjects.call(request) # Extract image data from response if response.content: # For base64 format, return in expected format if outputFormat == "base64": result = { "success": True, "image_data": response.content, "documents": [{ "documentName": "generated_image.png", "documentData": response.content, "mimeType": "image/png", "title": title or "Generated Image" }] } else: # Return raw content for other formats result = response.content # Emit stats for image generation self.services.workflow.storeWorkflowStat( self.services.currentWorkflow, response, f"ai.generate.image" ) self.services.workflow.progressLogUpdate(aiOperationId, 0.9, "Image generated") self.services.workflow.progressLogFinish(aiOperationId, True) return result else: errorMsg = f"No image data returned: {response.content}" logger.error(f"Error in AI image generation: {errorMsg}") self.services.workflow.progressLogFinish(aiOperationId, False) return {"success": False, "error": errorMsg} # CRITICAL: For document generation with JSON templates, NEVER compress the prompt # Compressing would truncate the template structure and confuse the AI if outputFormat: # Document generation with structured output if not options: options = AiCallOptions() options.compressPrompt = False # JSON templates must NOT be truncated options.compressContext = False # Context also should not be compressed # Handle document generation with specific output format using unified approach if outputFormat: # Use unified generation method for all document generation if documents and len(documents) > 0: self.services.workflow.progressLogUpdate(aiOperationId, 0.2, f"Extracting content from {len(documents)} documents") extracted_content = await self.callAiText(prompt, documents, options, aiOperationId) else: self.services.workflow.progressLogUpdate(aiOperationId, 0.2, "Preparing for direct generation") extracted_content = None self.services.workflow.progressLogUpdate(aiOperationId, 0.3, "Building generation prompt") from modules.services.serviceGeneration.subPromptBuilderGeneration import buildGenerationPrompt # First call without continuation context generation_prompt = await buildGenerationPrompt(outputFormat, prompt, title, extracted_content, None) # Prepare prompt builder arguments for continuation promptArgs = { "outputFormat": outputFormat, "userPrompt": prompt, "title": title, "extracted_content": extracted_content } self.services.workflow.progressLogUpdate(aiOperationId, 0.4, "Calling AI for content generation") generated_json = await self._callAiWithLooping( generation_prompt, options, "document_generation", buildGenerationPrompt, promptArgs, aiOperationId ) self.services.workflow.progressLogUpdate(aiOperationId, 0.7, "Parsing generated JSON") # Parse the generated JSON (extract fenced/embedded JSON first) try: extracted_json = self.services.utils.jsonExtractString(generated_json) generated_data = json.loads(extracted_json) except json.JSONDecodeError as e: logger.error(f"Failed to parse generated JSON: {str(e)}") logger.error(f"JSON content length: {len(generated_json)}") logger.error(f"JSON content preview (last 200 chars): ...{generated_json[-200:]}") logger.error(f"JSON content around error position: {generated_json[max(0, e.pos-50):e.pos+50]}") # Write the problematic JSON to debug file self.services.utils.writeDebugFile(generated_json, "failed_json_parsing") self.services.workflow.progressLogFinish(aiOperationId, False) return {"success": False, "error": f"Generated content is not valid JSON: {str(e)}"} self.services.workflow.progressLogUpdate(aiOperationId, 0.8, f"Rendering to {outputFormat} format") # Render to final format using the existing renderer try: from modules.services.serviceGeneration.mainServiceGeneration import GenerationService generationService = GenerationService(self.services) rendered_content, mime_type = await generationService.renderReport( generated_data, outputFormat, title or "Generated Document", prompt, self ) # Build result in the expected format result = { "success": True, "content": generated_data, "documents": [{ "documentName": f"generated.{outputFormat}", "documentData": rendered_content, "mimeType": mime_type, "title": title or "Generated Document" }], "is_multi_file": False, "format": outputFormat, "title": title, "split_strategy": "single", "total_documents": 1, "processed_documents": 1 } # Log AI response for debugging self.services.utils.writeDebugFile(str(result), "document_generation_response", documents) self.services.workflow.progressLogFinish(aiOperationId, True) return result except Exception as e: logger.error(f"Error rendering document: {str(e)}") self.services.workflow.progressLogFinish(aiOperationId, False) return {"success": False, "error": f"Rendering failed: {str(e)}"} # Handle text calls (no output format specified) self.services.workflow.progressLogUpdate(aiOperationId, 0.5, "Processing text call") if documents: # Use document processing for text calls with documents result = await self.callAiText(prompt, documents, options, aiOperationId) else: # Use shared core function for direct text calls result = await self._callAiWithLooping(prompt, options, "text", None, None, aiOperationId) self.services.workflow.progressLogFinish(aiOperationId, True) return result except Exception as e: logger.error(f"Error in callAiDocuments: {str(e)}") self.services.workflow.progressLogFinish(aiOperationId, False) raise async def callAiText( self, prompt: str, documents: Optional[List[ChatDocument]], options: AiCallOptions, operationId: Optional[str] = None ) -> str: """ Handle text calls with document processing through ExtractionService. UNIFIED PROCESSING: Always use per-chunk processing for consistency. """ await self._ensureAiObjectsInitialized() return await self.extractionService.processDocumentsPerChunk(documents, prompt, self.aiObjects, options, operationId)