Add to_docx to report and tweak report

lumen/ai/export.py

@@ -1,22 +1,32 @@
 import base64
+import io
 import os
 import re
+import tempfile
+import warnings
 
+from datetime import datetime
 from io import BytesIO
+from pathlib import Path
 from textwrap import dedent
 from typing import Any
 
 import nbformat
 import yaml
 
+from docx.shared import Mm
+from docxtpl import (
+    DocxTemplate, InlineImage, R, RichText,
+)
 from panel import Column
 from panel.chat import ChatMessage, ChatStep
 from panel.pane.image import ImageBase
+from panel_material_ui import Typography
 
 from ..config import config
 from ..pipeline import Pipeline
 from ..views import View
-from .views import LumenOutput
+from .views import LumenOutput, VegaLiteOutput
 
 
 def make_md_cell(text: str):
@@ -135,3 +145,182 @@ def export_notebook(messages: list[ChatMessage], preamble: str = ""):
     cells, extensions = render_cells(messages)
     cells = make_preamble(preamble, extensions=extensions) + cells
     return write_notebook(cells)
+
+
+def render_docx_template(
+    sections: list,
+    docx_template_path: str | Path,
+    **docx_context: dict,
+) -> io.BytesIO:
+    """
+    Export outputs to a Word document (.docx) format.
+
+    Arguments
+    ---------
+    sections: list
+        List of sections to process
+    docx_context : dict | None
+        Context dictionary for docx template rendering. If keys are not provided,
+        the following defaults will be used:
+
+        - 'title': title parameter or 'Lumen Report'
+        - 'subtitle': 'Generated on {date}' (e.g., 'Generated on October 27, 2025')
+        - 'cover_page_header': '' (empty string)
+        - 'cover_page_footer': '' (empty string)
+        - 'content_page_header': '' (empty string)
+
+        The following keys are always auto-generated and cannot be overridden:
+        - 'sections': List of section dicts with 'title', 'image', and 'caption'
+        - 'page_break': R('\f') for page breaks
+    docx_template_path : str | None
+        Path to the docx template file. If None, uses the default Lumen template.
+
+    Returns
+    -------
+    BytesIO
+        A BytesIO buffer containing the rendered docx document.
+
+    Raises
+    ------
+    RuntimeError
+        If the outputs list is empty.
+    FileNotFoundError
+        If the template file is not found.
+
+    Example
+    -------
+        buffer = to_docx(
+            outputs=report.outputs,
+            **{
+                'subtitle': 'Quarterly Analysis',
+                'cover_page_header': 'ACME Corporation'
+            }
+        )
+    """
+    # Set default template path
+    if docx_template_path is None:
+        docx_template_path = str(Path(__file__).parent / "assets" / "lumen_template.docx")
+
+    # Load template
+    template_path = Path(docx_template_path)
+    if not template_path.exists():
+        raise FileNotFoundError(f"Template file not found: {template_path}")
+
+    doc = DocxTemplate(str(template_path))
+
+    # Start with copy of docx_context or empty dict
+    context = dict(docx_context) if docx_context else {}
+
+    # Set defaults for missing keys
+    if 'title' not in context:
+        context['title'] = "Lumen Report"
+
+    if 'subtitle' not in context:
+        date_string = datetime.now().strftime("%B %d, %Y")
+        context['subtitle'] = f"Generated on {date_string}"
+
+    if 'cover_page_header' not in context:
+        context['cover_page_header'] = ""
+
+    if 'cover_page_footer' not in context:
+        context['cover_page_footer'] = ""
+
+    if 'content_page_header' not in context:
+        context['content_page_header'] = ""
+
+    # Always generate sections
+    context['sections'] = generate_docx_sections(doc, sections)
+
+    # Always set page_break
+    context['page_break'] = R("\f")
+
+    # Render template
+    doc.render(context)
+
+    # Return as BytesIO
+    buffer = io.BytesIO()
+    doc.save(buffer)
+    buffer.seek(0)
+    return buffer
+
+
+def generate_docx_sections(doc: DocxTemplate, sections: list) -> list[dict]:
+    """
+    Generate sections list from report tasks for docx template.
+
+    Arguments
+    ---------
+    doc : DocxTemplate
+        The document template instance (needed for InlineImage creation)
+    sections : list
+        List of sections to process
+
+    Returns
+    -------
+    list[dict]
+        List of section dictionaries with title, image, and caption
+    """
+    output_sections = []
+    for section in sections:
+        section_dict = {
+            "title": section.title or "Untitled Section",
+            "image": None,
+            "caption": RichText("")
+        }
+
+        # Process section outputs to find visualizations and captions
+        image_found = False
+        for i, out in enumerate(section.outputs):
+            if isinstance(out, VegaLiteOutput) and not image_found:
+                # Convert LumenOutput to image
+                image_path = output_to_image(out)
+                if image_path:
+                    section_dict["image"] = InlineImage(doc, image_path, width=Mm(160))
+                    image_found = True
+
+                    # Check if next output is a Typography for caption
+                    if i + 1 < len(section.outputs):
+                        next_out = section.outputs[i + 1]
+                        if isinstance(next_out, Typography):
+                            section_dict["caption"] = RichText(next_out.object)
+                    break
+
+        if section_dict["image"]:  # Only add section if it has an image
+            output_sections.append(section_dict)
+    return output_sections
+
+
+def output_to_image(output: VegaLiteOutput) -> str | None:
+    """
+    Convert a VegaLiteOutput to an image file path.
+
+    Arguments
+    ---------
+    output : VegaLiteOutput
+        The output to convert
+
+    Returns
+    -------
+    str | None
+        Path to temporary image file, or None if conversion failed
+    """
+    # Create a temporary file for the image
+    tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+    tmp_path = tmp.name
+    tmp.close()
+    try:
+        # Render the component and save as image
+        component = output.component
+        with open(tmp_path, 'wb') as f:
+            vega_pane = component.__panel__()._pane
+            vega_pane.param.update(
+                width=650,
+                height=400,
+            )
+            image_bytes = vega_pane.export("png", scale=2, ppi=144)
+            f.write(image_bytes)
+            return tmp_path
+    except Exception as e:
+        warnings.warn(f"Failed to convert output to image: {e}", stacklevel=2)
+        os.unlink(tmp_path)
+        return None