Spaces:
Runtime error
Runtime error
| # Live PDF Generator Module | |
| # Enables real-time PDF updates as content changes | |
| import asyncio | |
| import io | |
| import tempfile | |
| import os | |
| from typing import Dict, Any, Optional, AsyncGenerator | |
| try: | |
| from fpdf2 import FPDF | |
| except ImportError: | |
| try: | |
| from fpdf import FPDF | |
| except ImportError: | |
| raise ImportError("Neither fpdf2 nor fpdf package found. Please install fpdf2.") | |
| import base64 | |
| import sys | |
| import os | |
| sys.path.append(os.path.dirname(__file__)) | |
| from pdf_generator import create_pdf_document | |
| from models import ExportRequest | |
| class LivePDFGenerator: | |
| """Generates and streams PDF updates in real-time""" | |
| def __init__(self): | |
| self.active_sessions: Dict[str, 'PDFSession'] = {} | |
| async def create_live_session(self, business_plan_id: str, initial_request: ExportRequest) -> str: | |
| """Create a new live PDF session""" | |
| session_id = f"pdf_session_{business_plan_id}_{int(asyncio.get_event_loop().time())}" | |
| session = PDFSession(session_id, business_plan_id, initial_request) | |
| self.active_sessions[session_id] = session | |
| # Generate initial PDF | |
| await session.update_pdf(initial_request) | |
| return session_id | |
| async def update_session(self, session_id: str, updated_request: ExportRequest) -> bool: | |
| """Update an existing PDF session with new content""" | |
| if session_id not in self.active_sessions: | |
| return False | |
| session = self.active_sessions[session_id] | |
| await session.update_pdf(updated_request) | |
| return True | |
| async def stream_pdf_updates(self, session_id: str) -> AsyncGenerator[bytes, None]: | |
| """Stream PDF updates as they occur""" | |
| if session_id not in self.active_sessions: | |
| return | |
| session = self.active_sessions[session_id] | |
| async for pdf_chunk in session.stream_updates(): | |
| yield pdf_chunk | |
| def close_session(self, session_id: str): | |
| """Close and cleanup a PDF session""" | |
| if session_id in self.active_sessions: | |
| del self.active_sessions[session_id] | |
| class PDFSession: | |
| """Individual PDF session that manages real-time updates""" | |
| def __init__(self, session_id: str, business_plan_id: str, initial_request: ExportRequest): | |
| self.session_id = session_id | |
| self.business_plan_id = business_plan_id | |
| self.current_request = initial_request | |
| self.update_queue = asyncio.Queue() | |
| self.is_active = True | |
| self.last_pdf_bytes: Optional[bytes] = None | |
| async def update_pdf(self, new_request: ExportRequest): | |
| """Update the PDF with new content""" | |
| self.current_request = new_request | |
| try: | |
| # Generate new PDF | |
| pdf_bytes, filename = create_pdf_document(new_request) | |
| self.last_pdf_bytes = pdf_bytes | |
| # Notify subscribers of update | |
| await self.update_queue.put({ | |
| 'type': 'pdf_update', | |
| 'data': base64.b64encode(pdf_bytes).decode('utf-8'), | |
| 'filename': filename, | |
| 'timestamp': asyncio.get_event_loop().time() | |
| }) | |
| except Exception as e: | |
| # Notify subscribers of error | |
| await self.update_queue.put({ | |
| 'type': 'pdf_error', | |
| 'error': str(e), | |
| 'timestamp': asyncio.get_event_loop().time() | |
| }) | |
| async def stream_updates(self) -> AsyncGenerator[bytes, None]: | |
| """Stream PDF updates as they occur""" | |
| while self.is_active: | |
| try: | |
| # Wait for updates with timeout | |
| update = await asyncio.wait_for(self.update_queue.get(), timeout=30.0) | |
| if update['type'] == 'pdf_update': | |
| # Send SSE event for PDF update | |
| event_data = f"event: pdf_update\ndata: {{\"data\": \"{update['data']}\", \"filename\": \"{update['filename']}\"}}\n\n" | |
| yield event_data.encode('utf-8') | |
| elif update['type'] == 'pdf_error': | |
| # Send SSE event for error | |
| error_data = f"event: pdf_error\ndata: {{\"error\": \"{update['error']}\"}}\n\n" | |
| yield error_data.encode('utf-8') | |
| except asyncio.TimeoutError: | |
| # Send keep-alive | |
| yield b": keep-alive\n\n" | |
| except Exception as e: | |
| # Send error event | |
| error_data = f"event: stream_error\ndata: {str(e)}\n\n" | |
| yield error_data.encode('utf-8') | |
| break | |
| def close(self): | |
| """Close this session""" | |
| self.is_active = False | |
| # Global instance | |
| live_pdf_generator = LivePDFGenerator() | |