# Copyright (c) 2025 Patrick Motsch # All rights reserved. """ AI Call Looping Module Handles AI calls with looping and repair logic, including: - Looping with JSON repair and continuation - KPI definition and tracking - Progress tracking and iteration management """ import json import logging from typing import Dict, Any, List, Optional, Callable from modules.datamodels.datamodelAi import AiCallRequest, AiCallOptions, OperationTypeEnum, PriorityEnum, ProcessingModeEnum, JsonAccumulationState from modules.datamodels.datamodelExtraction import ContentPart from modules.shared.jsonUtils import buildContinuationContext, extractJsonString, tryParseJson from modules.services.serviceAi.subJsonResponseHandling import JsonResponseHandler from modules.services.serviceAi.subLoopingUseCases import LoopingUseCaseRegistry from modules.workflows.processing.shared.stateTools import checkWorkflowStopped logger = logging.getLogger(__name__) class AiCallLooper: """Handles AI calls with looping and repair logic.""" def __init__(self, services, aiService, responseParser): """Initialize AiCallLooper with service center, AI service, and response parser access.""" self.services = services self.aiService = aiService self.responseParser = responseParser self.useCaseRegistry = LoopingUseCaseRegistry() # Initialize use case registry 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, userPrompt: Optional[str] = None, contentParts: Optional[List[ContentPart]] = None, # ARCHITECTURE: Support ContentParts for large content useCaseId: str = None # REQUIRED: Explicit use case ID - no auto-detection, no fallback ) -> 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 userPrompt: Optional user prompt for KPI definition contentParts: Optional content parts for first iteration useCaseId: REQUIRED: Explicit use case ID - no auto-detection, no fallback Returns: Complete AI response after all iterations """ # REQUIRED: useCaseId must be provided - no auto-detection, no fallback if not useCaseId: errorMsg = ( "useCaseId is REQUIRED for callAiWithLooping. " "No auto-detection - must explicitly specify use case ID. " f"Available use cases: {list(self.useCaseRegistry.useCases.keys())}" ) logger.error(errorMsg) raise ValueError(errorMsg) # Validate use case exists useCase = self.useCaseRegistry.get(useCaseId) if not useCase: errorMsg = ( f"Use case '{useCaseId}' not found in registry. " f"Available use cases: {list(self.useCaseRegistry.useCases.keys())}" ) logger.error(errorMsg) raise ValueError(errorMsg) maxIterations = 50 # Prevent infinite loops iteration = 0 allSections = [] # Accumulate all sections across iterations lastRawResponse = None # Store last raw JSON response for continuation accumulatedDirectJson = [] # Accumulate JSON strings for direct return use cases (chapter_structure, code_structure) # Get parent operation ID for iteration operations (parentId should be operationId, not log entry ID) parentOperationId = operationId # Use the parent's operationId directly while iteration < maxIterations: iteration += 1 # Create separate operation for each iteration with parent reference iterationOperationId = None if operationId: iterationOperationId = f"{operationId}_iter_{iteration}" self.services.chat.progressLogStart( iterationOperationId, "AI Call", f"Iteration {iteration}", "", parentOperationId=parentOperationId ) # Build iteration prompt # CRITICAL: Build continuation prompt if we have sections OR if we have a previous response (even if broken) # This ensures continuation prompts are built even when JSON is so broken that no sections can be extracted if (len(allSections) > 0 or lastRawResponse) and promptBuilder and promptArgs: # This is a continuation - build continuation context with raw JSON and rebuild prompt continuationContext = buildContinuationContext(allSections, lastRawResponse, useCaseId) if not lastRawResponse: logger.warning(f"Iteration {iteration}: No previous response available for continuation!") # Unified prompt builder call: All prompt builders accept continuationContext and **kwargs # Each builder extracts only the parameters it needs from kwargs # This ensures consistent architecture across all use cases if not promptArgs.get('services') and hasattr(self, 'services'): promptArgs['services'] = self.services iterationPrompt = await promptBuilder(continuationContext=continuationContext, **promptArgs) else: # First iteration - use original prompt iterationPrompt = prompt # Make AI call try: checkWorkflowStopped(self.services) if iterationOperationId: self.services.chat.progressLogUpdate(iterationOperationId, 0.3, "Calling AI model") # ARCHITECTURE: Pass ContentParts directly to AiCallRequest # This allows model-aware chunking to handle large content properly # ContentParts are only passed in first iteration (continuations don't need them) request = AiCallRequest( prompt=iterationPrompt, context="", options=options, contentParts=contentParts if iteration == 1 else None # Only pass ContentParts in first iteration ) # Write the ACTUAL prompt sent to AI # For section content generation: write prompt for first iteration and continuation iterations # For document generation: write prompt for each iteration isSectionContent = "_section_" in debugPrefix if iteration == 1: self.services.utils.writeDebugFile(iterationPrompt, f"{debugPrefix}_prompt") elif isSectionContent: # Save continuation prompts for section_content debugging self.services.utils.writeDebugFile(iterationPrompt, f"{debugPrefix}_prompt_iteration_{iteration}") else: # Document generation - save all iteration prompts self.services.utils.writeDebugFile(iterationPrompt, f"{debugPrefix}_prompt_iteration_{iteration}") response = await self.aiService.callAi(request) result = response.content # Track bytes for progress reporting bytesReceived = len(result.encode('utf-8')) if result else 0 totalBytesSoFar = sum(len(section.get('content', '').encode('utf-8')) if isinstance(section.get('content'), str) else 0 for section in allSections) + bytesReceived # Update progress after AI call with byte information if iterationOperationId: # Format bytes for display (kB or MB) if totalBytesSoFar < 1024: bytesDisplay = f"{totalBytesSoFar}B" elif totalBytesSoFar < 1024 * 1024: bytesDisplay = f"{totalBytesSoFar / 1024:.1f}kB" else: bytesDisplay = f"{totalBytesSoFar / (1024 * 1024):.1f}MB" self.services.chat.progressLogUpdate(iterationOperationId, 0.6, f"AI response received ({bytesDisplay})") # Write raw AI response to debug file # For section content generation: write response for first iteration and continuation iterations # For document generation: write response for each iteration if iteration == 1: self.services.utils.writeDebugFile(result, f"{debugPrefix}_response") elif isSectionContent: # Save continuation responses for section_content debugging self.services.utils.writeDebugFile(result, f"{debugPrefix}_response_iteration_{iteration}") else: # Document generation - save all iteration responses self.services.utils.writeDebugFile(result, f"{debugPrefix}_response_iteration_{iteration}") # Emit stats for this iteration (only if workflow exists and has id) if self.services.workflow and hasattr(self.services.workflow, 'id') and self.services.workflow.id: try: self.services.chat.storeWorkflowStat( self.services.workflow, response, f"ai.call.{debugPrefix}.iteration_{iteration}" ) except Exception as statError: # Don't break the main loop if stat storage fails logger.warning(f"Failed to store workflow stat: {str(statError)}") # Check for error response using generic error detection (errorCount > 0 or modelName == "error") if hasattr(response, 'errorCount') and response.errorCount > 0: errorMsg = f"Iteration {iteration}: Error response detected (errorCount={response.errorCount}), stopping loop: {result[:200] if result else 'empty'}" logger.error(errorMsg) break if hasattr(response, 'modelName') and response.modelName == "error": errorMsg = f"Iteration {iteration}: Error response detected (modelName=error), stopping loop: {result[:200] if result else 'empty'}" logger.error(errorMsg) break if not result or not result.strip(): logger.warning(f"Iteration {iteration}: Empty response, stopping") break # Check if this is a text response (not document generation) # Text responses don't need JSON parsing - return immediately after first successful response isTextResponse = (promptBuilder is None and promptArgs is None) or debugPrefix == "text" if isTextResponse: # For text responses, return the text immediately - no JSON parsing needed logger.info(f"Iteration {iteration}: Text response received, returning immediately") if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, True) return result # Store raw response for continuation (even if broken) lastRawResponse = result # Parse JSON for use case handling parsedJsonForUseCase = None extractedJsonForUseCase = None try: extractedJsonForUseCase = extractJsonString(result) parsedJson, parseError, _ = tryParseJson(extractedJsonForUseCase) if parseError is None and parsedJson: parsedJsonForUseCase = parsedJson except Exception: pass # Handle use cases that return JSON directly (no section extraction needed) directReturnUseCases = ["section_content", "chapter_structure", "code_structure", "code_content"] if useCaseId in directReturnUseCases: # For chapter_structure, code_structure, section_content, and code_content, check completeness and support looping loopingUseCases = ["chapter_structure", "code_structure", "section_content", "code_content"] if useCaseId in loopingUseCases: # CRITICAL: Check if JSON string is incomplete BEFORE parsing # If JSON is truncated, it will be closed for parsing, making it appear complete # So we need to check the original string, not the parsed JSON isStringIncomplete = self._isJsonStringIncomplete(extractedJsonForUseCase if extractedJsonForUseCase else result) # If parsing failed (e.g., invalid JSON with comments or truncated JSON), continue looping to get valid JSON if not parsedJsonForUseCase: logger.info(f"Iteration {iteration}: Use case '{useCaseId}' - JSON parsing failed (likely incomplete/truncated), continuing iteration to complete") # Accumulate response for merging in next iteration accumulatedDirectJson.append(result) # Continue to next iteration - continuation prompt builder will handle the rest if iterationOperationId: self.services.chat.progressLogUpdate(iterationOperationId, 0.7, "JSON incomplete, requesting continuation") self.services.chat.progressLogFinish(iterationOperationId, True) continue # Check completeness: Use string-based check if available, otherwise fall back to parsed JSON check if isStringIncomplete: isComplete = False else: # Check completeness if we have parsed JSON isComplete = JsonResponseHandler.isJsonComplete(parsedJsonForUseCase) if not isComplete: logger.warning(f"Iteration {iteration}: Use case '{useCaseId}' - JSON is incomplete, continuing for continuation") # Accumulate response for merging in next iteration accumulatedDirectJson.append(result) # Continue to next iteration - continuation prompt builder will handle the rest if iterationOperationId: self.services.chat.progressLogUpdate(iterationOperationId, 0.7, "JSON incomplete, requesting continuation") self.services.chat.progressLogFinish(iterationOperationId, True) continue else: # JSON is complete - merge accumulated responses if any if accumulatedDirectJson: logger.info(f"Iteration {iteration}: Merging {len(accumulatedDirectJson) + 1} accumulated responses") # Use generic data-based merging for all use cases try: # Strategy: Merge strings first for incomplete JSON, then parse and merge parsed objects # This ensures incomplete JSON from part 1 is preserved allJsonStrings = accumulatedDirectJson + [result] # Step 1: Merge all JSON strings using existing overlap detection mergedJsonString = allJsonStrings[0] if allJsonStrings else "" hasOverlap = True # Track if any overlap was found for jsonStr in allJsonStrings[1:]: mergedJsonString, hasOverlapInMerge = JsonResponseHandler.mergeJsonStringsWithOverlap(mergedJsonString, jsonStr) # If no overlap found in any merge, stop iterations if not hasOverlapInMerge: hasOverlap = False logger.info(f"Iteration {iteration}: No overlap found during merge - stopping iterations and closing JSON") break # If no overlap was found, mark as complete and use closed JSON if not hasOverlap: isComplete = True # JSON is already closed by mergeJsonStringsWithOverlap when no overlap # Use the merged (closed) JSON string directly result = mergedJsonString # Try to parse it to get parsedJsonForUseCase try: extracted = extractJsonString(mergedJsonString) parsed, parseErr, _ = tryParseJson(extracted) if parseErr is None and parsed: normalized = self._normalizeJsonStructure(parsed, useCaseId) parsedJsonForUseCase = normalized result = json.dumps(normalized, indent=2, ensure_ascii=False) except Exception: pass # Use string result if parsing fails else: # Overlap found - continue with normal processing # Step 2: Try to parse the merged string extracted = extractJsonString(mergedJsonString) parsed, parseErr, _ = tryParseJson(extracted) if parseErr is None and parsed: # Parsing succeeded - normalize and use normalized = self._normalizeJsonStructure(parsed, useCaseId) parsedJsonForUseCase = normalized result = json.dumps(normalized, indent=2, ensure_ascii=False) else: # Parsing failed - try to extract partial data using Deep-Structure-Merging # This fallback works for all use cases: parse what we can from each part allParsed = [] for jsonStr in allJsonStrings: extracted = extractJsonString(jsonStr) parsed, parseErr, _ = tryParseJson(extracted) if parseErr is None and parsed: normalized = self._normalizeJsonStructure(parsed, useCaseId) allParsed.append(normalized) if allParsed: # Use mergeDeepStructures for intelligent merging across all use cases if len(allParsed) > 1: mergedJsonObj = allParsed[0] for nextObj in allParsed[1:]: mergedJsonObj = JsonResponseHandler.mergeDeepStructures( mergedJsonObj, nextObj, iteration, f"{useCaseId}.merge" ) else: mergedJsonObj = allParsed[0] parsedJsonForUseCase = mergedJsonObj result = json.dumps(mergedJsonObj, indent=2, ensure_ascii=False) else: # All parsing failed - use string merge result result = mergedJsonString except Exception as e: logger.warning(f"Failed data-based merge, falling back to string merging: {e}") # Fallback to string merging mergedJsonString = accumulatedDirectJson[0] if accumulatedDirectJson else result hasOverlap = True # Track if any overlap was found for prevJson in accumulatedDirectJson[1:]: mergedJsonString, hasOverlapInMerge = JsonResponseHandler.mergeJsonStringsWithOverlap(mergedJsonString, prevJson) if not hasOverlapInMerge: hasOverlap = False logger.info(f"Iteration {iteration}: No overlap found during fallback merge - stopping iterations") break if hasOverlap: mergedJsonString, hasOverlapInMerge = JsonResponseHandler.mergeJsonStringsWithOverlap(mergedJsonString, result) if not hasOverlapInMerge: hasOverlap = False logger.info(f"Iteration {iteration}: No overlap found in final fallback merge - stopping iterations") result = mergedJsonString # If no overlap was found, mark as complete and use closed JSON if not hasOverlap: isComplete = True # JSON is already closed by mergeJsonStringsWithOverlap when no overlap # Try to parse it to get parsedJsonForUseCase try: extractedMerged = extractJsonString(result) parsedMerged, parseError, _ = tryParseJson(extractedMerged) if parseError is None and parsedMerged: parsedJsonForUseCase = parsedMerged except Exception: pass # Use string result if parsing fails # Try to parse the string-merged result try: extractedMerged = extractJsonString(result) parsedMerged, parseError, _ = tryParseJson(extractedMerged) if parseError is None and parsedMerged: parsedJsonForUseCase = parsedMerged except Exception: pass logger.info(f"Iteration {iteration}: Use case '{useCaseId}' - JSON is complete") logger.info(f"Iteration {iteration}: Use case '{useCaseId}' - returning JSON directly") if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, True) # For section_content, return raw result to allow merging of multiple JSON blocks # The merging logic in subStructureFilling.py will handle extraction and merging if useCaseId == "section_content": final_json = result # Return raw response to preserve all JSON blocks # Write final merged result for section_content (overwrites iteration 1 response with complete merged result) self.services.utils.writeDebugFile(final_json, f"{debugPrefix}_response") else: final_json = json.dumps(parsedJsonForUseCase, indent=2, ensure_ascii=False) if parsedJsonForUseCase else (extractedJsonForUseCase or result) # Write final result for chapter structure and code structure if useCaseId in ["chapter_structure", "code_structure"]: self.services.utils.writeDebugFile(final_json, f"{debugPrefix}_final_result") return final_json except Exception as e: logger.error(f"Error in AI call iteration {iteration}: {str(e)}") if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, False) break if iteration >= maxIterations: logger.warning(f"AI call stopped after maximum iterations ({maxIterations})") # This code path is never reached because all use cases are in directReturnUseCases # and return early at line 417. This code would only execute for use cases that # require section extraction, but no such use cases are currently registered. logger.error(f"Unexpected code path: reached end of loop without return for use case '{useCaseId}'") return result if result else "" def _isJsonStringIncomplete(self, jsonString: str) -> bool: """ Check if JSON string is incomplete (truncated) BEFORE closing/parsing. This is critical because if JSON is truncated, closing it makes it appear complete, but we need to detect the truncation to continue iteration. Args: jsonString: JSON string to check Returns: True if JSON string appears incomplete/truncated, False otherwise """ if not jsonString or not jsonString.strip(): return False from modules.shared.jsonUtils import stripCodeFences, normalizeJsonText # Normalize JSON string normalized = stripCodeFences(normalizeJsonText(jsonString)).strip() if not normalized: return False # Find first '{' or '[' to start startIdx = -1 for i, char in enumerate(normalized): if char in '{[': startIdx = i break if startIdx == -1: return False jsonContent = normalized[startIdx:] # Check if structures are balanced (all opened structures are closed) braceCount = 0 bracketCount = 0 inString = False escapeNext = False for char in jsonContent: if escapeNext: escapeNext = False continue if char == '\\': escapeNext = True continue if char == '"': inString = not inString continue if not inString: if char == '{': braceCount += 1 elif char == '}': braceCount -= 1 elif char == '[': bracketCount += 1 elif char == ']': bracketCount -= 1 # If structures are unbalanced, JSON is incomplete if braceCount > 0 or bracketCount > 0: return True # Check if JSON ends with incomplete value (e.g., unclosed string, incomplete number, trailing comma) trimmed = jsonContent.rstrip() if not trimmed: return False # Check for trailing comma (might indicate incomplete) if trimmed.endswith(','): # Trailing comma might indicate incomplete, but could also be valid # Check if there's a closing bracket/brace after the comma return False # Trailing comma alone doesn't mean incomplete # Check if ends with incomplete string (odd number of quotes) quoteCount = jsonContent.count('"') if quoteCount % 2 == 1: # Odd number of quotes - string is not closed return True # Check if ends mid-value (e.g., ends with "417 instead of "4170. 41719"]) # Look for patterns that suggest truncation: # - Ends with incomplete number (e.g., "417) # - Ends with incomplete array element (e.g., ["417) # - Ends with incomplete object property (e.g., {"key": "val) # If JSON parses successfully without closing, it's complete from modules.shared.jsonUtils import tryParseJson parsed, parseErr, _ = tryParseJson(jsonContent) if parseErr is None: # Parses successfully - it's complete return False # If it doesn't parse, try closing it and see if that helps from modules.shared.jsonUtils import closeJsonStructures closed = closeJsonStructures(jsonContent) parsedClosed, parseErrClosed, _ = tryParseJson(closed) if parseErrClosed is None: # Only parses after closing - it was incomplete return True # Doesn't parse even after closing - might be malformed, but assume incomplete to be safe return True def _normalizeJsonStructure(self, parsed: Any, useCaseId: str) -> Any: """ Normalize JSON structure to ensure consistent format before merging. Handles different response formats and converts them to expected structure. Args: parsed: Parsed JSON object (can be dict, list, or primitive) useCaseId: Use case ID to determine expected structure Returns: Normalized JSON structure """ # For section_content, expect {"elements": [...]} structure if useCaseId == "section_content": if isinstance(parsed, list): # Check if list contains strings (invalid format) or element objects if parsed and isinstance(parsed[0], str): # Invalid format - list of strings instead of elements # Try to convert strings to paragraph elements as fallback # This can happen if AI returns raw text instead of structured JSON logger.debug(f"Received list of strings instead of elements array, converting to paragraph elements") elements = [] for text in parsed: if isinstance(text, str) and text.strip(): elements.append({ "type": "paragraph", "content": { "text": text.strip() } }) return {"elements": elements} if elements else {"elements": []} else: # Convert plain list of elements to elements structure return {"elements": parsed} elif isinstance(parsed, dict): # If it already has "elements", return as-is if "elements" in parsed: return parsed # If it has "type" and looks like an element, wrap in elements array elif parsed.get("type"): return {"elements": [parsed]} # Otherwise, assume it's already in correct format else: return parsed # For other use cases, return as-is (they have their own structures) return parsed