treesearchlib 1.1.0__py3-none-any.whl

treesearch/parsers/image_metadata.py

@@ -0,0 +1,345 @@
+# -*- coding: utf-8 -*-
+"""图像元数据网关（deep module）—— vision 解读结果写回 / 读回 + 内容指纹。
+
+载体（ADR-0009 + 工单 01 spike 实证）：
+- JPEG/JPG：EXIF ``XPComment`` (0x9c9c, UTF-16LE)，经 piexif 无损写入；
+  Windows 资源管理器「备注」可读完整中文。
+- PNG / WebP：工单 05/06 实现，本期 ``read_back`` 返回 None、其余抛 NotImplementedError。
+
+``XPComment`` payload 为「markdown 主体 + 末尾 HTML 注释元数据」，人机共用——
+Windows「备注」里人眼能看到描述主体，宿主解析末尾注释取版本::
+
+    {markdown}
+
+    <!--doclens model={model_tag} prompt={prompt_version}-->
+
+设计目标：解读结果「跟文件走」；写回不改像素、不改内容指纹（避免增量死循环）。
+"""
+from __future__ import annotations
+
+import hashlib
+import logging
+import os
+import re
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# 走 vision 解读 + 写回的图像格式（ADR-0009：仅这四种）
+INTERPRETED_IMAGE_EXTS: frozenset[str] = frozenset({".jpg", ".jpeg", ".png", ".webp"})
+
+# 全图像格式集合（含本期不解读的，供「其他图像格式跳过」判断，工单 07）
+ALL_IMAGE_EXTS: frozenset[str] = frozenset({
+    ".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp", ".tiff", ".tif",
+})
+
+_JPEG_EXTS: frozenset[str] = frozenset({".jpg", ".jpeg"})
+
+# EXIF XPComment tag（Microsoft 扩展，UTF-16LE 字符串）
+_XPCOMMENT_TAG = 0x9C9C
+
+# payload 末尾元数据标记（HTML 注释，markdown 不渲染、Windows 备注可见但不喧宾夺主）
+_META_MARKER = "<!--doclens"
+_META_RE = re.compile(r"model=(?P<model>\S+)\s+prompt=(?P<prompt>\S+?)\s*-->")
+
+
+def _ext(path: str) -> str:
+    return os.path.splitext(path)[1].lower()
+
+
+def _encode_xpcomment(text: str) -> bytes:
+    """XPComment：UTF-16LE + 双字节 null 终止。"""
+    return text.encode("utf-16-le") + b"\x00\x00"
+
+
+def _decode_xpcomment(raw) -> str:
+    if not raw:
+        return ""
+    if isinstance(raw, (tuple, list)):  # piexif.load 对 BYTE 类返回 tuple[int, ...]
+        raw = bytes(raw)
+    # 先在完整 bytes 上 decode（保 UTF-16LE 双字节对齐），再在 str 层去尾部 null 字符。
+    # 不能先 rstrip 字节——会误删高字节为 0x00 的字符（如 '>' = 0x003E），破坏对齐。
+    try:
+        text = raw.decode("utf-16-le")
+    except UnicodeDecodeError:
+        text = raw.decode("utf-16-le", errors="replace")
+    return text.rstrip("\x00")
+
+
+def _format_payload(markdown: str, model_tag: str, prompt_version: str) -> str:
+    return f"{markdown}\n\n{_META_MARKER} model={model_tag} prompt={prompt_version}-->"
+
+
+def _parse_payload(text: str) -> dict:
+    """从 XPComment 文本解析 payload；无元数据标记则视为纯描述（版本字段为空）。"""
+    idx = text.rfind(_META_MARKER)
+    if idx < 0:
+        return {"markdown": text, "model_tag": "", "prompt_version": ""}
+    markdown = text[:idx].rstrip("\n").rstrip()
+    m = _META_RE.search(text[idx:])
+    return {
+        "markdown": markdown,
+        "model_tag": m.group("model") if m else "",
+        "prompt_version": m.group("prompt") if m else "",
+    }
+
+
+# 期望版本：宿主在 reindex 前经 set_expected_version 设置；read_back 默认据此校验，
+# payload 版本不符则当无解读 → image_to_tree 占位入队重解读（工单 08）。
+_expected_model_tag: Optional[str] = None
+_expected_prompt_version: Optional[str] = None
+
+
+def set_expected_version(
+    model_tag: Optional[str] = None, prompt_version: Optional[str] = None
+) -> None:
+    """设置当前期望的解读版本（宿主在 reindex 前调用）。"""
+    global _expected_model_tag, _expected_prompt_version
+    _expected_model_tag = model_tag
+    _expected_prompt_version = prompt_version
+
+
+# write_back 失败计数（工单 09 可观测；经 writeback_failure_count 读、reset 清零）
+_writeback_failures: int = 0
+
+
+def writeback_failure_count() -> int:
+    """返回累计的 write_back 失败次数（供 status API 展示）。"""
+    return _writeback_failures
+
+
+def reset_writeback_failures() -> None:
+    global _writeback_failures
+    _writeback_failures = 0
+
+
+def read_back(
+    image_path: str,
+    *,
+    model_tag: Optional[str] = None,
+    prompt_version: Optional[str] = None,
+) -> Optional[dict]:
+    """读图像元数据里的解读 payload。
+
+    返回 ``{"markdown", "model_tag", "prompt_version"}``；无元数据返回 None。
+    传入 ``model_tag`` / ``prompt_version`` 且与 payload 不符时返回 None（版本不符）。
+    """
+    ext = _ext(image_path)
+    if ext in _JPEG_EXTS:
+        text = _read_jpeg_xpcomment(image_path)
+    elif ext == ".png":
+        text = _read_png(image_path)
+    elif ext == ".webp":
+        text = _read_webp(image_path)
+    else:
+        # 其他格式不解读
+        return None
+    if not text:
+        return None
+    payload = _parse_payload(text)
+    # 手动备注（model_tag == "manual"）是用户显式覆盖，不参与版本过期校验，
+    # 持久保留——否则切模型/重建会被 read_back 判过期 → 占位重解析覆盖。
+    if payload.get("model_tag") == "manual":
+        return payload
+    mt = model_tag if model_tag is not None else _expected_model_tag
+    pv = prompt_version if prompt_version is not None else _expected_prompt_version
+    if mt is not None and payload["model_tag"] != mt:
+        return None
+    if pv is not None and payload["prompt_version"] != pv:
+        return None
+    return payload
+
+
+def _read_jpeg_xpcomment(image_path: str) -> Optional[str]:
+    try:
+        import piexif
+    except ImportError:
+        logger.warning("piexif 未安装，无法读 JPEG 元数据")
+        return None
+    try:
+        exif_dict = piexif.load(image_path)
+    except Exception as e:  # 无 EXIF / 损坏 → piexif 抛异常
+        logger.debug("piexif.load 失败 %s: %s", image_path, e)
+        return None
+    raw = exif_dict.get("0th", {}).get(_XPCOMMENT_TAG)
+    if raw is None:
+        return None
+    return _decode_xpcomment(raw)
+
+
+def write_back(
+    image_path: str,
+    markdown: str,
+    *,
+    model_tag: str,
+    prompt_version: str,
+) -> bool:
+    """无损写入解读 payload 到图像元数据。失败返回 False（不抛），并计入失败计数（工单 09）。
+
+    JPEG：合并已有 EXIF（只改 XPComment，不覆盖其他字段），piexif 无损 insert。
+    """
+    ext = _ext(image_path)
+    if ext in _JPEG_EXTS:
+        ok = _write_jpeg(image_path, markdown, model_tag, prompt_version)
+    elif ext == ".png":
+        ok = _write_png(image_path, markdown, model_tag, prompt_version)
+    elif ext == ".webp":
+        ok = _write_webp(image_path, markdown, model_tag, prompt_version)
+    else:
+        logger.warning("write_back 暂不支持 %s", ext)
+        ok = False
+    if not ok:
+        global _writeback_failures
+        _writeback_failures += 1
+    return ok
+
+
+def _write_jpeg(
+    image_path: str, markdown: str, model_tag: str, prompt_version: str
+) -> bool:
+    try:
+        import piexif
+    except ImportError:
+        logger.error("piexif 未安装，无法写 JPEG 元数据")
+        return False
+    payload = _format_payload(markdown, model_tag, prompt_version)
+    try:
+        # 合并已有 EXIF：无 EXIF 时 piexif.load 抛异常 → 用空 0th
+        try:
+            exif_dict = piexif.load(image_path)
+            zeroth = dict(exif_dict.get("0th", {}))
+        except Exception:
+            zeroth = {}
+        zeroth[_XPCOMMENT_TAG] = _encode_xpcomment(payload)
+        exif_bytes = piexif.dump({"0th": zeroth})
+        piexif.insert(exif_bytes, image_path)  # 无损：只动 APP1 EXIF，不碰像素
+        return True
+    except Exception as e:
+        logger.warning("write_back JPEG 失败 %s: %s", image_path, e)
+        return False
+
+
+# PNG iTXt chunk 的 key（UTF-8，doclens 自读自写；Windows 资源管理器不读 PNG 元数据）
+_PNG_KEY = "Description"
+
+
+def _read_png(image_path: str) -> Optional[str]:
+    try:
+        from PIL import Image
+        with Image.open(image_path) as img:
+            img.load()
+            text = img.info.get(_PNG_KEY)
+        return text if isinstance(text, str) else None
+    except Exception as e:
+        logger.debug("读 PNG 元数据失败 %s: %s", image_path, e)
+        return None
+
+
+def _write_png(
+    image_path: str, markdown: str, model_tag: str, prompt_version: str
+) -> bool:
+    """PNG：Pillow ``PngInfo`` iTXt（UTF-8）。
+
+    ``save`` 会重编码，但 PNG 无损 → 像素不变 → ``content_fingerprint`` 不变（不死循环）。
+    仅写 ``_PNG_KEY``，不保留原图其他 tEXt（PNG 的 tEXt 多为软件元信息，可接受丢失）。
+    """
+    try:
+        from PIL import Image
+        from PIL.PngImagePlugin import PngInfo
+        payload = _format_payload(markdown, model_tag, prompt_version)
+        with Image.open(image_path) as img:
+            img.load()
+            pi = PngInfo()
+            pi.add_itxt(_PNG_KEY, payload)
+            img.save(image_path, "PNG", pnginfo=pi)
+        return True
+    except Exception as e:
+        logger.warning("write_back PNG 失败 %s: %s", image_path, e)
+        return False
+
+
+def _webp_xmp_packet(payload: str) -> bytes:
+    """把 payload 包成 XMP packet（dc:description）。XML-escape 防止 payload 里的
+    ``<!--`` 破坏 XML 结构；``getxmp`` 读回时自动 unescape。"""
+    import xml.sax.saxutils as su
+    desc = su.escape(payload)
+    return (
+        b'<?xpacket begin="\xef\xbb\xbf" id="W5M0MpCehiHzreSzNTczkc9d"?>'
+        b'<x:xmpmeta xmlns:x="adobe:ns:meta/">'
+        b'<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">'
+        b'<rdf:Description rdf:about="" xmlns:dc="http://purl.org/dc/elements/1.1/">'
+        b'<dc:description><rdf:Alt>'
+        b'<rdf:li xml:lang="x-default">' + desc.encode("utf-8") + b'</rdf:li>'
+        b'</rdf:Alt></dc:description>'
+        b'</rdf:Description></rdf:RDF></x:xmpmeta>'
+        b'<?xpacket end="w"?>'
+    )
+
+
+def _read_webp(image_path: str) -> Optional[str]:
+    """读 WebP XMP 的 dc:description 文本（getxmp 自动 unescape）。"""
+    try:
+        from PIL import Image
+        with Image.open(image_path) as img:
+            img.load()
+            try:
+                xmp = img.getxmp()
+            except Exception:
+                return None
+        li = (
+            (((xmp.get("xmpmeta") or {}).get("RDF") or {}).get("Description") or {})
+            .get("description", {}).get("Alt", {}).get("li")
+        )
+        if isinstance(li, dict):
+            t = li.get("text")
+            if isinstance(t, str):
+                return t
+        return None
+    except Exception as e:
+        logger.debug("读 WebP XMP 失败 %s: %s", image_path, e)
+        return None
+
+
+def _write_webp(
+    image_path: str, markdown: str, model_tag: str, prompt_version: str
+) -> bool:
+    """WebP：Pillow ``save(xmp=, lossless=True)``。lossless 保像素 → content_fingerprint 不变。"""
+    try:
+        from PIL import Image
+        payload = _format_payload(markdown, model_tag, prompt_version)
+        xmp = _webp_xmp_packet(payload)
+        with Image.open(image_path) as img:
+            img.load()
+            img.save(image_path, "WEBP", lossless=True, xmp=xmp)
+        return True
+    except Exception as e:
+        logger.warning("write_back WebP 失败 %s: %s", image_path, e)
+        return False
+
+
+def content_fingerprint(image_path: str) -> str:
+    """对图像核心内容（解码像素）算指纹（md5）。
+
+    用解码像素而非文件字节，确保写回元数据（EXIF/XMP，含 piexif 写入时对段结构
+    的调整）不改变指纹 —— 这是避免「写回 → file_hash 变 → 增量重解析」死循环的关键
+    （工单 03 把它注入 file_hash_fn）。对所有图像格式通用（JPEG/PNG/WebP）。
+
+    注：spike 曾试「剥离 EXIF 段算字节 hash」，但 piexif.insert 会移除原图 APP0/JFIF
+    段（副作用、像素不变），导致字节 hash 漂移；像素 hash 完全不受段结构影响，故改用。
+    """
+    ext = _ext(image_path)
+    if ext in ALL_IMAGE_EXTS:
+        return _pixel_hash(image_path)
+    raise NotImplementedError(f"content_fingerprint 不支持 {ext}")
+
+
+def _pixel_hash(image_path: str) -> str:
+    """解码像素（RGB）的 md5。元数据/段结构不影响，对所有图像格式通用。"""
+    try:
+        from PIL import Image
+        with Image.open(image_path) as img:
+            return hashlib.md5(img.convert("RGB").tobytes()).hexdigest()
+    except Exception as e:
+        logger.warning("像素指纹失败 %s: %s，回退文件字节 hash", image_path, e)
+        with open(image_path, "rb") as f:
+            return hashlib.md5(f.read()).hexdigest()

