# 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 FLOW LOGIC VARIABLES: - jsonBase: str (merged JSON so far, starts empty) - lastValidCompletePart: str (fallback for failures) - mergeFailCount: int = 0 (max 3) FLOW: ┌─────────────────────────────────────────────────────────────────┐ │ 1. BUILD PROMPT │ │ - First: original prompt │ │ - Next: buildContinuationContext(lastRawResponse) │ ├─────────────────────────────────────────────────────────────────┤ │ 2. CALL AI → response fragment │ ├─────────────────────────────────────────────────────────────────┤ │ 4. MERGE jsonBase + response │ │ ├─ FAILS: repeat prompt, fails++ (if >=3 return fallback) │ │ └─ SUCCEEDS: try parse │ │ ├─ SUCCEEDS: FINISHED │ │ └─ FAILS: → step 5 │ ├─────────────────────────────────────────────────────────────────┤ │ 5. GET CONTEXTS (merge OK, parse failed) │ │ getContexts(mergedJson) → │ │ - If no cut point: overlapContext = "" │ │ - Store contexts for next iteration │ ├─────────────────────────────────────────────────────────────────┤ │ 6. DECIDE │ │ ├─ jsonParsingSuccess=true AND overlapContext="": │ │ │ FINISHED. return completePart │ │ ├─ jsonParsingSuccess=true AND overlapContext!="": │ │ │ CONTINUE, fails=0 │ │ └─ ELSE: repeat prompt, fails++ │ └─────────────────────────────────────────────────────────────────┘ """ import json import logging from typing import Dict, Any, List, Optional, Callable from modules.datamodels.datamodelAi import ( AiCallRequest, AiCallOptions ) from modules.datamodels.datamodelExtraction import ContentPart from modules.services.serviceAi.subJsonResponseHandling import JsonResponseHandler from modules.services.serviceAi.subLoopingUseCases import LoopingUseCaseRegistry from modules.workflows.processing.shared.stateTools import checkWorkflowStopped from modules.shared.jsonContinuation import getContexts from modules.shared.jsonUtils import buildContinuationContext, extractJsonString, tryParseJson from modules.shared.jsonUtils import tryParseJson from modules.shared.jsonUtils import closeJsonStructures from modules.shared.jsonUtils import stripCodeFences, normalizeJsonText 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 # JSON Base Iteration System: # - jsonBase: the merged JSON string (replaces accumulatedDirectJson array) # - After each iteration, new response is merged with jsonBase # - On merge success: check if complete, store contexts for next iteration # - On merge fail: retry with same prompt, increment fails jsonBase = None # Merged JSON string (starts None, set on first response) # Merge fail tracking - stop after 3 consecutive merge failures MAX_MERGE_FAILS = 3 mergeFailCount = 0 # Global counter for merge failures across entire loop lastValidCompletePart = None # Store last successfully parsed completePart for fallback # 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: # Extract templateStructure and basePrompt from promptArgs (REQUIRED) templateStructure = promptArgs.get("templateStructure") if not templateStructure: raise ValueError( f"templateStructure is REQUIRED in promptArgs for use case '{useCaseId}'. " "Prompt creation functions must return (prompt, templateStructure) tuple." ) basePrompt = promptArgs.get("basePrompt") if not basePrompt: # Fallback: use prompt parameter (should be the same) basePrompt = prompt logger.warning( f"basePrompt not found in promptArgs for use case '{useCaseId}', " "using prompt parameter instead. This may indicate a bug." ) # This is a continuation - build continuation context with raw JSON and rebuild prompt continuationContext = buildContinuationContext( allSections, lastRawResponse, useCaseId, templateStructure ) if not lastRawResponse: logger.warning(f"Iteration {iteration}: No previous response available for continuation!") # Store valid completePart from continuation context for fallback on merge failures # Use getContexts to check if completePart is parseable and store it if lastRawResponse and not lastValidCompletePart: try: contexts = getContexts(lastRawResponse) if contexts.jsonParsingSuccess and contexts.completePart: lastValidCompletePart = contexts.completePart logger.debug(f"Iteration {iteration}: Stored initial valid completePart ({len(lastValidCompletePart)} chars)") except Exception as e: logger.debug(f"Iteration {iteration}: Failed to extract completePart: {e}") # Unified prompt builder call: Continuation builders only need continuationContext, templateStructure, and basePrompt # All initial context (section, userPrompt, etc.) is already in basePrompt, so promptArgs is not needed # Extract templateStructure and basePrompt from promptArgs (they're explicit parameters) iterationPrompt = await promptBuilder( continuationContext=continuationContext, templateStructure=templateStructure, basePrompt=basePrompt ) 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 # Handle use cases that return JSON directly (no section extraction needed) # Check if use case supports direct return (all registered use cases do) if useCase and not useCase.requiresExtraction: # ===================================================================== # ITERATION FLOW (Simplified) # ===================================================================== # Step 4: MERGE jsonBase + new response # - FAILS: repeat prompt, increment fails cont (if >=3 return fallback) # - SUCCEEDS: try parse # - SUCCEEDS: FINISHED # - FAILS: proceed to Step 5 # Step 5: GET CONTEXTS (merge OK, parse failed) # - getContexts() with repair # - If no cut point: overlapContext = "" # Step 6: DECIDE # - jsonParsingSuccess=true AND overlapContext="": FINISHED # - jsonParsingSuccess=true AND overlapContext!="": continue, fails=0 # - ELSE: repeat prompt, increment fails count # ===================================================================== # STEP 4: MERGE jsonBase + new response # Use candidateJson to hold merged result until we confirm it's valid candidateJson = None if jsonBase is None: # First iteration - candidate is the current result candidateJson = result logger.debug(f"Iteration {iteration}: First response, candidateJson ({len(candidateJson)} chars)") else: # Merge jsonBase with new response logger.info(f"Iteration {iteration}: Merging jsonBase ({len(jsonBase)} chars) with new response ({len(result)} chars)") mergedJsonString, hasOverlap = JsonResponseHandler.mergeJsonStringsWithOverlap(jsonBase, result) if not hasOverlap: # MERGE FAILED - repeat prompt with unchanged jsonBase mergeFailCount += 1 logger.warning( f"Iteration {iteration}: Merge failed, no overlap found " f"(fail {mergeFailCount}/{MAX_MERGE_FAILS})" ) if mergeFailCount >= MAX_MERGE_FAILS: # Max failures reached - return last valid completePart logger.error( f"Iteration {iteration}: Max merge failures ({MAX_MERGE_FAILS}) reached, " "returning last valid completePart" ) if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, False) if lastValidCompletePart: try: extracted = extractJsonString(lastValidCompletePart) parsed, parseErr, _ = tryParseJson(extracted) if parseErr is None and parsed: normalized = self._normalizeJsonStructure(parsed, useCase) return json.dumps(normalized, indent=2, ensure_ascii=False) except Exception: pass return lastValidCompletePart else: # No valid fallback - return whatever we have return jsonBase if jsonBase else "" # Not at max failures - retry with same prompt (jsonBase unchanged) if iterationOperationId: self.services.chat.progressLogUpdate( iterationOperationId, 0.7, f"Merge failed ({mergeFailCount}/{MAX_MERGE_FAILS}), retrying" ) self.services.chat.progressLogFinish(iterationOperationId, True) continue # MERGE SUCCEEDED - set candidate (don't update jsonBase yet!) candidateJson = mergedJsonString logger.debug(f"Iteration {iteration}: Merge succeeded, candidateJson ({len(candidateJson)} chars)") # Try direct parse of candidate try: extracted = extractJsonString(candidateJson) parsed, parseErr, _ = tryParseJson(extracted) if parseErr is None and parsed: # Direct parse succeeded - FINISHED # Commit candidate to jsonBase jsonBase = candidateJson logger.info(f"Iteration {iteration}: Direct parse succeeded, JSON is complete") normalized = self._normalizeJsonStructure(parsed, useCase) result = json.dumps(normalized, indent=2, ensure_ascii=False) if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, True) if not useCase.finalResultHandler: raise ValueError( f"Use case '{useCaseId}' is missing required 'finalResultHandler' callback." ) return useCase.finalResultHandler( result, normalized, extracted, debugPrefix, self.services ) except Exception as e: logger.debug(f"Iteration {iteration}: Direct parse failed: {e}") # STEP 5: GET CONTEXTS (merge OK, parse failed = cut JSON) # Use candidateJson for context extraction contexts = getContexts(candidateJson) logger.debug( f"Iteration {iteration}: getContexts() -> " f"jsonParsingSuccess={contexts.jsonParsingSuccess}, " f"overlapContext={'\"\"' if not contexts.overlapContext else f'({len(contexts.overlapContext)} chars)'}" ) # STEP 6: DECIDE based on jsonParsingSuccess and overlapContext if contexts.jsonParsingSuccess and contexts.overlapContext == "": # JSON is complete (no cut point) - FINISHED # Use completePart for final result (closed, repaired JSON) # No more merging needed, so we don't need the cut version jsonBase = contexts.completePart logger.info(f"Iteration {iteration}: jsonParsingSuccess=true, overlapContext='', JSON complete") # Store and parse completePart lastValidCompletePart = contexts.completePart try: extracted = extractJsonString(contexts.completePart) parsed, parseErr, _ = tryParseJson(extracted) if parseErr is None and parsed: normalized = self._normalizeJsonStructure(parsed, useCase) result = json.dumps(normalized, indent=2, ensure_ascii=False) if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, True) if not useCase.finalResultHandler: raise ValueError( f"Use case '{useCaseId}' is missing required 'finalResultHandler' callback." ) return useCase.finalResultHandler( result, normalized, extracted, debugPrefix, self.services ) except Exception as e: logger.warning(f"Iteration {iteration}: Failed to parse completePart: {e}") # Fallback: return completePart as-is if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, True) return contexts.completePart elif contexts.jsonParsingSuccess and contexts.overlapContext != "": # JSON parseable but has cut point - CONTINUE to next iteration # CRITICAL: Use hierarchyContext (CUT json) as jsonBase for next merge! # - hierarchyContext = the truncated JSON at cut point (needed for overlap matching) # - completePart = closed JSON (for validation/fallback only) # The next AI fragment's overlap must match the CUT point, not closed structures jsonBase = contexts.hierarchyContext logger.info( f"Iteration {iteration}: jsonParsingSuccess=true, overlapContext not empty, " f"continuing iteration (jsonBase updated to hierarchyContext: {len(jsonBase)} chars)" ) # Store valid completePart as fallback (different from jsonBase!) lastValidCompletePart = contexts.completePart # Reset fail counter on successful progress mergeFailCount = 0 # Update lastRawResponse for continuation prompt building # Use the CUT version for prompt context as well lastRawResponse = jsonBase if iterationOperationId: self.services.chat.progressLogUpdate(iterationOperationId, 0.7, "JSON incomplete, requesting continuation") self.services.chat.progressLogFinish(iterationOperationId, True) continue else: # JSON not parseable after repair - repeat prompt, increment fails # Do NOT update jsonBase - keep previous valid state mergeFailCount += 1 logger.warning( f"Iteration {iteration}: jsonParsingSuccess=false, " f"repeat prompt (fail {mergeFailCount}/{MAX_MERGE_FAILS})" ) if mergeFailCount >= MAX_MERGE_FAILS: # Max failures reached - return last valid completePart logger.error( f"Iteration {iteration}: Max failures ({MAX_MERGE_FAILS}) reached, " "returning last valid completePart" ) if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, False) if lastValidCompletePart: try: extracted = extractJsonString(lastValidCompletePart) parsed, parseErr, _ = tryParseJson(extracted) if parseErr is None and parsed: normalized = self._normalizeJsonStructure(parsed, useCase) return json.dumps(normalized, indent=2, ensure_ascii=False) except Exception: pass return lastValidCompletePart else: return jsonBase if jsonBase else "" # Not at max - retry with same prompt # Do NOT update jsonBase or lastRawResponse - keep previous for retry if iterationOperationId: self.services.chat.progressLogUpdate( iterationOperationId, 0.7, f"Parse failed ({mergeFailCount}/{MAX_MERGE_FAILS}), retrying" ) self.services.chat.progressLogFinish(iterationOperationId, True) continue 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 should never be reached because all registered use cases # return early when JSON is complete. This 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 # 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 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 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, useCase) -> 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) useCase: LoopingUseCase instance with jsonNormalizer callback Returns: Normalized JSON structure """ # Use callback to normalize JSON structure (REQUIRED - no fallback) if not useCase or not useCase.jsonNormalizer: raise ValueError( f"Use case '{useCase.useCaseId if useCase else 'unknown'}' is missing required 'jsonNormalizer' callback. " "All use cases must provide a jsonNormalizer function." ) return useCase.jsonNormalizer(parsed, useCase.useCaseId)