336 lines
15 KiB
Python
336 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"""# TASK: Generate Chapter Structure
|
|
|
|
This is a PLANNING task. Return EXACTLY ONE complete JSON object. Do not generate multiple JSON objects, alternatives, or variations. Do not use separators like "---" between JSON objects.
|
|
|
|
## 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}
|
|
|
|
## 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 that reflects the user's intent for the specific data. If contentParts is empty, must be detailed. If contentParts are present, the hint should guide how to extract and structure the data according to the user's requirements (e.g., specific columns, format, structure)
|
|
|
|
## OUTPUT FORMAT
|
|
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
|
|
|