| from datetime import datetime |
| from io import BytesIO |
| from pathlib import Path |
| from typing import Dict, Any, List |
| from html import escape |
|
|
| from markdown import markdown |
|
|
| import site |
| from fpdf import FPDF |
|
|
| from open_webui.env import STATIC_DIR, FONTS_DIR |
| from open_webui.models.chats import ChatTitleMessagesForm |
|
|
|
|
| class PDFGenerator: |
| """ |
| Description: |
| The `PDFGenerator` class is designed to create PDF documents from chat messages. |
| The process involves transforming markdown content into HTML and then into a PDF format |
| |
| Attributes: |
| - `form_data`: An instance of `ChatTitleMessagesForm` containing title and messages. |
| |
| """ |
|
|
| def __init__(self, form_data: ChatTitleMessagesForm): |
| self.html_body = None |
| self.messages_html = None |
| self.form_data = form_data |
|
|
| self.css = Path(STATIC_DIR / "assets" / "pdf-style.css").read_text() |
|
|
| def format_timestamp(self, timestamp: float) -> str: |
| """Convert a UNIX timestamp to a formatted date string.""" |
| try: |
| date_time = datetime.fromtimestamp(timestamp) |
| return date_time.strftime("%Y-%m-%d, %H:%M:%S") |
| except (ValueError, TypeError) as e: |
| |
| return "" |
|
|
| def _build_html_message(self, message: Dict[str, Any]) -> str: |
| """Build HTML for a single message.""" |
| role = escape(message.get("role", "user")) |
| content = escape(message.get("content", "")) |
| timestamp = message.get("timestamp") |
|
|
| model = escape(message.get("model") if role == "assistant" else "") |
|
|
| date_str = escape(self.format_timestamp(timestamp) if timestamp else "") |
|
|
| |
| |
| |
|
|
| content = content.replace("\n", "<br/>") |
| html_message = f""" |
| <div> |
| <div> |
| <h4> |
| <strong>{role.title()}</strong> |
| <span style="font-size: 12px;">{model}</span> |
| </h4> |
| <div> {date_str} </div> |
| </div> |
| <br/> |
| <br/> |
| |
| <div> |
| {content} |
| </div> |
| </div> |
| <br/> |
| """ |
| return html_message |
|
|
| def _generate_html_body(self) -> str: |
| """Generate the full HTML body for the PDF.""" |
| escaped_title = escape(self.form_data.title) |
| return f""" |
| <html> |
| <head> |
| <meta name="viewport" content="width=device-width, initial-scale=1.0" /> |
| </head> |
| <body> |
| <div> |
| <div> |
| <h2>{escaped_title}</h2> |
| {self.messages_html} |
| </div> |
| </div> |
| </body> |
| </html> |
| """ |
|
|
| def generate_chat_pdf(self) -> bytes: |
| """ |
| Generate a PDF from chat messages. |
| """ |
| try: |
| global FONTS_DIR |
|
|
| pdf = FPDF() |
| pdf.add_page() |
|
|
| |
| if not FONTS_DIR.exists(): |
| FONTS_DIR = Path(site.getsitepackages()[0]) / "static/fonts" |
| |
| |
| if not FONTS_DIR.exists(): |
| FONTS_DIR = Path(".") / "backend" / "static" / "fonts" |
|
|
| pdf.add_font("NotoSans", "", f"{FONTS_DIR}/NotoSans-Regular.ttf") |
| pdf.add_font("NotoSans", "b", f"{FONTS_DIR}/NotoSans-Bold.ttf") |
| pdf.add_font("NotoSans", "i", f"{FONTS_DIR}/NotoSans-Italic.ttf") |
| pdf.add_font("NotoSansKR", "", f"{FONTS_DIR}/NotoSansKR-Regular.ttf") |
| pdf.add_font("NotoSansJP", "", f"{FONTS_DIR}/NotoSansJP-Regular.ttf") |
| pdf.add_font("NotoSansSC", "", f"{FONTS_DIR}/NotoSansSC-Regular.ttf") |
| pdf.add_font("Twemoji", "", f"{FONTS_DIR}/Twemoji.ttf") |
|
|
| pdf.set_font("NotoSans", size=12) |
| pdf.set_fallback_fonts( |
| ["NotoSansKR", "NotoSansJP", "NotoSansSC", "Twemoji"] |
| ) |
|
|
| pdf.set_auto_page_break(auto=True, margin=15) |
|
|
| |
| messages_html_list: List[str] = [ |
| self._build_html_message(msg) for msg in self.form_data.messages |
| ] |
| self.messages_html = "<div>" + "".join(messages_html_list) + "</div>" |
|
|
| |
| self.html_body = self._generate_html_body() |
|
|
| pdf.write_html(self.html_body) |
|
|
| |
| pdf_bytes = pdf.output() |
|
|
| return bytes(pdf_bytes) |
| except Exception as e: |
| raise e |
|
|
