866 lines
48 KiB
Python
866 lines
48 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
|
|
documentMetadata = None # Store document metadata (title, filename) from first iteration
|
|
accumulationState = None # Track accumulation state for string accumulation
|
|
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)
|
|
if not lastRawResponse:
|
|
logger.warning(f"Iteration {iteration}: No previous response available for continuation!")
|
|
|
|
# For section_content, pass all promptArgs (it uses buildSectionPromptWithContinuation which needs all args)
|
|
# For other use cases (chapter_structure, code_structure), filter to only accepted parameters
|
|
if useCaseId == "section_content":
|
|
# Pass all promptArgs plus continuationContext for section_content
|
|
iterationPrompt = await promptBuilder(**promptArgs, continuationContext=continuationContext)
|
|
else:
|
|
# Filter promptArgs to only include parameters that buildGenerationPrompt accepts
|
|
# buildGenerationPrompt accepts: outputFormat, userPrompt, title, extracted_content, continuationContext, services
|
|
filteredPromptArgs = {
|
|
k: v for k, v in promptArgs.items()
|
|
if k in ['outputFormat', 'userPrompt', 'title', 'extracted_content', 'services']
|
|
}
|
|
# Always include services if available
|
|
if not filteredPromptArgs.get('services') and hasattr(self, 'services'):
|
|
filteredPromptArgs['services'] = self.services
|
|
|
|
# Rebuild prompt with continuation context using the provided prompt builder
|
|
iterationPrompt = await promptBuilder(**filteredPromptArgs, continuationContext=continuationContext)
|
|
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", "image_batch"]
|
|
if useCaseId in directReturnUseCases:
|
|
# For chapter_structure, code_structure, and section_content, check completeness and support looping
|
|
loopingUseCases = ["chapter_structure", "code_structure", "section_content"]
|
|
if useCaseId in loopingUseCases:
|
|
# 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 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:
|
|
# Parse all accumulated JSON strings and current response
|
|
allParsed = []
|
|
for jsonStr in accumulatedDirectJson + [result]:
|
|
extracted = extractJsonString(jsonStr)
|
|
parsed, parseErr, _ = tryParseJson(extracted)
|
|
if parseErr is None and parsed:
|
|
# Normalize structure: ensure consistent format
|
|
normalized = self._normalizeJsonStructure(parsed, useCaseId)
|
|
allParsed.append(normalized)
|
|
|
|
if allParsed and len(allParsed) > 1:
|
|
# Generic recursive merge of parsed JSON objects
|
|
mergedJsonObj = self._mergeJsonObjectsRecursively(allParsed)
|
|
|
|
# Reconstruct merged JSON string
|
|
mergedJsonString = json.dumps(mergedJsonObj, indent=2, ensure_ascii=False)
|
|
parsedJsonForUseCase = mergedJsonObj
|
|
result = mergedJsonString
|
|
logger.info(f"Successfully merged {len(accumulatedDirectJson) + 1} JSON fragments using generic recursive merge")
|
|
elif allParsed:
|
|
# Only one parsed JSON, use it directly
|
|
parsedJsonForUseCase = allParsed[0]
|
|
result = json.dumps(parsedJsonForUseCase, indent=2, ensure_ascii=False)
|
|
else:
|
|
# Fallback to string merging if parsing fails
|
|
logger.warning("Failed to parse all JSON fragments for data-based merge, falling back to string merging")
|
|
mergedJsonString = accumulatedDirectJson[0] if accumulatedDirectJson else result
|
|
for prevJson in accumulatedDirectJson[1:]:
|
|
mergedJsonString = JsonResponseHandler.mergeJsonStringsWithOverlap(mergedJsonString, prevJson)
|
|
mergedJsonString = JsonResponseHandler.mergeJsonStringsWithOverlap(mergedJsonString, 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
|
|
for prevJson in accumulatedDirectJson[1:]:
|
|
mergedJsonString = JsonResponseHandler.mergeJsonStringsWithOverlap(mergedJsonString, prevJson)
|
|
mergedJsonString = JsonResponseHandler.mergeJsonStringsWithOverlap(mergedJsonString, result)
|
|
result = mergedJsonString
|
|
|
|
# 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
|
|
|
|
# Extract sections from response (handles both valid and broken JSON)
|
|
# Only for document generation (JSON responses)
|
|
# CRITICAL: Pass allSections and accumulationState to enable string accumulation
|
|
extractedSections, wasJsonComplete, parsedResult, accumulationState = self.responseParser.extractSectionsFromResponse(
|
|
result, iteration, debugPrefix, allSections, accumulationState
|
|
)
|
|
|
|
# CRITICAL: Merge sections BEFORE KPI validation
|
|
# This ensures sections are preserved even if KPI validation fails
|
|
if extractedSections:
|
|
allSections = JsonResponseHandler.mergeSectionsIntelligently(allSections, extractedSections, iteration)
|
|
|
|
# Define KPIs if we just entered accumulation mode (iteration 1, incomplete JSON)
|
|
if accumulationState and accumulationState.isAccumulationMode and iteration == 1 and not accumulationState.kpis:
|
|
logger.info(f"Iteration {iteration}: Defining KPIs for accumulation tracking")
|
|
continuationContext = buildContinuationContext(allSections, result)
|
|
# Pass raw response string from first iteration for KPI definition
|
|
kpiDefinitions = await self._defineKpisFromPrompt(
|
|
userPrompt or prompt,
|
|
result, # Pass raw JSON string from first iteration
|
|
continuationContext,
|
|
debugPrefix
|
|
)
|
|
# Initialize KPIs with currentValue = 0
|
|
accumulationState.kpis = [{**kpi, "currentValue": 0} for kpi in kpiDefinitions]
|
|
logger.info(f"Defined {len(accumulationState.kpis)} KPIs: {[kpi.get('id') for kpi in accumulationState.kpis]}")
|
|
|
|
# Extract and validate KPIs (if in accumulation mode with KPIs defined)
|
|
if accumulationState and accumulationState.isAccumulationMode and accumulationState.kpis:
|
|
# For KPI extraction, prefer accumulated JSON string over repaired JSON
|
|
# because repairBrokenJson may lose data (e.g., empty rows array when JSON is incomplete)
|
|
updatedKpis = []
|
|
|
|
# First try to extract from parsedResult (repaired JSON)
|
|
if parsedResult:
|
|
try:
|
|
updatedKpis = JsonResponseHandler.extractKpiValuesFromJson(
|
|
parsedResult,
|
|
accumulationState.kpis
|
|
)
|
|
# Check if we got meaningful values (non-zero)
|
|
hasValidValues = any(kpi.get("currentValue", 0) > 0 for kpi in updatedKpis)
|
|
if not hasValidValues and accumulationState.accumulatedJsonString:
|
|
# Repaired JSON has empty values, try accumulated string
|
|
logger.debug("Repaired JSON has empty KPI values, trying accumulated JSON string")
|
|
updatedKpis = JsonResponseHandler.extractKpiValuesFromIncompleteJson(
|
|
accumulationState.accumulatedJsonString,
|
|
accumulationState.kpis
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"Error extracting KPIs from parsedResult: {e}")
|
|
updatedKpis = []
|
|
|
|
# If no parsedResult or extraction failed, try accumulated string
|
|
if not updatedKpis and accumulationState.accumulatedJsonString:
|
|
try:
|
|
updatedKpis = JsonResponseHandler.extractKpiValuesFromIncompleteJson(
|
|
accumulationState.accumulatedJsonString,
|
|
accumulationState.kpis
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"Error extracting KPIs from accumulated JSON string: {e}")
|
|
updatedKpis = []
|
|
|
|
if updatedKpis:
|
|
shouldProceed, reason = JsonResponseHandler.validateKpiProgression(
|
|
accumulationState,
|
|
updatedKpis
|
|
)
|
|
|
|
if not shouldProceed:
|
|
logger.warning(f"Iteration {iteration}: KPI validation failed: {reason}")
|
|
if iterationOperationId:
|
|
self.services.chat.progressLogFinish(iterationOperationId, False)
|
|
if operationId:
|
|
self.services.chat.progressLogUpdate(operationId, 0.9, f"KPI validation failed: {reason} ({iteration} iterations)")
|
|
break
|
|
|
|
# Update KPIs in accumulation state
|
|
accumulationState.kpis = updatedKpis
|
|
logger.info(f"Iteration {iteration}: KPIs updated: {[(kpi.get('id'), kpi.get('currentValue')) for kpi in updatedKpis]}")
|
|
|
|
# Check if all KPIs completed
|
|
allCompleted = True
|
|
for kpi in updatedKpis:
|
|
targetValue = kpi.get("targetValue", 0)
|
|
currentValue = kpi.get("currentValue", 0)
|
|
if currentValue < targetValue:
|
|
allCompleted = False
|
|
break
|
|
|
|
if allCompleted:
|
|
logger.info(f"Iteration {iteration}: All KPIs completed, finishing accumulation")
|
|
wasJsonComplete = True # Mark as complete to exit loop
|
|
|
|
# CRITICAL: Handle JSON fragments (continuation content)
|
|
# Fragment merging happens inside extractSectionsFromResponse
|
|
# If merge fails (returns wasJsonComplete=True), stop iterations and complete JSON
|
|
if not extractedSections and allSections:
|
|
if wasJsonComplete:
|
|
# Merge failed - stop iterations, complete JSON with available data
|
|
logger.error(f"Iteration {iteration}: ❌ MERGE FAILED - Stopping iterations, completing JSON with available data")
|
|
if iterationOperationId:
|
|
self.services.chat.progressLogFinish(iterationOperationId, False)
|
|
if operationId:
|
|
self.services.chat.progressLogUpdate(operationId, 0.9, f"Merge failed, completing JSON ({iteration} iterations)")
|
|
break
|
|
|
|
# Fragment was detected and merged successfully
|
|
logger.info(f"Iteration {iteration}: JSON fragment detected and merged, continuing")
|
|
# Don't break - fragment was merged, continue to get more content if needed
|
|
# Check if we should continue based on JSON completeness
|
|
shouldContinue = self.responseParser.shouldContinueGeneration(
|
|
allSections,
|
|
iteration,
|
|
wasJsonComplete,
|
|
result
|
|
)
|
|
if shouldContinue:
|
|
if iterationOperationId:
|
|
self.services.chat.progressLogUpdate(iterationOperationId, 0.8, "Fragment merged, continuing")
|
|
self.services.chat.progressLogFinish(iterationOperationId, True)
|
|
continue
|
|
else:
|
|
# Done - fragment was merged and JSON is complete
|
|
if iterationOperationId:
|
|
self.services.chat.progressLogFinish(iterationOperationId, True)
|
|
if operationId:
|
|
self.services.chat.progressLogUpdate(operationId, 0.95, f"Generation complete ({iteration} iterations, fragment merged)")
|
|
logger.info(f"Generation complete after {iteration} iterations: fragment merged")
|
|
break
|
|
|
|
# Extract document metadata from first iteration if available
|
|
if iteration == 1 and parsedResult and not documentMetadata:
|
|
documentMetadata = self.responseParser.extractDocumentMetadata(parsedResult)
|
|
|
|
# Update progress after parsing
|
|
if iterationOperationId:
|
|
if extractedSections:
|
|
self.services.chat.progressLogUpdate(iterationOperationId, 0.8, f"Extracted {len(extractedSections)} sections")
|
|
|
|
if not extractedSections:
|
|
# CRITICAL: If JSON was incomplete/broken, continue even if no sections extracted
|
|
# This allows the AI to retry and complete the broken JSON
|
|
if not wasJsonComplete:
|
|
logger.warning(f"Iteration {iteration}: No sections extracted from broken JSON, continuing for another attempt")
|
|
continue
|
|
# If JSON was complete but no sections extracted - check if it was a fragment
|
|
# Fragments are handled above, so if we get here and it's complete, it's an error
|
|
logger.warning(f"Iteration {iteration}: No sections extracted from complete JSON, stopping")
|
|
break
|
|
|
|
# NOTE: Section merging now happens BEFORE KPI validation (see above)
|
|
# This ensures sections are preserved even if KPI validation fails
|
|
|
|
# Calculate total bytes in merged content for progress display
|
|
merged_json_str = json.dumps(allSections, indent=2, ensure_ascii=False)
|
|
totalBytesGenerated = len(merged_json_str.encode('utf-8'))
|
|
|
|
# Update main operation with byte progress
|
|
if operationId:
|
|
# Format bytes for display
|
|
if totalBytesGenerated < 1024:
|
|
bytesDisplay = f"{totalBytesGenerated}B"
|
|
elif totalBytesGenerated < 1024 * 1024:
|
|
bytesDisplay = f"{totalBytesGenerated / 1024:.1f}kB"
|
|
else:
|
|
bytesDisplay = f"{totalBytesGenerated / (1024 * 1024):.1f}MB"
|
|
# Estimate progress based on iterations (rough estimate)
|
|
estimatedProgress = min(0.9, 0.4 + (iteration * 0.1))
|
|
self.services.chat.progressLogUpdate(operationId, estimatedProgress, f"Pipeline: {bytesDisplay} (iteration {iteration})")
|
|
|
|
# Log merged sections for debugging
|
|
# For section content generation: skip merged sections debug files (only one prompt/response needed)
|
|
isSectionContent = "_section_" in debugPrefix
|
|
if not isSectionContent:
|
|
self.services.utils.writeDebugFile(merged_json_str, f"{debugPrefix}_merged_sections_iteration_{iteration}")
|
|
|
|
# Check if we should continue (completion detection)
|
|
# Simple logic: JSON completeness determines continuation
|
|
shouldContinue = self.responseParser.shouldContinueGeneration(
|
|
allSections,
|
|
iteration,
|
|
wasJsonComplete,
|
|
result
|
|
)
|
|
|
|
if shouldContinue:
|
|
# Finish iteration operation (will continue with next iteration)
|
|
if iterationOperationId:
|
|
# Show byte progress in iteration completion
|
|
iterBytes = len(result.encode('utf-8')) if result else 0
|
|
if iterBytes < 1024:
|
|
iterBytesDisplay = f"{iterBytes}B"
|
|
elif iterBytes < 1024 * 1024:
|
|
iterBytesDisplay = f"{iterBytes / 1024:.1f}kB"
|
|
else:
|
|
iterBytesDisplay = f"{iterBytes / (1024 * 1024):.1f}MB"
|
|
self.services.chat.progressLogUpdate(iterationOperationId, 0.95, f"Completed ({iterBytesDisplay})")
|
|
self.services.chat.progressLogFinish(iterationOperationId, True)
|
|
continue
|
|
else:
|
|
# Done - finish iteration and update main operation
|
|
if iterationOperationId:
|
|
# Show final byte count
|
|
finalBytes = len(merged_json_str.encode('utf-8'))
|
|
if finalBytes < 1024:
|
|
finalBytesDisplay = f"{finalBytes}B"
|
|
elif finalBytes < 1024 * 1024:
|
|
finalBytesDisplay = f"{finalBytes / 1024:.1f}kB"
|
|
else:
|
|
finalBytesDisplay = f"{finalBytes / (1024 * 1024):.1f}MB"
|
|
self.services.chat.progressLogUpdate(iterationOperationId, 0.95, f"Complete ({finalBytesDisplay})")
|
|
self.services.chat.progressLogFinish(iterationOperationId, True)
|
|
if operationId:
|
|
# Show final size in main operation
|
|
finalBytes = len(merged_json_str.encode('utf-8'))
|
|
if finalBytes < 1024:
|
|
finalBytesDisplay = f"{finalBytes}B"
|
|
elif finalBytes < 1024 * 1024:
|
|
finalBytesDisplay = f"{finalBytes / 1024:.1f}kB"
|
|
else:
|
|
finalBytesDisplay = f"{finalBytes / (1024 * 1024):.1f}MB"
|
|
self.services.chat.progressLogUpdate(operationId, 0.95, f"Generation complete: {finalBytesDisplay} ({iteration} iterations, {len(allSections)} sections)")
|
|
logger.info(f"Generation complete after {iteration} iterations: {len(allSections)} sections")
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in AI call iteration {iteration}: {str(e)}")
|
|
if iterationOperationId:
|
|
self.services.chat.progressLogFinish(iterationOperationId, False)
|
|
break
|
|
|
|
if iteration >= maxIterations:
|
|
logger.warning(f"AI call stopped after maximum iterations ({maxIterations})")
|
|
|
|
# CRITICAL: Complete any incomplete structures in sections before building final result
|
|
# This ensures JSON is properly closed even if merge failed or iterations stopped early
|
|
allSections = JsonResponseHandler.completeIncompleteStructures(allSections)
|
|
|
|
# Build final result from accumulated sections
|
|
final_result = self.responseParser.buildFinalResultFromSections(allSections, documentMetadata)
|
|
|
|
# Write final result to debug file
|
|
# For section content generation: skip final_result debug file (response already written)
|
|
isSectionContent = "_section_" in debugPrefix
|
|
if not isSectionContent:
|
|
self.services.utils.writeDebugFile(final_result, f"{debugPrefix}_final_result")
|
|
|
|
return final_result
|
|
|
|
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
|
|
# This shouldn't happen, but we'll log a warning and return empty structure
|
|
logger.warning(f"Invalid response format: received list of strings instead of elements array. Expected {{'elements': [...]}} structure.")
|
|
return {"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
|
|
|
|
def _mergeJsonObjectsRecursively(self, jsonObjects: List[Any]) -> Any:
|
|
"""
|
|
GENERIC recursive merge function for JSON objects.
|
|
Works for ANY JSON structure - handles lists, dicts, and primitives intelligently.
|
|
|
|
Merge strategy:
|
|
- Lists/Arrays: Merge by removing duplicates based on content (works for rows, items, elements, etc.)
|
|
- Dicts/Objects: Merge properties, recursively merging nested structures
|
|
- Primitives: Use the latest value
|
|
|
|
Args:
|
|
jsonObjects: List of parsed JSON objects to merge
|
|
|
|
Returns:
|
|
Merged JSON object
|
|
"""
|
|
if not jsonObjects:
|
|
return None
|
|
if len(jsonObjects) == 1:
|
|
return jsonObjects[0]
|
|
|
|
# Start with first object and merge others into it
|
|
merged = jsonObjects[0]
|
|
|
|
for obj in jsonObjects[1:]:
|
|
merged = self._mergeTwoObjects(merged, obj)
|
|
|
|
return merged
|
|
|
|
def _mergeTwoObjects(self, obj1: Any, obj2: Any) -> Any:
|
|
"""
|
|
Merge two JSON objects recursively.
|
|
|
|
Args:
|
|
obj1: First object (base)
|
|
obj2: Second object (to merge into obj1)
|
|
|
|
Returns:
|
|
Merged object
|
|
"""
|
|
# Handle None values
|
|
if obj1 is None:
|
|
return obj2
|
|
if obj2 is None:
|
|
return obj1
|
|
|
|
# Handle different types
|
|
if isinstance(obj1, dict) and isinstance(obj2, dict):
|
|
# Merge dictionaries
|
|
merged = dict(obj1) # Start with copy of obj1
|
|
|
|
for key, value2 in obj2.items():
|
|
if key in merged:
|
|
# Key exists in both - recursively merge
|
|
merged[key] = self._mergeTwoObjects(merged[key], value2)
|
|
else:
|
|
# New key - add it
|
|
merged[key] = value2
|
|
|
|
return merged
|
|
|
|
elif isinstance(obj1, list) and isinstance(obj2, list):
|
|
# Merge lists by removing duplicates based on content
|
|
merged = list(obj1) # Start with copy of obj1
|
|
seenItems = set() # Track seen items to avoid duplicates
|
|
|
|
# Add all items from obj1 with their keys
|
|
for item in merged:
|
|
itemKey = self._createItemKey(item)
|
|
seenItems.add(itemKey)
|
|
|
|
# Add items from obj2 that aren't duplicates
|
|
for item in obj2:
|
|
itemKey = self._createItemKey(item)
|
|
if itemKey not in seenItems:
|
|
seenItems.add(itemKey)
|
|
merged.append(item)
|
|
|
|
return merged
|
|
|
|
else:
|
|
# Different types or primitives - use obj2 (latest value)
|
|
return obj2
|
|
|
|
def _createItemKey(self, item: Any) -> Any:
|
|
"""
|
|
Create a key for an item to detect duplicates.
|
|
Works generically for any JSON structure.
|
|
|
|
Args:
|
|
item: Item to create key for
|
|
|
|
Returns:
|
|
Key that can be used for duplicate detection
|
|
"""
|
|
if isinstance(item, dict):
|
|
# For dicts, create key from all values (or specific identifying fields)
|
|
# Try to find common identifying fields first
|
|
if "id" in item:
|
|
return ("id", item["id"])
|
|
elif "type" in item and "content" in item:
|
|
# For elements with type and content, use type + content hash
|
|
content = item.get("content", {})
|
|
if isinstance(content, dict):
|
|
# For tables/lists, use type + first few rows/items for key
|
|
if "rows" in content:
|
|
rows = content.get("rows", [])
|
|
return ("type", item["type"], "rows", tuple(rows[:3]) if rows else ())
|
|
elif "items" in content:
|
|
items = content.get("items", [])
|
|
return ("type", item["type"], "items", tuple(items[:3]) if items else ())
|
|
return ("type", item["type"], tuple(sorted(item.items())))
|
|
else:
|
|
# Generic: use sorted items tuple
|
|
return tuple(sorted(item.items()))
|
|
elif isinstance(item, (list, tuple)):
|
|
# For lists/tuples, use the tuple itself as key
|
|
return tuple(item) if isinstance(item, list) else item
|
|
else:
|
|
# For primitives, use the value itself
|
|
return item
|
|
|
|
async def _defineKpisFromPrompt(
|
|
self,
|
|
userPrompt: str,
|
|
rawJsonString: Optional[str],
|
|
continuationContext: Dict[str, Any],
|
|
debugPrefix: str = "kpi"
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Make separate AI call to define KPIs based on user prompt and incomplete JSON.
|
|
|
|
Args:
|
|
userPrompt: Original user prompt
|
|
rawJsonString: Raw JSON string from first iteration response
|
|
continuationContext: Continuation context (not used for JSON, kept for compatibility)
|
|
debugPrefix: Prefix for debug file names
|
|
|
|
Returns:
|
|
List of KPI definitions: [{"id": str, "description": str, "jsonPath": str, "targetValue": int}, ...]
|
|
"""
|
|
# Use raw JSON string from first iteration response
|
|
if rawJsonString:
|
|
# Remove markdown code fences if present
|
|
from modules.shared.jsonUtils import stripCodeFences
|
|
incompleteJson = stripCodeFences(rawJsonString.strip())
|
|
else:
|
|
incompleteJson = "Not available"
|
|
|
|
kpiDefinitionPrompt = f"""Analyze the user request and incomplete JSON to define KPIs (Key Performance Indicators) for tracking progress.
|
|
|
|
User Request:
|
|
{userPrompt}
|
|
|
|
Delivered JSON part:
|
|
{incompleteJson}
|
|
|
|
Task: Define which JSON items should be tracked to measure completion progress.
|
|
|
|
IMPORTANT: Analyze the Delivered JSON part structure to understand what is being tracked:
|
|
1. Identify the structure type (table with rows, list with items, etc.)
|
|
2. Determine what the jsonPath actually counts (number of rows, number of items, etc.)
|
|
3. Calculate targetValue based on what is being tracked, NOT the total quantity requested
|
|
|
|
For each trackable item, provide:
|
|
- id: Unique identifier (use descriptive name)
|
|
- description: What this KPI measures (be specific about what is counted)
|
|
- jsonPath: Path to extract value from JSON (use dot notation with array indices, e.g., "documents[0].sections[1].elements[0].rows")
|
|
- targetValue: Target value to reach (integer) - MUST match what jsonPath actually tracks (rows count, items count, etc.)
|
|
|
|
Return ONLY valid JSON in this format:
|
|
{{
|
|
"kpis": [
|
|
{{
|
|
"id": "unique_id",
|
|
"description": "Description of what is measured",
|
|
"jsonPath": "path.to.value",
|
|
"targetValue": 0
|
|
}}
|
|
]
|
|
}}
|
|
|
|
If no trackable items can be identified, return: {{"kpis": []}}
|
|
"""
|
|
|
|
try:
|
|
request = AiCallRequest(
|
|
prompt=kpiDefinitionPrompt,
|
|
options=AiCallOptions(
|
|
operationType=OperationTypeEnum.DATA_ANALYSE,
|
|
priority=PriorityEnum.SPEED,
|
|
processingMode=ProcessingModeEnum.BASIC
|
|
)
|
|
)
|
|
|
|
# Write KPI definition prompt to debug file
|
|
self.services.utils.writeDebugFile(kpiDefinitionPrompt, f"{debugPrefix}_kpi_definition_prompt")
|
|
|
|
checkWorkflowStopped(self.services)
|
|
response = await self.aiService.callAi(request)
|
|
|
|
# Write KPI definition response to debug file
|
|
self.services.utils.writeDebugFile(response.content, f"{debugPrefix}_kpi_definition_response")
|
|
|
|
# Parse response
|
|
extracted = extractJsonString(response.content)
|
|
kpiResponse = json.loads(extracted)
|
|
|
|
kpiDefinitions = kpiResponse.get("kpis", [])
|
|
logger.info(f"Defined {len(kpiDefinitions)} KPIs for tracking")
|
|
|
|
return kpiDefinitions
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to define KPIs: {e}, continuing without KPI tracking")
|
|
return []
|
|
|