Files

234 lines
8.3 KiB
Python

"""Document-production tools (phase 2 #92) — WRITE, confirmation required.
The assistant can generate real files inside a vault:
* ``create_xlsx`` — spreadsheet (openpyxl);
* ``create_docx`` — Word document (python-docx);
* ``create_csv`` — CSV (stdlib);
* ``create_pdf`` — PDF (reportlab, from markdown-ish content).
Every tool is ``WRITE`` (two-step confirm in the UI / propose-apply over MCP),
vault-scoped through ``requires_vault`` and saved via the shared mutation
service (path safety, read-only check, backup on overwrite).
"""
from __future__ import annotations
import csv as csv_lib
import io
import logging
import re
from typing import Any, cast
# saxutils.escape uniquement (échappement de chaînes, aucun parsing XML).
from xml.sax import saxutils # nosec B406
from backend.services.errors import ServiceError
from backend.services.mutations import save_raw_file
from backend.tools.context import ToolContext, ToolError, ToolRisk
from backend.tools.registry import tool
from backend.tools.schemas import CsvInput, DocxInput, PdfInput, SpreadsheetInput
logger = logging.getLogger("obsigate.tools.documents")
MAX_PDF_CHARS = 200_000
MAX_ROWS = 5_000
def _save(vault: str, path: str, content: bytes, overwrite: bool) -> dict[str, Any]:
"""Shared save helper (maps ServiceError to ToolError)."""
try:
return save_raw_file(vault, path, content, overwrite=overwrite, allow_docs=True)
except ServiceError as e:
raise ToolError(e.message, code=e.code, details=e.details) from e
def _check_rows(rows: list[list[Any]]) -> None:
if not rows:
raise ToolError("Aucune ligne fournie", code="invalid_arguments")
if len(rows) > MAX_ROWS:
raise ToolError(
f"Trop de lignes ({len(rows)} > {MAX_ROWS})", code="invalid_arguments"
)
def _check_extension(path: str, expected: str) -> str:
"""Enforce the document extension; return the normalized path."""
path = (path or "").strip()
if not path.lower().endswith(expected):
raise ToolError(
f"Extension attendue : {expected}", code="invalid_arguments"
)
return path
@tool(
name="create_xlsx",
description=(
"Create an .xlsx spreadsheet in a vault from rows of cell values "
"(first row = header). Use for tables, budgets, checklists the user "
"asked to turn into an Excel file."
),
input_model=SpreadsheetInput,
risk=ToolRisk.WRITE,
requires_vault=True,
)
def create_xlsx(ctx: ToolContext, params: SpreadsheetInput) -> dict[str, Any]:
"""Build the workbook with openpyxl and save it into the vault."""
from openpyxl import Workbook
_check_rows(params.rows)
path = _check_extension(params.path, ".xlsx")
wb = Workbook()
ws = wb.active
ws.title = params.sheet_name[:31] or "Feuille1"
for row in params.rows:
ws.append(list(row))
buffer = io.BytesIO()
wb.save(buffer)
return _save(params.vault, path, buffer.getvalue(), params.overwrite)
@tool(
name="create_docx",
description=(
"Create a .docx Word document in a vault from an optional title and "
"ordered paragraphs. Use for letters, reports, structured drafts."
),
input_model=DocxInput,
risk=ToolRisk.WRITE,
requires_vault=True,
)
def create_docx(ctx: ToolContext, params: DocxInput) -> dict[str, Any]:
"""Build the document with python-docx and save it into the vault."""
from docx import Document
if not params.paragraphs:
raise ToolError("Aucun paragraphe fourni", code="invalid_arguments")
path = _check_extension(params.path, ".docx")
doc = Document()
if params.title.strip():
doc.add_heading(params.title.strip(), level=1)
for paragraph in params.paragraphs:
doc.add_paragraph(paragraph)
buffer = io.BytesIO()
doc.save(buffer)
return _save(params.vault, path, buffer.getvalue(), params.overwrite)
@tool(
name="create_csv",
description=(
"Create a .csv file in a vault from rows of cell values (first row = "
"header). Use for flat data exports, simple tables."
),
input_model=CsvInput,
risk=ToolRisk.WRITE,
requires_vault=True,
)
def create_csv(ctx: ToolContext, params: CsvInput) -> dict[str, Any]:
"""Serialize the rows and save the CSV into the vault."""
_check_rows(params.rows)
path = _check_extension(params.path, ".csv")
delimiter = params.delimiter if params.delimiter in (",", ";", "\t") else ","
buffer = io.StringIO()
writer = csv_lib.writer(buffer, delimiter=delimiter, lineterminator="\n")
writer.writerows(params.rows)
return _save(params.vault, path, buffer.getvalue().encode("utf-8"), params.overwrite)
_HEADING_RE = re.compile(r"^(#{1,6})\s+(.*)$")
def _markdown_to_flowables(content: str) -> list[tuple[str, str]]:
"""Split markdown-ish content into (style, text) blocks for reportlab."""
blocks: list[tuple[str, str]] = []
for raw_line in content.splitlines():
line = raw_line.rstrip()
if not line.strip():
continue
heading = _HEADING_RE.match(line)
if heading:
blocks.append((f"H{min(3, len(heading.group(1)))}", heading.group(2).strip()))
else:
blocks.append(("P", line.strip()))
return blocks
def _render_markdown_pdf(content: str, title: str) -> bytes | None:
"""Render markdown → HTML → PDF through the document-page pipeline.
Uses the same stack as the « Download PDF » button of the document viewer
(mistune with the table plugin + WeasyPrint print CSS), so tables, code
blocks and lists are laid out correctly. Returns ``None`` when WeasyPrint
is not importable (missing GTK on some hosts) so the caller can fall back
to the simplified reportlab renderer.
"""
try:
import mistune
from backend.pdf_export import build_pdf_html, generate_pdf
renderer = mistune.create_markdown(
escape=False,
plugins=["table", "strikethrough", "footnotes", "task_lists"],
)
# mistune 3.3 types `Markdown.__call__` as `str | list[...]` (le
# renderer HTML renvoie toujours `str` à l'exécution).
html = cast(str, renderer(content))
return generate_pdf(build_pdf_html(html, title), title)
except Exception as e:
# WeasyPrint loads GTK lazily: a missing native library can surface at
# import OR render time. Fall back to the simple renderer either way.
logger.warning("WeasyPrint pipeline unavailable for create_pdf: %s", e)
return None
def _render_reportlab_pdf(content: str, title: str) -> bytes:
"""Fallback renderer (no WeasyPrint): headings + paragraphs, no tables."""
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer
styles = getSampleStyleSheet()
style_map = {
"P": styles["BodyText"],
"H1": styles["Heading1"],
"H2": styles["Heading2"],
"H3": styles["Heading3"],
}
buffer = io.BytesIO()
doc = SimpleDocTemplate(buffer, pagesize=A4, title=title[:200])
story: list[Any] = [Paragraph(saxutils.escape(title[:300]), styles["Title"])]
for style, line in _markdown_to_flowables(content):
story.append(Spacer(1, 4))
story.append(Paragraph(saxutils.escape(line), style_map[style]))
doc.build(story)
return buffer.getvalue()
@tool(
name="create_pdf",
description=(
"Create a .pdf document in a vault from markdown content (headings, "
"paragraphs, tables, code blocks, lists). Use for printable "
"deliverables; tables are laid out like the document-page PDF export."
),
input_model=PdfInput,
risk=ToolRisk.WRITE,
requires_vault=True,
)
def create_pdf(ctx: ToolContext, params: PdfInput) -> dict[str, Any]:
"""Render the content and save the PDF into the vault.
Primary path: mistune (tables) + WeasyPrint — identical to the viewer's
« Download PDF » export. Fallback (WeasyPrint unavailable): simplified
reportlab layout without tables.
"""
path = _check_extension(params.path, ".pdf")
content = params.content[:MAX_PDF_CHARS]
pdf_bytes = _render_markdown_pdf(content, params.title[:300])
if pdf_bytes is None:
pdf_bytes = _render_reportlab_pdf(content, params.title[:300])
return _save(params.vault, path, pdf_bytes, params.overwrite)