264 lines
9.7 KiB
Python
264 lines
9.7 KiB
Python
"""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 = "<!-- deerflow:ai-rewrite-image:start -->"
|
|
_END_MARKER = "<!-- deerflow:ai-rewrite-image:end -->"
|
|
_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"\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",
|
|
]
|