# 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 .subJsonResponseHandling import JsonResponseHandler from .subLoopingUseCases import LoopingUseCaseRegistry from modules.workflows.processing.shared.stateTools import checkWorkflowStopped from modules.shared.jsonContinuation import getContexts from modules.shared.jsonUtils import buildContinuationContext, 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 = 10 iteration = 0 result = "" allSections = [] lastRawResponse = None # 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 MAX_CONSECUTIVE_EMPTY_RESPONSES = 3 consecutive_empty_responses = 0 # 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}") # Note: Stats are now stored centrally in callAi() - no need to duplicate here # 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(): consecutive_empty_responses += 1 logger.warning( "Iteration %s: Empty AI response (consecutive %s/%s) modelName=%s errorCount=%s", iteration, consecutive_empty_responses, MAX_CONSECUTIVE_EMPTY_RESPONSES, getattr(response, "modelName", None), getattr(response, "errorCount", None), ) if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, False) if consecutive_empty_responses >= MAX_CONSECUTIVE_EMPTY_RESPONSES: logger.error( "Stopping loop: %s consecutive empty responses from model", consecutive_empty_responses, ) break continue consecutive_empty_responses = 0 # 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 # NOTE: Do NOT update lastRawResponse here! # lastRawResponse should only be updated after successful merge # This ensures retry iterations use the correct base context # 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: parsed, parseErr, _ = tryParseJson(lastValidCompletePart) if parseErr is None: 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)") # Update lastRawResponse ONLY after we have a valid candidateJson # (first iteration or successful merge - NOT on merge failure!) # This ensures retry iterations use the correct base context lastRawResponse = candidateJson # Try direct parse of candidate (same pipeline as structure filling / getContexts) try: parsed, parseErr, extracted = tryParseJson(candidateJson) if parseErr is None: # 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) overlapInfo = "(empty=complete)" if contexts.overlapContext == "" else f"({len(contexts.overlapContext)} chars)" logger.debug( f"Iteration {iteration}: getContexts() -> " f"jsonParsingSuccess={contexts.jsonParsingSuccess}, " f"overlapContext={overlapInfo}" ) # STEP 6: DECIDE based on jsonParsingSuccess and overlapContext if contexts.jsonParsingSuccess and contexts.overlapContext == "": # getContexts and downstream must agree with tryParseJson (same as structure filling). logger.info(f"Iteration {iteration}: jsonParsingSuccess=true, overlapContext='', JSON complete") lastValidCompletePart = contexts.completePart try: parsed, parseErr, extracted = tryParseJson(contexts.completePart) if parseErr is not None: raise ValueError(str(parseErr)) normalized = self._normalizeJsonStructure(parsed, useCase) result = json.dumps(normalized, indent=2, ensure_ascii=False) jsonBase = contexts.completePart 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}: completePart not serializable after getContexts success: {e}" ) mergeFailCount += 1 if mergeFailCount >= MAX_MERGE_FAILS: logger.error( f"Iteration {iteration}: Max failures ({MAX_MERGE_FAILS}) " "after output pipeline mismatch" ) if iterationOperationId: self.services.chat.progressLogFinish(iterationOperationId, False) return jsonBase if jsonBase else "" if iterationOperationId: self.services.chat.progressLogUpdate( iterationOperationId, 0.7, f"Output pipeline failed ({mergeFailCount}/{MAX_MERGE_FAILS}), retrying", ) self.services.chat.progressLogFinish(iterationOperationId, True) continue 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: parsed, parseErr, _ = tryParseJson(lastValidCompletePart) if parseErr is None: 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})") # Prefer last repaired complete JSON from getContexts (raw `result` is only the last fragment). if lastValidCompletePart and useCase and not useCase.requiresExtraction: try: parsed, parseErr, extracted = tryParseJson(lastValidCompletePart) if parseErr is None: normalized = self._normalizeJsonStructure(parsed, useCase) out = json.dumps(normalized, indent=2, ensure_ascii=False) if useCase.finalResultHandler: logger.warning( "callAiWithLooping: max iterations — returning last valid completePart for %r", useCaseId, ) return useCase.finalResultHandler( out, normalized, extracted, debugPrefix, self.services ) except Exception as e: logger.debug("Max-iterations fallback on completePart failed: %s", e) logger.error( "End of callAiWithLooping without success for use case %r (iterations=%s, lastResultLen=%s)", useCaseId, iteration, len(result) if isinstance(result, str) else 0, ) 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)