treesearch/parsers/image_parser.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+"""独立图像文件 parser —— 占位节点，真正的视觉解析由后台 Vision Worker 完成。
+
+架构（见 docs/adr/0001-vision-image-indexing.md）：
+视觉识别耗时长，不能拖慢主索引。本 parser 不调任何视觉 API，只产出
+「占位节点」树（标题=文件名，保证文件名立即可搜索），并通过
+``result["vision_pending"]`` 标记通知 indexer 把文件写入 vision_queue 表；
+宿主侧的 VisionWorker（如 doclens）后台串行消费队列，拿到视觉模型输出的
+Markdown 后复用 ``md_to_tree`` 建树，原位替换占位节点。
+"""
+import logging
+import os
+
+from ..indexer import md_to_tree
+
+logger = logging.getLogger(__name__)
+
+# 独立图像文件扩展名 —— 仅这四种走 vision 解读 + 元数据写回（ADR-0009 / 工单 07）。
+# gif/bmp/tiff/tif 不再作为图像处理（不 image_to_tree、不占位、不入 vision 队列）；
+# svg 是文本走 code/xml，不在此列。
+IMAGE_EXTENSIONS: frozenset[str] = frozenset({
+    ".png", ".jpg", ".jpeg", ".webp",
+})
+
+# 占位节点正文（解析完成后被视觉模型输出的 Markdown 原位替换）
+PLACEHOLDER_TEXT = "（图像文件已登记，等待后台视觉解析；当前仅文件名可搜索）"
+
+
+async def image_to_tree(image_path: str, **kwargs) -> dict:
+    """为图像文件构建树：已有解读就读回、否则占位入队。
+
+    - **已有 vision 解读元数据**（``read_back`` 命中）：用解读 Markdown 走 ``md_to_tree``
+      建树，**不设** ``vision_pending``（indexer 据此跳过入队）——这是 force 重建不重花
+      API 的闭环（ADR-0009 / 工单 04）。
+    - **无解读**：单个 ``# <文件名>`` 根节点 + 占位正文，设 ``vision_pending`` 入队
+      让 Vision Worker 后台解析。
+
+    走 ``md_to_tree`` 以获得与其他类型一致的 node_id / summary 后处理。
+    """
+    from .image_metadata import read_back
+
+    abs_path = os.path.abspath(image_path)
+    doc_name = os.path.splitext(os.path.basename(image_path))[0]
+
+    payload = read_back(abs_path)
+    markdown = payload["markdown"] if payload and payload.get("markdown") else None
+    pending = markdown is None
+    if pending:
+        markdown = f"# {doc_name}\n\n{PLACEHOLDER_TEXT}\n"
+
+    result = await md_to_tree(md_content=markdown, **kwargs)
+    # md_to_tree(md_content=...) 的 doc_name 是 "untitled"，且没有 source_path，
+    # 这里补上真实值，保证增量索引 / preview 反查正常工作。
+    result["doc_name"] = doc_name
+    result["source_path"] = abs_path
+    if pending:
+        # 通知 indexer._index_one 写入 vision_queue（worker 据此后台解析）
+        result["vision_pending"] = True
+    return result

