Spaces:

AccuratePaperExtraction
/

MADTOOL

Configuration error

MADTOOL / system_chart_service.py

Initial commit - chatbot online version

ceaf3bb 8 months ago

13.3 kB

import tempfile
import os
import shutil
from io import BytesIO
from zipfile import ZipFile
from typing import Dict, Callable, Any, List
from datetime import datetime
from system_chart_extraction import system_chart_extraction

class SystemChartService:
    """
    Service class for System Chart Extraction.
    Handles the orchestration of system chart extraction, visualization creation,
    and result packaging.
    """
    
    def __init__(self):
        pass
    
    async def extract_system_charts(
        self,
        category: str,
        llm,
        pdf_dict: Dict[str, str],
        enable_visualization: bool = False,
        run_system_chart_normalization: bool = False,
        progress_callback: Callable[[Dict[str, Any]], None] = None,
        extraction_id: str = None
    ) -> Dict[str, Any]:
        """
        Main system chart extraction pipeline.
        
        Args:
            category: Database category for extraction
            llm: Language model for extraction
            pdf_dict: Dictionary mapping document titles to full text content
            enable_visualization: Whether to create visualization diagrams
            progress_callback: Function to report progress updates
            
        Returns:
            Dict containing extraction results and metadata
        """
        try:
            # Validate input
            if not pdf_dict:
                return {
                    "success": False,
                    "message": "No files found in the specified category.",
                    "zip_data": None,
                    "filename": None
                }
            
            print(f"📄 Found {len(pdf_dict)} papers to process")
            
            # Initial progress update
            if progress_callback:
                progress_callback({
                    "status": "starting",
                    "progress": 0,
                    "message": "Getting PDF files..."
                })
            
            # Create progress callback function for the extraction process
            def extraction_progress_callback(progress_data):
                if progress_callback:
                    print(f"📊 Progress update: {progress_data}")
                    progress_callback(progress_data)
                    print(f"🔄 Progress updated: {progress_data}")
            
            # Update progress before starting extraction
            if progress_callback:
                progress_callback({
                    "status": "starting",
                    "progress": 5,
                    "message": f"Found {len(pdf_dict)} papers. Starting extraction..."
                })
            
            print(f"🔬 Starting system chart extraction...")
            
            # Create temporary directory for visualization files if needed
            temp_dir = None
            if enable_visualization:
                temp_dir = tempfile.mkdtemp(prefix="system_chart_viz_")
                print(f"📁 Created temporary directory: {temp_dir}")
            
            # Run system chart extraction with progress tracking
            output_system_charts, output_all_details, created_visualizations = await system_chart_extraction(
                pdf_dict, 
                llm, 
                extraction_progress_callback, 
                enable_visualization,
                temp_dir,
                extraction_id,  # Pass extraction_id to save conversations
                run_normalization=run_system_chart_normalization  # ← only normalize when asked/needed
            )
            
            # Create ZIP file in memory
            zip_data = self._create_zip_file(
                output_system_charts, 
                output_all_details, 
                created_visualizations
            )
            
            # Clean up temporary directory
            if temp_dir:
                self._cleanup_temp_directory(temp_dir)
            
            # Prepare result
            viz_message = f" with {len(created_visualizations)} visualizations" if created_visualizations else ""
            
            return {
                "success": True,
                "message": f"System chart extraction completed successfully{viz_message}!",
                "zip_data": zip_data,
                "filename": "system_chart_extraction.zip",
                "charts_count": len(pdf_dict),
                "visualizations_count": len(created_visualizations) if created_visualizations else 0
            }
            
        except Exception as e:
            # Clean up on error
            if 'temp_dir' in locals() and temp_dir:
                self._cleanup_temp_directory(temp_dir)
            
            return {
                "success": False,
                "message": f"Error during system chart extraction: {str(e)}",
                "zip_data": None,
                "filename": None
            }
    
    def _create_zip_file(
        self, 
        output_system_charts: str, 
        output_all_details: str, 
        created_visualizations: List[str]
    ) -> bytes:
        """Create ZIP file containing extraction results and visualizations."""
        zip_buffer = BytesIO()
        
        with ZipFile(zip_buffer, "w") as zipf:
            # Add text results
            zipf.writestr("system_charts.txt", output_system_charts)
            zipf.writestr("all_details.txt", output_all_details)
            
            # Add visualization files if any were created
            if created_visualizations:
                for viz_path in created_visualizations:
                    if os.path.exists(viz_path):
                        # Extract just the filename for the ZIP entry
                        viz_filename = os.path.basename(viz_path)
                        zipf.write(viz_path, f"visualizations/{viz_filename}")
                        print(f"📦 Added visualization to ZIP: {viz_filename}")
        
        # Get ZIP data and close buffer
        zip_data = zip_buffer.getvalue()
        zip_buffer.close()
        
        return zip_data
    
    def _cleanup_temp_directory(self, temp_dir: str) -> None:
        """Clean up temporary directory and its contents."""
        if temp_dir and os.path.exists(temp_dir):
            try:
                shutil.rmtree(temp_dir)
                print(f"🗑️ Cleaned up temporary directory: {temp_dir}")
            except OSError as e:
                print(f"⚠️ Could not clean up temporary directory {temp_dir}: {e}")
                # Don't fail the whole process if cleanup fails
                pass
    
    @staticmethod
    def parse_system_charts_file(file_content: str) -> List[Dict[str, str]]:
        """
        Parse a delimited system charts file into individual chart objects.
        
        Args:
            file_content: Content of the system_charts.txt file
            
        Returns:
            List of dictionaries, each containing:
            - 'chart_number': The chart number (auto-generated)
            - 'title': Paper title
            - 'citation': Citation format (same as title for compatibility)
            - 'content': System chart table content
        """
        charts = []
        lines = file_content.split('\n')
        
        current_chart = None
        chart_number = 1
        
        for i, line in enumerate(lines):
            line = line.strip()
            
            # Check for delimiter line (100 '=' characters)
            if line == '=' * 100:
                # Look for PAPER: line after delimiter
                if i + 1 < len(lines):
                    next_line = lines[i + 1].strip()
                    if next_line.startswith('PAPER:'):
                        # Save previous chart if exists
                        if current_chart:
                            charts.append(current_chart)
                        
                        # Start new chart
                        paper_title = next_line.replace('PAPER:', '').strip()
                        current_chart = {
                            'chart_number': str(chart_number),
                            'title': paper_title,
                            'citation': f'[{paper_title}]',
                            'content': ''
                        }
                        chart_number += 1
                        
            elif line.startswith('SYSTEM_CHART:') and current_chart:
                # Extract the system chart content
                chart_content = line.replace('SYSTEM_CHART:', '').strip()
                current_chart['content'] = chart_content
        
        # Handle last chart if file doesn't end with proper delimiter
        if current_chart:
            charts.append(current_chart)
        
        print(f"📊 Parsed {len(charts)} system charts from file")
        return charts
    
    @staticmethod
    def combine_system_charts_files(file_contents: List[str]) -> str:
        """
        Combine multiple system chart txt files into a single formatted string.
        Preserves the one-to-one relationship between system charts and titles.
        
        Args:
            file_contents: List of system chart file contents to combine
            
        Returns:
            Combined system chart content in the new format
        """
        combined_charts = []
        
        # Parse each file and collect all charts
        for file_content in file_contents:
            charts = SystemChartService.parse_system_charts_file(file_content)
            combined_charts.extend(charts)
        
        # Generate combined output in new format
        combined_output = "COMBINED SYSTEM CHARTS EXTRACTION RESULTS\n"
        combined_output += "=" * 100 + "\n"
        
        for chart in combined_charts:
            delimiter = "=" * 100
            combined_output += f'\n{delimiter}\n'
            combined_output += f'PAPER: {chart["title"]}\n'
            combined_output += f'SYSTEM_CHART: {chart["content"]}\n'
            combined_output += f'{delimiter}\n\n'
        
        # Add summary
        summary_delimiter = "=" * 100
        combined_output += f'\n{summary_delimiter}\n'
        combined_output += f'COMBINATION SUMMARY\n'
        combined_output += f'{summary_delimiter}\n'
        combined_output += f'Total Files Combined: {len(file_contents)}\n'
        combined_output += f'Total System Charts: {len(combined_charts)}\n'
        combined_output += f'Combination Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n'
        
        print(f"🔗 Combined {len(file_contents)} files into {len(combined_charts)} system charts")
        return combined_output
    
    async def combine_multiple_system_chart_files(
        self,
        file_paths: List[str],
        output_filename: str = "combined_system_charts.txt"
    ) -> Dict[str, Any]:
        """
        Combine multiple system chart txt files into a single ZIP file.
        
        Args:
            file_paths: List of file paths to system chart txt files
            output_filename: Name for the output file
            
        Returns:
            Dict containing success status, message, and ZIP data
        """
        try:
            file_contents = []
            
            # Read all files
            for file_path in file_paths:
                if not os.path.exists(file_path):
                    return {
                        "success": False,
                        "message": f"File not found: {file_path}",
                        "zip_data": None,
                        "filename": None
                    }
                
                with open(file_path, 'r', encoding='utf-8') as f:
                    file_contents.append(f.read())
            
            # Combine the files
            combined_content = self.combine_system_charts_files(file_contents)
            
            # Create ZIP file
            zip_buffer = BytesIO()
            with ZipFile(zip_buffer, "w") as zipf:
                zipf.writestr(output_filename, combined_content)
            
            zip_data = zip_buffer.getvalue()
            zip_buffer.close()
            
            return {
                "success": True,
                "message": f"Successfully combined {len(file_paths)} system chart files!",
                "zip_data": zip_data,
                "filename": "combined_system_charts.zip",
                "files_combined": len(file_paths)
            }
            
        except Exception as e:
            return {
                "success": False,
                "message": f"Error combining system chart files: {str(e)}",
                "zip_data": None,
                "filename": None
            }
    
    @staticmethod
    def get_chart_citations(file_content: str) -> List[str]:
        """
        Extract just the citations from a system charts file.
        
        Args:
            file_content: Content of the system_charts.txt file
            
        Returns:
            List of citation strings
        """
        charts = SystemChartService.parse_system_charts_file(file_content)
        return [chart.get('citation', chart.get('title', f"Chart {chart['chart_number']}")) 
                for chart in charts if chart.get('chart_number')] 

	import tempfile
	import os
	import shutil
	from io import BytesIO
	from zipfile import ZipFile
	from typing import Dict, Callable, Any, List
	from datetime import datetime
	from system_chart_extraction import system_chart_extraction

	class SystemChartService:
	"""
	Service class for System Chart Extraction.
	Handles the orchestration of system chart extraction, visualization creation,
	and result packaging.
	"""

	def __init__(self):
	pass

	async def extract_system_charts(
	self,
	category: str,
	llm,
	pdf_dict: Dict[str, str],
	enable_visualization: bool = False,
	run_system_chart_normalization: bool = False,
	progress_callback: Callable[[Dict[str, Any]], None] = None,
	extraction_id: str = None
	) -> Dict[str, Any]:
	"""
	Main system chart extraction pipeline.

	Args:
	category: Database category for extraction
	llm: Language model for extraction
	pdf_dict: Dictionary mapping document titles to full text content
	enable_visualization: Whether to create visualization diagrams
	progress_callback: Function to report progress updates

	Returns:
	Dict containing extraction results and metadata
	"""
	try:
	# Validate input
	if not pdf_dict:
	return {
	"success": False,
	"message": "No files found in the specified category.",
	"zip_data": None,
	"filename": None
	}

	print(f"📄 Found {len(pdf_dict)} papers to process")

	# Initial progress update
	if progress_callback:
	progress_callback({
	"status": "starting",
	"progress": 0,
	"message": "Getting PDF files..."
	})

	# Create progress callback function for the extraction process
	def extraction_progress_callback(progress_data):
	if progress_callback:
	print(f"📊 Progress update: {progress_data}")
	progress_callback(progress_data)
	print(f"🔄 Progress updated: {progress_data}")

	# Update progress before starting extraction
	if progress_callback:
	progress_callback({
	"status": "starting",
	"progress": 5,
	"message": f"Found {len(pdf_dict)} papers. Starting extraction..."
	})

	print(f"🔬 Starting system chart extraction...")

	# Create temporary directory for visualization files if needed
	temp_dir = None
	if enable_visualization:
	temp_dir = tempfile.mkdtemp(prefix="system_chart_viz_")
	print(f"📁 Created temporary directory: {temp_dir}")

	# Run system chart extraction with progress tracking
	output_system_charts, output_all_details, created_visualizations = await system_chart_extraction(
	pdf_dict,
	llm,
	extraction_progress_callback,
	enable_visualization,
	temp_dir,
	extraction_id, # Pass extraction_id to save conversations
	run_normalization=run_system_chart_normalization # ← only normalize when asked/needed
	)

	# Create ZIP file in memory
	zip_data = self._create_zip_file(
	output_system_charts,
	output_all_details,
	created_visualizations
	)

	# Clean up temporary directory
	if temp_dir:
	self._cleanup_temp_directory(temp_dir)

	# Prepare result
	viz_message = f" with {len(created_visualizations)} visualizations" if created_visualizations else ""

	return {
	"success": True,
	"message": f"System chart extraction completed successfully{viz_message}!",
	"zip_data": zip_data,
	"filename": "system_chart_extraction.zip",
	"charts_count": len(pdf_dict),
	"visualizations_count": len(created_visualizations) if created_visualizations else 0
	}

	except Exception as e:
	# Clean up on error
	if 'temp_dir' in locals() and temp_dir:
	self._cleanup_temp_directory(temp_dir)

	return {
	"success": False,
	"message": f"Error during system chart extraction: {str(e)}",
	"zip_data": None,
	"filename": None
	}

	def _create_zip_file(
	self,
	output_system_charts: str,
	output_all_details: str,
	created_visualizations: List[str]
	) -> bytes:
	"""Create ZIP file containing extraction results and visualizations."""
	zip_buffer = BytesIO()

	with ZipFile(zip_buffer, "w") as zipf:
	# Add text results
	zipf.writestr("system_charts.txt", output_system_charts)
	zipf.writestr("all_details.txt", output_all_details)

	# Add visualization files if any were created
	if created_visualizations:
	for viz_path in created_visualizations:
	if os.path.exists(viz_path):
	# Extract just the filename for the ZIP entry
	viz_filename = os.path.basename(viz_path)
	zipf.write(viz_path, f"visualizations/{viz_filename}")
	print(f"📦 Added visualization to ZIP: {viz_filename}")

	# Get ZIP data and close buffer
	zip_data = zip_buffer.getvalue()
	zip_buffer.close()

	return zip_data

	def _cleanup_temp_directory(self, temp_dir: str) -> None:
	"""Clean up temporary directory and its contents."""
	if temp_dir and os.path.exists(temp_dir):
	try:
	shutil.rmtree(temp_dir)
	print(f"🗑️ Cleaned up temporary directory: {temp_dir}")
	except OSError as e:
	print(f"⚠️ Could not clean up temporary directory {temp_dir}: {e}")
	# Don't fail the whole process if cleanup fails
	pass

	@staticmethod
	def parse_system_charts_file(file_content: str) -> List[Dict[str, str]]:
	"""
	Parse a delimited system charts file into individual chart objects.

	Args:
	file_content: Content of the system_charts.txt file

	Returns:
	List of dictionaries, each containing:
	- 'chart_number': The chart number (auto-generated)
	- 'title': Paper title
	- 'citation': Citation format (same as title for compatibility)
	- 'content': System chart table content
	"""
	charts = []
	lines = file_content.split('\n')

	current_chart = None
	chart_number = 1

	for i, line in enumerate(lines):
	line = line.strip()

	# Check for delimiter line (100 '=' characters)
	if line == '=' * 100:
	# Look for PAPER: line after delimiter
	if i + 1 < len(lines):
	next_line = lines[i + 1].strip()
	if next_line.startswith('PAPER:'):
	# Save previous chart if exists
	if current_chart:
	charts.append(current_chart)

	# Start new chart
	paper_title = next_line.replace('PAPER:', '').strip()
	current_chart = {
	'chart_number': str(chart_number),
	'title': paper_title,
	'citation': f'[{paper_title}]',
	'content': ''
	}
	chart_number += 1

	elif line.startswith('SYSTEM_CHART:') and current_chart:
	# Extract the system chart content
	chart_content = line.replace('SYSTEM_CHART:', '').strip()
	current_chart['content'] = chart_content

	# Handle last chart if file doesn't end with proper delimiter
	if current_chart:
	charts.append(current_chart)

	print(f"📊 Parsed {len(charts)} system charts from file")
	return charts

	@staticmethod
	def combine_system_charts_files(file_contents: List[str]) -> str:
	"""
	Combine multiple system chart txt files into a single formatted string.
	Preserves the one-to-one relationship between system charts and titles.

	Args:
	file_contents: List of system chart file contents to combine

	Returns:
	Combined system chart content in the new format
	"""
	combined_charts = []

	# Parse each file and collect all charts
	for file_content in file_contents:
	charts = SystemChartService.parse_system_charts_file(file_content)
	combined_charts.extend(charts)

	# Generate combined output in new format
	combined_output = "COMBINED SYSTEM CHARTS EXTRACTION RESULTS\n"
	combined_output += "=" * 100 + "\n"

	for chart in combined_charts:
	delimiter = "=" * 100
	combined_output += f'\n{delimiter}\n'
	combined_output += f'PAPER: {chart["title"]}\n'
	combined_output += f'SYSTEM_CHART: {chart["content"]}\n'
	combined_output += f'{delimiter}\n\n'

	# Add summary
	summary_delimiter = "=" * 100
	combined_output += f'\n{summary_delimiter}\n'
	combined_output += f'COMBINATION SUMMARY\n'
	combined_output += f'{summary_delimiter}\n'
	combined_output += f'Total Files Combined: {len(file_contents)}\n'
	combined_output += f'Total System Charts: {len(combined_charts)}\n'
	combined_output += f'Combination Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n'

	print(f"🔗 Combined {len(file_contents)} files into {len(combined_charts)} system charts")
	return combined_output

	async def combine_multiple_system_chart_files(
	self,
	file_paths: List[str],
	output_filename: str = "combined_system_charts.txt"
	) -> Dict[str, Any]:
	"""
	Combine multiple system chart txt files into a single ZIP file.

	Args:
	file_paths: List of file paths to system chart txt files
	output_filename: Name for the output file

	Returns:
	Dict containing success status, message, and ZIP data
	"""
	try:
	file_contents = []

	# Read all files
	for file_path in file_paths:
	if not os.path.exists(file_path):
	return {
	"success": False,
	"message": f"File not found: {file_path}",
	"zip_data": None,
	"filename": None
	}

	with open(file_path, 'r', encoding='utf-8') as f:
	file_contents.append(f.read())

	# Combine the files
	combined_content = self.combine_system_charts_files(file_contents)

	# Create ZIP file
	zip_buffer = BytesIO()
	with ZipFile(zip_buffer, "w") as zipf:
	zipf.writestr(output_filename, combined_content)

	zip_data = zip_buffer.getvalue()
	zip_buffer.close()

	return {
	"success": True,
	"message": f"Successfully combined {len(file_paths)} system chart files!",
	"zip_data": zip_data,
	"filename": "combined_system_charts.zip",
	"files_combined": len(file_paths)
	}

	except Exception as e:
	return {
	"success": False,
	"message": f"Error combining system chart files: {str(e)}",
	"zip_data": None,
	"filename": None
	}

	@staticmethod
	def get_chart_citations(file_content: str) -> List[str]:
	"""
	Extract just the citations from a system charts file.

	Args:
	file_content: Content of the system_charts.txt file

	Returns:
	List of citation strings
	"""
	charts = SystemChartService.parse_system_charts_file(file_content)
	return [chart.get('citation', chart.get('title', f"Chart {chart['chart_number']}"))
	for chart in charts if chart.get('chart_number')]