deepresearch-flow 0.1.1__py3-none-any.whl

deepresearch_flow/recognize/markdown.py

@@ -0,0 +1,277 @@
+"""Markdown image helpers for recognize commands."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import logging
+import mimetypes
+import re
+from pathlib import Path
+from typing import Awaitable, Callable, Optional
+from urllib.parse import urlparse, unquote
+
+import httpx
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_USER_AGENT = (
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+    "AppleWebKit/537.36 (KHTML, like Gecko) "
+    "Chrome/122.0.0.0 Safari/537.36"
+)
+HTTP_TIMEOUT_SECONDS = 60.0
+
+ALLOWED_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".svg"}
+EXTENSION_OVERRIDES = {
+    ".jpe": ".jpg",
+}
+
+IMAGE_PATTERN = re.compile(r"!\[([^\]]*)\]\(([^)]+)\)")
+DATA_URL_PATTERN = re.compile(r"^data:([^;,]+)(;base64)?,(.*)$", re.DOTALL)
+
+
+class NameRegistry:
+    def __init__(self, output_dir: Path) -> None:
+        self.output_dir = output_dir
+        if output_dir.exists():
+            self.used = {path.name for path in output_dir.iterdir() if path.is_file()}
+        else:
+            self.used = set()
+        self.lock = asyncio.Lock()
+
+    def reserve(self, base: str, ext: str) -> str:
+        base = sanitize_filename(base) or "file"
+        ext = ext if ext.startswith(".") else f".{ext}"
+        candidate = f"{base}{ext}"
+        counter = 0
+        while candidate in self.used or (self.output_dir / candidate).exists():
+            counter += 1
+            candidate = f"{base}_{counter}{ext}"
+        self.used.add(candidate)
+        return candidate
+
+    async def reserve_async(self, base: str, ext: str) -> str:
+        async with self.lock:
+            return self.reserve(base, ext)
+
+
+def sanitize_filename(value: str) -> str:
+    cleaned = re.sub(r"[^\w.\-]+", "_", value.strip())
+    return cleaned.strip("._-")
+
+
+def read_text(path: Path) -> str:
+    try:
+        return path.read_text(encoding="utf-8")
+    except UnicodeDecodeError:
+        return path.read_text(encoding="latin-1")
+
+
+def split_link_target(raw_link: str) -> tuple[str, str, str, str]:
+    link = raw_link.strip()
+    if link.startswith("<"):
+        end = link.find(">")
+        if end != -1:
+            return link[1:end], link[end + 1 :], "<", ">"
+    parts = link.split()
+    if not parts:
+        return "", "", "", ""
+    target = parts[0]
+    suffix = link[len(target) :]
+    return target, suffix, "", ""
+
+
+def is_data_url(target: str) -> bool:
+    return target.startswith("data:")
+
+
+def is_http_url(target: str) -> bool:
+    parsed = urlparse(target)
+    return parsed.scheme in {"http", "https"}
+
+
+def resolve_local_path(md_path: Path, target: str) -> Path:
+    target_path = Path(unquote(target))
+    if target_path.is_absolute():
+        return target_path
+    return (md_path.parent / target_path).resolve()
+
+
+def extension_from_mime(mime: str) -> Optional[str]:
+    ext = mimetypes.guess_extension(mime, strict=False)
+    if ext in EXTENSION_OVERRIDES:
+        return EXTENSION_OVERRIDES[ext]
+    return ext
+
+
+def mime_from_path(path: Path) -> Optional[str]:
+    mime, _ = mimetypes.guess_type(path.name)
+    if mime:
+        return mime
+    if path.suffix.lower() in {".jpg", ".jpeg"}:
+        return "image/jpeg"
+    if path.suffix.lower() == ".png":
+        return "image/png"
+    if path.suffix.lower() == ".gif":
+        return "image/gif"
+    if path.suffix.lower() == ".webp":
+        return "image/webp"
+    if path.suffix.lower() == ".svg":
+        return "image/svg+xml"
+    return None
+
+
+def parse_data_url(target: str) -> Optional[tuple[str, bytes]]:
+    match = DATA_URL_PATTERN.match(target)
+    if not match:
+        return None
+    mime = match.group(1) or ""
+    if not mime.startswith("image/"):
+        return None
+    if match.group(2) != ";base64":
+        return None
+    payload = match.group(3) or ""
+    try:
+        return mime, base64.b64decode(payload)
+    except Exception as exc:  # pragma: no cover - defensive
+        logger.warning("Failed to decode base64 image: %s", exc)
+        return None
+
+
+def data_url_from_bytes(mime: str, data: bytes) -> str:
+    encoded = base64.b64encode(data).decode("ascii")
+    return f"data:{mime};base64,{encoded}"
+
+
+def base_name_from_alt(alt_text: str) -> str:
+    if not alt_text:
+        return ""
+    candidate = sanitize_filename(alt_text)
+    if not candidate:
+        return ""
+    suffix = Path(candidate).suffix.lower()
+    if suffix in ALLOWED_IMAGE_EXTS:
+        return Path(candidate).stem
+    return candidate
+
+
+def hash_name_from_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()[:12]
+
+
+async def rewrite_markdown_images(
+    content: str,
+    replacer: Callable[[str, str], Awaitable[Optional[str]]],
+) -> str:
+    output: list[str] = []
+    last_idx = 0
+    for match in IMAGE_PATTERN.finditer(content):
+        output.append(content[last_idx : match.start()])
+        alt_text = match.group(1)
+        raw_link = match.group(2)
+        target, suffix, prefix, postfix = split_link_target(raw_link)
+        new_target = await replacer(alt_text, target)
+        if new_target is None:
+            output.append(match.group(0))
+        else:
+            new_link = f"{prefix}{new_target}{postfix}{suffix}"
+            output.append(f"![{alt_text}]({new_link})")
+        last_idx = match.end()
+    output.append(content[last_idx:])
+    return "".join(output)
+
+
+def count_markdown_images(content: str) -> dict[str, int]:
+    counts = {"total": 0, "data": 0, "http": 0, "local": 0}
+    for match in IMAGE_PATTERN.finditer(content):
+        counts["total"] += 1
+        raw_link = match.group(2)
+        target, _, _, _ = split_link_target(raw_link)
+        if not target:
+            continue
+        if is_data_url(target):
+            counts["data"] += 1
+        elif is_http_url(target):
+            counts["http"] += 1
+        else:
+            counts["local"] += 1
+    return counts
+
+
+async def embed_markdown_images(
+    content: str,
+    md_path: Path,
+    enable_http: bool,
+    http_client: Optional[httpx.AsyncClient],
+) -> str:
+    async def replacer(alt_text: str, target: str) -> Optional[str]:
+        if not target:
+            return None
+        if is_data_url(target):
+            return None
+        if is_http_url(target):
+            if not enable_http or http_client is None:
+                return None
+            try:
+                response = await http_client.get(target)
+            except Exception as exc:
+                logger.warning("Failed to fetch %s: %s", target, exc)
+                return None
+            if response.status_code >= 400:
+                logger.warning("Failed to fetch %s: HTTP %d", target, response.status_code)
+                return None
+            content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
+            if not content_type.startswith("image/"):
+                guessed = mime_from_path(Path(urlparse(target).path))
+                if not guessed or not guessed.startswith("image/"):
+                    logger.warning(
+                        "Skipping non-image URL %s (Content-Type %s)", target, content_type
+                    )
+                    return None
+                content_type = guessed
+            return data_url_from_bytes(content_type, response.content)
+
+        local_path = resolve_local_path(md_path, target)
+        if not local_path.exists() or not local_path.is_file():
+            logger.warning("Image not found: %s", local_path)
+            return None
+        mime = mime_from_path(local_path)
+        if not mime or not mime.startswith("image/"):
+            logger.warning("Unsupported image type: %s", local_path)
+            return None
+        data = await asyncio.to_thread(local_path.read_bytes)
+        return data_url_from_bytes(mime, data)
+
+    return await rewrite_markdown_images(content, replacer)
+
+
+async def unpack_markdown_images(
+    content: str,
+    images_dir: Path,
+    name_registry: NameRegistry,
+) -> str:
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    async def replacer(alt_text: str, target: str) -> Optional[str]:
+        if not is_data_url(target):
+            return None
+        parsed = parse_data_url(target)
+        if parsed is None:
+            return None
+        mime, data = parsed
+        ext = extension_from_mime(mime)
+        if not ext:
+            logger.warning("Unsupported MIME type: %s", mime)
+            return None
+        base_name = base_name_from_alt(alt_text)
+        if not base_name:
+            base_name = hash_name_from_bytes(data)
+        filename = await name_registry.reserve_async(base_name, ext)
+        dest_path = images_dir / filename
+        await asyncio.to_thread(dest_path.write_bytes, data)
+        return f"images/{filename}"
+
+    return await rewrite_markdown_images(content, replacer)

deepresearch_flow/recognize/organize.py

@@ -0,0 +1,95 @@
+"""OCR output organizers for recognize commands."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import shutil
+from pathlib import Path
+from typing import Iterable
+
+from deepresearch_flow.recognize.markdown import (
+    NameRegistry,
+    embed_markdown_images,
+    read_text,
+    rewrite_markdown_images,
+    resolve_local_path,
+    is_data_url,
+    is_http_url,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def discover_mineru_dirs(inputs: Iterable[str], recursive: bool) -> list[Path]:
+    results: set[Path] = set()
+    for raw in inputs:
+        path = Path(raw)
+        if path.is_file():
+            if path.name != "full.md":
+                raise FileNotFoundError(f"Expected full.md file but got: {path}")
+            parent = path.parent.resolve()
+            if (parent / "images").is_dir():
+                results.add(parent)
+            else:
+                logger.warning("Skipping %s (missing images/)", parent)
+            continue
+        if not path.exists():
+            raise FileNotFoundError(f"Input path not found: {path}")
+        if path.is_dir():
+            if (path / "full.md").is_file():
+                if (path / "images").is_dir():
+                    results.add(path.resolve())
+                else:
+                    logger.warning("Skipping %s (missing images/)", path)
+            pattern = path.rglob("full.md") if recursive else path.glob("full.md")
+            for full_path in pattern:
+                parent = full_path.parent.resolve()
+                if (parent / "images").is_dir():
+                    results.add(parent)
+                else:
+                    logger.warning("Skipping %s (missing images/)", parent)
+            continue
+        raise FileNotFoundError(f"Input path not found: {path}")
+    return sorted(results)
+
+
+async def organize_mineru_dir(
+    layout_dir: Path,
+    output_simple: Path | None,
+    output_base64: Path | None,
+    output_filename: str,
+    image_registry: NameRegistry | None,
+) -> None:
+    md_path = layout_dir / "full.md"
+    content = await asyncio.to_thread(read_text, md_path)
+
+    if output_simple is not None and image_registry is not None:
+        images_dir = output_simple / "images"
+        images_dir.mkdir(parents=True, exist_ok=True)
+        image_map: dict[Path, str] = {}
+
+        async def replace_simple(_: str, target: str) -> str | None:
+            if not target or is_data_url(target) or is_http_url(target):
+                return None
+            source_path = resolve_local_path(md_path, target)
+            if not source_path.exists() or not source_path.is_file():
+                logger.warning("Image not found: %s", source_path)
+                return None
+            if source_path in image_map:
+                return f"images/{image_map[source_path]}"
+            filename = await image_registry.reserve_async(source_path.stem, source_path.suffix)
+            dest_path = images_dir / filename
+            await asyncio.to_thread(shutil.copy2, source_path, dest_path)
+            image_map[source_path] = filename
+            return f"images/{filename}"
+
+        updated = await rewrite_markdown_images(content, replace_simple)
+        output_path = output_simple / output_filename
+        await asyncio.to_thread(output_path.write_text, updated, encoding="utf-8")
+
+    if output_base64 is not None:
+        updated = await embed_markdown_images(content, md_path, False, None)
+        output_path = output_base64 / output_filename
+        await asyncio.to_thread(output_path.write_text, updated, encoding="utf-8")