treesearch/parsers/image_store.py

@@ -0,0 +1,182 @@
+# -*- coding: utf-8 -*-
+"""图片落盘存储 —— 把富文档里提取出的图片写成磁盘文件并维护元信息。
+
+供 docx / pptx parser 在解析阶段调用：落盘 + 文档内按 blob sha256 去重，
+返回每个图片引用的 markdown 内联语法。预览端点用 ``resolve`` 反查文件。
+"""
+import hashlib
+import json
+import logging
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from urllib.parse import quote
+
+logger = logging.getLogger(__name__)
+
+# 扩展名 → MIME（公共 API：宿主按扩展名取 media type；缺省 application/octet-stream）
+EXT_TO_MEDIA: dict[str, str] = {
+    "png": "image/png",
+    "jpg": "image/jpeg",
+    "jpeg": "image/jpeg",
+    "gif": "image/gif",
+    "bmp": "image/bmp",
+    "webp": "image/webp",
+    "tiff": "image/tiff",
+    "tif": "image/tiff",
+    "svg": "image/svg+xml",
+    "emf": "image/emf",
+    "wmf": "image/wmf",
+    "ico": "image/x-icon",
+}
+
+_META_FILENAME = "_meta.json"
+
+
+@dataclass(frozen=True)
+class ImagePart:
+    """待落盘的一张图片。
+
+    Attributes:
+        blob: 图片二进制。
+        ext: 扩展名（不含点，如 "png"）。
+        source_ref: 溯源键，docx 用 rId，pptx 用 "slide{idx}:{shape_id}"。
+        disp_w: 文档内显示宽（px），供前端 icon 判定；None=未知（前端退回 naturalWidth）。
+    """
+    blob: bytes
+    ext: str
+    source_ref: str
+    disp_w: int | None = None
+
+
+@dataclass(frozen=True)
+class ImageRef:
+    """落盘后返回的图片引用。"""
+    seq: int
+    inline_md: str
+
+
+def doc_hash_for(rel_path: str) -> str:
+    """rel_path → 12 位 sha256，用作图片子目录名。"""
+    return hashlib.sha256(rel_path.encode("utf-8")).hexdigest()[:12]
+
+
+def _normalize_ext(ext: str) -> str:
+    return (ext or "").lower().lstrip(".")
+
+
+def _inline_md(seq: int, rel_path: str, disp_w: int | None = None) -> str:
+    url = f"/api/preview/asset?path={quote(rel_path, safe='')}&id={seq}"
+    if disp_w is not None:
+        url += f"&dw={disp_w}"
+    return f"![图片 {seq}]({url})"
+
+
+class ImageStore:
+    """图片落盘存储（无运行时可变状态，仅持有 images_root）。
+
+    所有方法基于磁盘文件操作，可被索引阶段与预览端点各自独立实例化，
+    只要 ``images_root`` 指向同一目录即可。
+    """
+
+    def __init__(self, images_root: Path):
+        self._root = Path(images_root)
+
+    @property
+    def root(self) -> Path:
+        return self._root
+
+    def _doc_dir(self, doc_hash: str) -> Path:
+        return self._root / doc_hash
+
+    def extract_for_doc(
+        self,
+        rel_path: str,
+        parts: list[ImagePart],
+    ) -> dict[str, ImageRef]:
+        """把一个文档的图片落盘 + 去重。
+
+        幂等：每次调用先清空该文档的图目录，再重新提取（保证重索引时无旧图残留）。
+        文档内按 ``sha256(blob)`` 去重，同图只落一份，多个 source_ref 指向同一 seq。
+        """
+        dh = doc_hash_for(rel_path)
+        doc_dir = self._doc_dir(dh)
+        if doc_dir.exists():
+            shutil.rmtree(doc_dir, ignore_errors=True)
+        if not parts:
+            return {}
+        doc_dir.mkdir(parents=True, exist_ok=True)
+
+        meta: dict[str, dict] = {}
+        refs: dict[str, ImageRef] = {}
+        blob_to_seq: dict[str, int] = {}
+        seq = 0
+        for part in parts:
+            sha = hashlib.sha256(part.blob).hexdigest()
+            if sha in blob_to_seq:
+                s = blob_to_seq[sha]
+            else:
+                seq += 1
+                s = seq
+                ext = _normalize_ext(part.ext) or "png"
+                fname = f"{s}.{ext}"
+                try:
+                    (doc_dir / fname).write_bytes(part.blob)
+                except OSError as e:
+                    logger.warning("Failed to write image %s: %s", fname, e)
+                    continue
+                entry = {
+                    "sha256": sha,
+                    "media_type": EXT_TO_MEDIA.get(ext, "application/octet-stream"),
+                    "filename": fname,
+                }
+                if part.disp_w is not None:
+                    entry["disp_w"] = part.disp_w
+                meta[str(s)] = entry
+                blob_to_seq[sha] = s
+            refs[part.source_ref] = ImageRef(seq=s, inline_md=_inline_md(s, rel_path, part.disp_w))
+
+        self._write_meta(doc_dir, meta)
+        return refs
+
+    def resolve(self, doc_hash: str, seq: int) -> tuple[Path, str] | None:
+        """端点用：按 doc_hash + seq 返回 (文件路径, media_type)，缺失返回 None。"""
+        doc_dir = self._doc_dir(doc_hash)
+        meta = self._load_meta(doc_dir)
+        entry = meta.get(str(seq))
+        if not entry:
+            return None
+        path = doc_dir / entry["filename"]
+        if not path.exists():
+            return None
+        return path, entry.get("media_type", "application/octet-stream")
+
+    def purge_doc(self, rel_path: str) -> None:
+        """删除单个文档的图目录（用于源文件被删除的清理）。"""
+        doc_dir = self._doc_dir(doc_hash_for(rel_path))
+        if doc_dir.exists():
+            shutil.rmtree(doc_dir, ignore_errors=True)
+
+    def purge_all(self) -> None:
+        """清空整个 images_root（用于 force 全量重建）。"""
+        if self._root.exists():
+            shutil.rmtree(self._root, ignore_errors=True)
+        self._root.mkdir(parents=True, exist_ok=True)
+
+    def _load_meta(self, doc_dir: Path) -> dict:
+        mp = doc_dir / _META_FILENAME
+        if not mp.exists():
+            return {}
+        try:
+            return json.loads(mp.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            logger.warning("image meta corrupted, ignoring: %s", mp)
+            return {}
+
+    def _write_meta(self, doc_dir: Path, meta: dict) -> None:
+        """原子写 _meta.json：临时文件 + os.replace。"""
+        mp = doc_dir / _META_FILENAME
+        tmp = mp.with_suffix(".tmp")
+        tmp.write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8")
+        os.replace(tmp, mp)