"""Safe, optional illustrations for a sandbox Markdown full rewrite. The writer streams prose first. When requested, this module derives a small set of heading-scoped illustration prompts, persists the returned bytes inside the same user-owned artifact tree, then adds only local ``/mnt`` references to the Markdown. Provider URLs never reach the browser or the saved document. """ from __future__ import annotations import os import re import tempfile from dataclasses import dataclass from pathlib import Path, PurePosixPath from deerflow.agents.deep_research.adapters.image import ImageGenerationError, ImageProvider _HEADING_RE = re.compile(r"^##[ \t]+([^\r\n]+?)[ \t]*$", re.MULTILINE) _TITLE_RE = re.compile(r"^#[ \t]+([^\r\n]+?)[ \t]*$", re.MULTILINE) _START_MARKER = "" _END_MARKER = "" _MANAGED_IMAGE_BLOCK_RE = re.compile( rf"\n?{re.escape(_START_MARKER)}.*?{re.escape(_END_MARKER)}\n?", re.DOTALL, ) _MIME_EXTENSIONS = { "image/png": "png", "image/jpeg": "jpg", "image/webp": "webp", } @dataclass(frozen=True) class DocumentIllustration: """A provider result that has not been exposed or persisted yet.""" section: str position: int data: bytes mime_type: str alt: str @dataclass(frozen=True) class PersistedDocumentIllustration: """A generated image with a virtual artifact path safe for Markdown.""" section: str position: int virtual_path: str alt: str def strip_managed_document_illustrations(markdown: str) -> str: """Remove only DeerFlow-generated image blocks before a fresh image pass.""" return _MANAGED_IMAGE_BLOCK_RE.sub("\n", markdown).strip() def select_illustration_sections(markdown: str, maximum: int) -> list[str]: """Pick distinct H2 sections; fall back to the title for short documents.""" limit = max(0, min(maximum, 4)) if not limit: return [] sections: list[str] = [] seen: set[str] = set() for match in _HEADING_RE.finditer(markdown): section = _clean_heading(match.group(1)) if section and section not in seen: seen.add(section) sections.append(section) if len(sections) >= limit: return sections if sections: return sections title_match = _TITLE_RE.search(markdown) title = _clean_heading(title_match.group(1)) if title_match else "文档主题" return [title or "文档主题"] async def generate_document_illustrations( provider: ImageProvider, *, markdown: str, document_name: str, maximum: int, on_started=None, on_generated=None, ) -> tuple[list[DocumentIllustration], list[str]]: """Generate a bounded, sequential image set and keep failures recoverable.""" sections = select_illustration_sections(markdown, maximum) title_match = _TITLE_RE.search(markdown) document_title = _clean_heading(title_match.group(1)) if title_match else Path(document_name).stem generated: list[DocumentIllustration] = [] warnings: list[str] = [] for position, section in enumerate(sections, 1): if on_started is not None: await on_started(section, position, len(sections)) try: result = await provider.generate( _illustration_prompt(document_title, section), section_id=section, ) data = result.get("data") if isinstance(result, dict) else None if not isinstance(data, bytes) or not data: raise ImageGenerationError("图片生成服务未返回可保存的图片数据") mime_type = str(result.get("mime_type") or "image/png").lower() if mime_type not in _MIME_EXTENSIONS: mime_type = "image/png" illustration = DocumentIllustration( section=section, position=position, data=data, mime_type=mime_type, alt=_clean_alt(str(result.get("alt") or f"{section}配图")), ) generated.append(illustration) if on_generated is not None: await on_generated(illustration, len(sections), len(generated)) except ImageGenerationError as exc: warnings.append(f"“{section}”配图未生成：{exc}") except Exception: # noqa: BLE001 # Generation remains best effort: an unavailable image provider # must never discard an otherwise complete document rewrite. warnings.append(f"“{section}”配图生成失败，已保留文字改写结果。") return generated, warnings def persist_document_illustrations( illustrations: list[DocumentIllustration], *, markdown_path: Path, virtual_markdown_path: str, job_id: str, ) -> tuple[list[PersistedDocumentIllustration], list[str]]: """Best-effort persist each asset beside the rewritten Markdown file. A single failed local write must not discard earlier successful images. In particular, Windows development reloaders can briefly recreate an output directory while a report is being written. The caller receives warnings for the timeline, while any persisted image remains usable in the final Markdown. """ if not illustrations: return [], [] # Keep this name intentionally short. On Windows the local sandbox path # already contains user/thread UUIDs; including a long report title plus a # UUID-based temporary filename exceeded MAX_PATH (and surfaced as a # misleading FileNotFoundError). Job ids make names unique across files, # so a single short hidden directory per outputs folder is sufficient. asset_directory_name = ".ai-images" asset_directory = markdown_path.parent / asset_directory_name virtual_parent = PurePosixPath(virtual_markdown_path).parent.as_posix().rstrip("/") persisted: list[PersistedDocumentIllustration] = [] warnings: list[str] = [] for image in illustrations: try: extension = _MIME_EXTENSIONS[image.mime_type] file_name = f"{job_id}-{image.position:02d}.{extension}" _atomic_write_bytes(asset_directory / file_name, image.data) persisted.append( PersistedDocumentIllustration( section=image.section, position=image.position, virtual_path=f"{virtual_parent}/{asset_directory_name}/{file_name}", alt=image.alt, ) ) except OSError: warnings.append(f"“{image.section}”配图保存失败，已继续保留其他成功配图。") return persisted, warnings def insert_document_illustrations( markdown: str, illustrations: list[PersistedDocumentIllustration], ) -> str: """Insert fenced local-image blocks directly below their selected sections.""" amended = strip_managed_document_illustrations(markdown) for image in illustrations: marker = f"## {image.section}" block = ( f"\n\n{_START_MARKER}\n" f"![{image.alt}]({image.virtual_path})\n" f"{_END_MARKER}\n" ) if marker in amended: amended = amended.replace(marker, f"{marker}{block}", 1) else: amended = f"{amended.rstrip()}{block}" return amended.strip() def _atomic_write_bytes(path: Path, content: bytes) -> None: """Write one image atomically with Windows-safe temporary path lengths.""" last_error: FileNotFoundError | None = None for _attempt in range(3): temporary: Path | None = None try: # A runtime cleanup can remove a freshly-created asset folder # between image calls. Recreate it immediately before writing and # retry this small local operation. NamedTemporaryFile deliberately # keeps the temporary basename short: appending a report title, # job id, and UUID can exceed Windows MAX_PATH before os.replace. path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) with tempfile.NamedTemporaryFile( mode="wb", dir=path.parent, prefix=".w-", suffix=".tmp", delete=False, ) as file: temporary = Path(file.name) file.write(content) file.flush() os.fsync(file.fileno()) os.replace(temporary, path) return except FileNotFoundError as exc: last_error = exc finally: try: if temporary is not None and temporary.exists(): temporary.unlink() except OSError: pass if last_error is not None: raise last_error def _illustration_prompt(document_title: str, section: str) -> str: return ( "为一份专业 Markdown 文档生成一张可嵌入的章节配图。" f"文档主题：{_clean_heading(document_title)}。章节：{section}。" "画面应信息性、简洁、与正文内容一致，可使用抽象示意、数据可视化风格或场景插画。" "不要在图片中放置可读文字、品牌标志、水印或人物肖像；不得虚构具体数据或事实。" ) def _clean_heading(value: str) -> str: return " ".join(value.replace("\r", " ").replace("\n", " ").split())[:160] def _clean_alt(value: str) -> str: return _clean_heading(value).replace("[", "").replace("]", "") or "文档配图" __all__ = [ "DocumentIllustration", "PersistedDocumentIllustration", "generate_document_illustrations", "insert_document_illustrations", "persist_document_illustrations", "select_illustration_sections", "strip_managed_document_illustrations", ]