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')]