335 lines
15 KiB
Python
335 lines
15 KiB
Python
# Copyright (c) 2025 Patrick Motsch
|
|
# All rights reserved.
|
|
"""
|
|
Structure Generation Module
|
|
|
|
Handles document structure generation, including:
|
|
- Generating document structure with sections
|
|
- Building structure prompts
|
|
"""
|
|
import json
|
|
import logging
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
from modules.datamodels.datamodelExtraction import ContentPart
|
|
from modules.datamodels.datamodelAi import AiCallOptions, OperationTypeEnum, PriorityEnum, ProcessingModeEnum
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StructureGenerator:
|
|
"""Handles document structure generation."""
|
|
|
|
def __init__(self, services, aiService):
|
|
"""Initialize StructureGenerator with service center and AI service access."""
|
|
self.services = services
|
|
self.aiService = aiService
|
|
|
|
def _getUserLanguage(self) -> str:
|
|
"""Get user language for document generation"""
|
|
try:
|
|
if self.services:
|
|
# Prefer detected language if available (from user intention analysis)
|
|
if hasattr(self.services, 'currentUserLanguage') and self.services.currentUserLanguage:
|
|
return self.services.currentUserLanguage
|
|
# Fallback to user's preferred language
|
|
elif hasattr(self.services, 'user') and self.services.user and hasattr(self.services.user, 'language'):
|
|
return self.services.user.language
|
|
except Exception:
|
|
pass
|
|
return 'en' # Default fallback
|
|
|
|
async def generateStructure(
|
|
self,
|
|
userPrompt: str,
|
|
contentParts: List[ContentPart],
|
|
outputFormat: str,
|
|
parentOperationId: str
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Phase 5C: Generiert Chapter-Struktur (Table of Contents).
|
|
Definiert für jedes Chapter:
|
|
- Level, Title
|
|
- contentParts (unified object with instruction and/or caption per part)
|
|
- generationHint
|
|
|
|
Args:
|
|
userPrompt: User-Anfrage
|
|
contentParts: Alle vorbereiteten ContentParts mit Metadaten
|
|
outputFormat: Ziel-Format (html, docx, pdf, etc.)
|
|
parentOperationId: Parent Operation-ID für ChatLog-Hierarchie
|
|
|
|
Returns:
|
|
Struktur-Dict mit documents und chapters (nicht sections!)
|
|
"""
|
|
# Erstelle Operation-ID für Struktur-Generierung
|
|
structureOperationId = f"{parentOperationId}_structure_generation"
|
|
|
|
# Starte ChatLog mit Parent-Referenz
|
|
self.services.chat.progressLogStart(
|
|
structureOperationId,
|
|
"Chapter Structure Generation",
|
|
"Structure",
|
|
f"Generating chapter structure for {outputFormat}",
|
|
parentOperationId=parentOperationId
|
|
)
|
|
|
|
try:
|
|
# Baue Chapter-Struktur-Prompt mit Content-Index
|
|
structurePrompt = self._buildChapterStructurePrompt(
|
|
userPrompt=userPrompt,
|
|
contentParts=contentParts,
|
|
outputFormat=outputFormat
|
|
)
|
|
|
|
# AI-Call für Chapter-Struktur-Generierung mit Looping-Unterstützung
|
|
# Use _callAiWithLooping instead of callAiPlanning to support continuation if response is cut
|
|
options = AiCallOptions(
|
|
operationType=OperationTypeEnum.DATA_GENERATE,
|
|
priority=PriorityEnum.QUALITY,
|
|
processingMode=ProcessingModeEnum.DETAILED,
|
|
compressPrompt=False,
|
|
compressContext=False,
|
|
resultFormat="json"
|
|
)
|
|
|
|
# Create prompt builder for continuation support
|
|
async def buildChapterStructurePromptWithContinuation(
|
|
continuationContext: Optional[Dict[str, Any]] = None,
|
|
**kwargs
|
|
) -> str:
|
|
"""Build chapter structure prompt with optional continuation context."""
|
|
basePrompt = self._buildChapterStructurePrompt(
|
|
userPrompt=userPrompt,
|
|
contentParts=contentParts,
|
|
outputFormat=outputFormat
|
|
)
|
|
|
|
if continuationContext:
|
|
# Add continuation instructions
|
|
deliveredSummary = continuationContext.get("delivered_summary", "")
|
|
elementBeforeCutoff = continuationContext.get("element_before_cutoff", "")
|
|
cutOffElement = continuationContext.get("cut_off_element", "")
|
|
|
|
continuationText = f"{deliveredSummary}\n\n"
|
|
continuationText += "⚠️ CONTINUATION: Response was cut off. Generate ONLY the remaining content that comes AFTER the reference elements below.\n\n"
|
|
|
|
if elementBeforeCutoff:
|
|
continuationText += "# REFERENCE: Last complete element (already delivered - DO NOT repeat):\n"
|
|
continuationText += f"{elementBeforeCutoff}\n\n"
|
|
|
|
if cutOffElement:
|
|
continuationText += "# REFERENCE: Incomplete element (cut off here - DO NOT repeat):\n"
|
|
continuationText += f"{cutOffElement}\n\n"
|
|
|
|
continuationText += "⚠️ CRITICAL: The elements above are REFERENCE ONLY. They are already delivered.\n"
|
|
continuationText += "Generate ONLY what comes AFTER these elements. DO NOT regenerate the entire JSON structure.\n"
|
|
continuationText += "Start directly with the next chapter that should follow.\n\n"
|
|
|
|
return f"""{basePrompt}
|
|
|
|
{continuationText}
|
|
|
|
Continue generating the remaining chapters now.
|
|
"""
|
|
else:
|
|
return basePrompt
|
|
|
|
# Call AI with looping support
|
|
# NOTE: Do NOT pass contentParts here - we only need metadata for structure generation
|
|
# The contentParts metadata is already included in the prompt (contentPartsIndex)
|
|
# Actual content extraction happens later during section generation
|
|
aiResponseJson = await self.aiService.callAiWithLooping(
|
|
prompt=structurePrompt,
|
|
options=options,
|
|
debugPrefix="chapter_structure_generation",
|
|
promptBuilder=buildChapterStructurePromptWithContinuation,
|
|
promptArgs={
|
|
"userPrompt": userPrompt,
|
|
"outputFormat": outputFormat,
|
|
"services": self.services
|
|
},
|
|
useCaseId="chapter_structure", # REQUIRED: Explicit use case ID
|
|
operationId=structureOperationId,
|
|
userPrompt=userPrompt,
|
|
contentParts=None # Do not pass ContentParts - only metadata needed, not content extraction
|
|
)
|
|
|
|
# Parse the complete JSON response (looping system already handles completion)
|
|
extractedJson = self.services.utils.jsonExtractString(aiResponseJson)
|
|
parsedJson, parseError, cleanedJson = self.services.utils.jsonTryParse(extractedJson)
|
|
|
|
if parseError is not None:
|
|
# Even with looping, try repair as fallback
|
|
logger.warning(f"JSON parsing failed after looping: {str(parseError)}. Attempting repair...")
|
|
from modules.shared import jsonUtils
|
|
repairedJson = jsonUtils.repairBrokenJson(extractedJson)
|
|
if repairedJson:
|
|
parsedJson, parseError, _ = self.services.utils.jsonTryParse(json.dumps(repairedJson))
|
|
if parseError is None:
|
|
logger.info("Successfully repaired and parsed JSON structure after looping")
|
|
structure = parsedJson
|
|
else:
|
|
logger.error(f"Failed to parse repaired JSON: {str(parseError)}")
|
|
raise ValueError(f"Failed to parse JSON structure after repair: {str(parseError)}")
|
|
else:
|
|
logger.error(f"Failed to repair JSON. Parse error: {str(parseError)}")
|
|
logger.error(f"Cleaned JSON preview (first 500 chars): {cleanedJson[:500]}")
|
|
raise ValueError(f"Failed to parse JSON structure: {str(parseError)}")
|
|
else:
|
|
structure = parsedJson
|
|
|
|
# ChatLog abschließen
|
|
self.services.chat.progressLogFinish(structureOperationId, True)
|
|
|
|
return structure
|
|
|
|
except Exception as e:
|
|
self.services.chat.progressLogFinish(structureOperationId, False)
|
|
logger.error(f"Error in generateStructure: {str(e)}")
|
|
raise
|
|
|
|
def _buildChapterStructurePrompt(
|
|
self,
|
|
userPrompt: str,
|
|
contentParts: List[ContentPart],
|
|
outputFormat: str
|
|
) -> str:
|
|
"""Baue Prompt für Chapter-Struktur-Generierung."""
|
|
# Baue ContentParts-Index - filtere leere Parts heraus
|
|
contentPartsIndex = ""
|
|
validParts = []
|
|
filteredParts = []
|
|
|
|
for part in contentParts:
|
|
contentFormat = part.metadata.get("contentFormat", "unknown")
|
|
|
|
# WICHTIG: Reference Parts haben absichtlich leere Daten - immer einschließen
|
|
if contentFormat == "reference":
|
|
validParts.append(part)
|
|
logger.debug(f"Including reference ContentPart {part.id} (intentionally empty data)")
|
|
continue
|
|
|
|
# Überspringe leere Parts (keine Daten oder nur Container ohne Inhalt)
|
|
# ABER: Reference Parts wurden bereits oben behandelt
|
|
if not part.data or (isinstance(part.data, str) and len(part.data.strip()) == 0):
|
|
# Überspringe Container-Parts ohne Daten
|
|
if part.typeGroup == "container" and not part.data:
|
|
filteredParts.append((part.id, "container without data"))
|
|
continue
|
|
# Überspringe andere leere Parts (aber nicht Reference, die wurden bereits behandelt)
|
|
if not part.data:
|
|
filteredParts.append((part.id, f"no data (format: {contentFormat})"))
|
|
continue
|
|
|
|
validParts.append(part)
|
|
logger.debug(f"Including ContentPart {part.id}: format={contentFormat}, type={part.typeGroup}, dataLength={len(str(part.data)) if part.data else 0}")
|
|
|
|
if filteredParts:
|
|
logger.debug(f"Filtered out {len(filteredParts)} empty ContentParts: {filteredParts}")
|
|
|
|
logger.info(f"Building structure prompt with {len(validParts)} valid ContentParts (from {len(contentParts)} total)")
|
|
|
|
# Baue Index nur für gültige Parts
|
|
for i, part in enumerate(validParts, 1):
|
|
contentFormat = part.metadata.get("contentFormat", "unknown")
|
|
originalFileName = part.metadata.get('originalFileName', 'N/A')
|
|
|
|
contentPartsIndex += f"\n{i}. ContentPart ID: {part.id}\n"
|
|
contentPartsIndex += f" Format: {contentFormat}\n"
|
|
contentPartsIndex += f" Type: {part.typeGroup}\n"
|
|
contentPartsIndex += f" MIME Type: {part.mimeType or 'N/A'}\n"
|
|
contentPartsIndex += f" Source: {part.metadata.get('documentId', 'unknown')}\n"
|
|
contentPartsIndex += f" Original file name: {originalFileName}\n"
|
|
contentPartsIndex += f" Usage hint: {part.metadata.get('usageHint', 'N/A')}\n"
|
|
|
|
if not contentPartsIndex:
|
|
contentPartsIndex = "\n(No content parts available)"
|
|
|
|
# Get language from services (user intention analysis)
|
|
language = self._getUserLanguage()
|
|
logger.debug(f"Using language from services (user intention analysis) for structure generation: {language}")
|
|
|
|
prompt = f"""CRITICAL OUTPUT REQUIREMENT: This is a PLANNING task, not a generation task. You MUST return EXACTLY ONE complete JSON object. Do NOT generate multiple JSON objects, alternatives, or variations. Do NOT use separators like "---" between JSON objects. Return the single best JSON structure that matches the requirements below.
|
|
|
|
USER REQUEST (for context):
|
|
```
|
|
{userPrompt}
|
|
```
|
|
|
|
LANGUAGE: Generate all content in {language.upper()} language. All text, titles, headings, paragraphs, and content must be written in {language.upper()}.
|
|
|
|
AVAILABLE CONTENT PARTS:
|
|
{contentPartsIndex}
|
|
|
|
TASK: Generate Chapter Structure for the documents to be generated.
|
|
|
|
IMPORTANT - CHAPTER INDEPENDENCE:
|
|
- Each chapter is independent and self-contained
|
|
- One chapter does NOT have information about another chapter
|
|
- Each chapter must provide its own context and be understandable alone
|
|
|
|
CONTENT ASSIGNMENT:
|
|
- Assign ContentParts to chapters via contentParts object
|
|
- For data extraction, the type of a contentPart (image, text, etc.) is NOT relevant - only what is specified in the instruction matters
|
|
- Include ALL relevant parts from same source when needed for structured data extraction
|
|
- Each contentPart can have either:
|
|
- "instruction": For AI extraction prompts (how to process/extract from this part)
|
|
- "caption": For user-facing presentation (how to display/reference this part in the document)
|
|
- Both can be present if needed
|
|
- Chapters without contentParts can only generate generic content (not document-specific)
|
|
|
|
FORMATTING:
|
|
- Formatting is handled automatically - focus on CONTENT and STRUCTURE only
|
|
|
|
CHAPTER STRUCTURE:
|
|
- chapter id, level (1, 2, 3, etc.), title
|
|
- contentParts: {{"partId": {{"instruction": "..."}} or {{"caption": "..."}} or both}} - Compact mapping of part IDs to their extraction instructions and/or presentation captions
|
|
- generationHint: Self-contained description (if contentParts is empty, must be VERY DETAILED)
|
|
|
|
RETURN JSON:
|
|
{{
|
|
"metadata": {{
|
|
"title": "Document Title",
|
|
"language": "{language}"
|
|
}},
|
|
"documents": [{{
|
|
"id": "doc_1",
|
|
"title": "Document Title",
|
|
"filename": "document.{outputFormat}",
|
|
"chapters": [
|
|
{{
|
|
"id": "chapter_1",
|
|
"level": 1,
|
|
"title": "Introduction",
|
|
"contentParts": {{
|
|
"part_ext_1": {{
|
|
"instruction": "Use full extracted text"
|
|
}},
|
|
"part_img_1": {{
|
|
"instruction": "Analyze image for additional details"
|
|
}},
|
|
"part_img_2": {{
|
|
"instruction": "Analyze image for additional details",
|
|
"caption": "Figure 1: Overview diagram"
|
|
}}
|
|
}},
|
|
"generationHint": "Create introduction section",
|
|
"sections": []
|
|
}},
|
|
{{
|
|
"id": "chapter_2",
|
|
"level": 1,
|
|
"title": "Main Title",
|
|
"contentParts": {{}},
|
|
"generationHint": "Create [specific content description] with [formatting details]. Include [required information]. Purpose: [explanation of what this chapter provides].",
|
|
"sections": []
|
|
}}
|
|
]
|
|
}}]
|
|
}}
|
|
|
|
OUTPUT FORMAT: Start with {{ and end with }}. Do NOT use markdown code fences (```json). Do NOT add explanatory text before or after the JSON. Return ONLY the JSON object itself.
|
|
"""
|
|
return prompt
|
|
|