gateway/modules/services/serviceAi/subAiCallLooping.py
ValueOn AG 64590aa61e fixes
2026-01-04 20:01:34 +01:00

589 lines
34 KiB
Python

# 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