chunksmith-agent 0.4.0__py3-none-any.whl

chunksmith_agent/__init__.py

@@ -0,0 +1,13 @@
+"""Reasoning-based Q&A over ChunkSmith multimodal indexes (outline tree + elements)."""
+
+from chunksmith_agent.agent import ChunkSmithAgent
+from chunksmith_agent.index_builder import build_document_index, build_document_index_from_saved
+from chunksmith_agent.models import AgentAnswer, DocumentIndex
+
+__all__ = [
+    "AgentAnswer",
+    "ChunkSmithAgent",
+    "DocumentIndex",
+    "build_document_index",
+    "build_document_index_from_saved",
+]

chunksmith_agent/agent.py

@@ -0,0 +1,59 @@
+"""ChunkSmith agent session: Q&A over a built document index."""
+
+from __future__ import annotations
+
+from typing import Any, Iterator
+
+from chunksmith_agent.index_builder import build_document_index
+from chunksmith_agent.models import AgentAnswer, DocumentIndex
+from chunksmith_agent.retrieval import answer_question, iter_answer_events
+from chunksmith_agent.session import AgentConversation
+from chunksmith_agent.settings import AgentSettings, load_settings
+
+
+class ChunkSmithAgent:
+    """Holds a document index and answers with session memory (reuse sections on follow-ups)."""
+
+    def __init__(
+        self,
+        index: DocumentIndex,
+        settings: AgentSettings | None = None,
+    ) -> None:
+        self.index = index
+        self.settings = settings or load_settings()
+        self.conversation = AgentConversation()
+
+    @classmethod
+    def from_multimodal_output(cls, out: dict[str, Any], settings: AgentSettings | None = None) -> ChunkSmithAgent:
+        return cls(build_document_index(out), settings=settings)
+
+    def reset_conversation(self) -> None:
+        """Clear chat history and cached section selection."""
+        self.conversation = AgentConversation()
+
+    def ask(self, query: str, *, stream_tokens: bool = False) -> AgentAnswer:
+        return answer_question(
+            self.index,
+            query,
+            self.settings,
+            conversation=self.conversation,
+        )
+
+    def ask_events(
+        self,
+        query: str,
+        *,
+        event_sink: Any | None = None,
+        emit_image_events: bool = True,
+        emit_table_events: bool = True,
+    ) -> Iterator[tuple[str, dict[str, Any]]]:
+        """Yield ``(event_name, payload)`` for CLI streaming."""
+        yield from iter_answer_events(
+            self.index,
+            query,
+            self.settings,
+            event_sink=event_sink,
+            emit_image_events=emit_image_events,
+            emit_table_events=emit_table_events,
+            conversation=self.conversation,
+        )

chunksmith_agent/element_retrieval.py

@@ -0,0 +1,164 @@
+"""Element helpers for building per-node media from canonical bundles (standalone)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+ANCHOR_START_KEYS = ("split_document_anchor_start", "anchor_start")
+ANCHOR_END_KEYS = ("split_document_anchor_end", "anchor_end")
+FALLBACK_START_KEYS = ("split_document_anchor", "anchor")
+
+
+def anchor_start_from_row(row: dict[str, Any]) -> str:
+    for key in ANCHOR_START_KEYS:
+        value = str(row.get(key) or "").strip()
+        if value:
+            return value[:240]
+    for key in FALLBACK_START_KEYS:
+        value = str(row.get(key) or "").strip()
+        if value:
+            return value[:240]
+    return ""
+
+
+def flatten_nodes(structure: list[Any]) -> list[dict[str, Any]]:
+    out: list[dict[str, Any]] = []
+
+    def _walk(nodes: list[Any]) -> None:
+        for node in nodes:
+            if not isinstance(node, dict):
+                continue
+            out.append(node)
+            child = node.get("nodes")
+            if isinstance(child, list):
+                _walk(child)
+
+    _walk(structure)
+    return out
+
+
+@dataclass
+class BundleElement:
+    element_type: str
+    text: str
+    text_as_html: str | None
+    page_number: int
+    element_id: int | str
+
+
+def _parse_element(raw: dict[str, Any]) -> BundleElement | None:
+    if not isinstance(raw, dict):
+        return None
+    pn = raw.get("page_number")
+    if pn is None:
+        return None
+    eid = raw.get("element_id", 0)
+    etype = str(raw.get("element_type") or raw.get("type") or "Text")
+    text = str(raw.get("text") or "")
+    html = raw.get("text_as_html")
+    html_str = str(html) if isinstance(html, str) and html.strip() else None
+    return BundleElement(
+        element_type=etype,
+        text=text,
+        text_as_html=html_str,
+        page_number=max(1, int(pn)),
+        element_id=eid,
+    )
+
+
+def _bundle_elements(bundle: dict[str, Any]) -> list[BundleElement]:
+    out: list[BundleElement] = []
+    for raw in bundle.get("elements") or []:
+        el = _parse_element(raw)
+        if el is not None:
+            out.append(el)
+    return out
+
+
+def _element_body(el: BundleElement) -> str:
+    if el.element_type.strip().lower() == "table" and el.text_as_html:
+        return el.text_as_html
+    return el.text or ""
+
+
+def _elements_in_page_span(
+    elements: list[BundleElement],
+    start_page: int,
+    end_page: int,
+) -> list[BundleElement]:
+    lo = max(1, int(start_page))
+    hi = max(lo, int(end_page))
+    return [el for el in elements if lo <= el.page_number <= hi]
+
+
+def _needle_from_node(node: dict[str, Any]) -> str:
+    pseudo = str(node.get("split_document_anchor") or "").strip()
+    if pseudo:
+        return pseudo[:240]
+    return anchor_start_from_row(node)
+
+
+def _next_needle_from_structure(
+    structure: list[dict[str, Any]],
+    node: dict[str, Any],
+) -> str:
+    flat = flatten_nodes(structure)
+    target_id = node.get("node_id")
+    for i, row in enumerate(flat):
+        if row.get("node_id") != target_id:
+            continue
+        if i + 1 >= len(flat):
+            return ""
+        nxt = flat[i + 1]
+        needle = anchor_start_from_row(nxt)
+        if needle:
+            return needle
+        title = str(nxt.get("title") or "").strip()
+        return title[:240] if title else ""
+    return ""
+
+
+def _span_blob_with_offsets(
+    span_els: list[BundleElement],
+) -> tuple[str, list[tuple[BundleElement, int, int]]]:
+    parts: list[str] = []
+    offsets: list[tuple[BundleElement, int, int]] = []
+    pos = 0
+    for el in span_els:
+        body = _element_body(el)
+        if not body:
+            continue
+        start = pos
+        parts.append(body)
+        pos += len(body)
+        parts.append("\n\n")
+        pos += 2
+        offsets.append((el, start, pos - 2))
+    return "".join(parts), offsets
+
+
+def _slice_section_range(
+    blob: str,
+    offsets: list[tuple[BundleElement, int, int]],
+    *,
+    anchor: str,
+    pseudo_node: dict[str, Any],
+    next_needle: str,
+) -> tuple[int, int]:
+    del pseudo_node  # reserved for future anchor refinement
+    if not blob.strip():
+        return 0, 0
+    start = 0
+    if anchor:
+        idx = blob.find(anchor)
+        if idx >= 0:
+            start = idx
+    end = len(blob)
+    if next_needle:
+        idx = blob.find(next_needle, start + max(1, len(anchor)))
+        if idx >= 0:
+            end = idx
+    if not offsets:
+        return start, end
+    return start, end

chunksmith_agent/index_builder.py

@@ -0,0 +1,325 @@
+"""Build ``DocumentIndex`` from a multimodal pipeline result or saved JSON."""
+
+from __future__ import annotations
+
+import base64
+from pathlib import Path
+from typing import Any
+
+from chunksmith_agent.index_context import flatten_structure
+from chunksmith_agent.models import DocumentIndex
+from chunksmith_agent.element_retrieval import (
+    _bundle_elements,
+    _element_body,
+    _elements_in_page_span,
+    _needle_from_node,
+    _next_needle_from_structure,
+    _slice_section_range,
+    _span_blob_with_offsets,
+)
+
+_TABLE_TYPES = frozenset({"table"})
+_IMAGE_TYPES = frozenset({"image", "figure", "chart"})
+
+
+def _node_page_span(node: dict[str, Any]) -> tuple[int, int] | None:
+    si = node.get("start_index")
+    ei = node.get("end_index")
+    if si is None or ei is None:
+        pi = node.get("physical_index")
+        if pi is not None:
+            p = int(pi)
+            return p, p
+        return None
+    return max(1, int(si)), max(int(si), int(ei))
+
+
+def _load_png_b64(path: Path, *, max_bytes: int = 4_000_000) -> str | None:
+    try:
+        raw = path.read_bytes()
+        if len(raw) > max_bytes:
+            return None
+        return base64.standard_b64encode(raw).decode("ascii")
+    except OSError:
+        return None
+
+
+def _media_for_node(
+    node: dict[str, Any],
+    *,
+    structure: list[dict[str, Any]],
+    elements: list[Any],
+    image_dir: Path | None,
+    png_cycle: list[Path],
+) -> dict[str, Any]:
+    span = _node_page_span(node)
+    text = str(node.get("text") or "").strip()
+    tables: list[dict[str, Any]] = []
+    images: list[dict[str, Any]] = []
+
+    if span and elements:
+        start_p, end_p = span
+        span_els = _elements_in_page_span(elements, start_p, end_p)
+        blob, offsets = _span_blob_with_offsets(span_els)
+        if blob.strip():
+            needle = _needle_from_node(node)
+            nxt = _next_needle_from_structure(structure, node)
+            pseudo = {
+                "split_document_anchor": needle,
+                "title": node.get("title") or "",
+            }
+            try:
+                sc, ec = _slice_section_range(
+                    blob,
+                    offsets,
+                    anchor=needle,
+                    pseudo_node=pseudo,
+                    next_needle=nxt,
+                )
+                section_els = [el for el, s, e in offsets if s < ec and e > sc]
+            except Exception:
+                section_els = span_els
+        else:
+            section_els = span_els
+
+        if not text:
+            parts = [_element_body(el).strip() for el in section_els if _element_body(el).strip()]
+            text = "\n\n".join(parts)
+
+        rank = 0
+        for el in section_els:
+            t = el.element_type.strip().lower()
+            if t in _TABLE_TYPES and el.text_as_html:
+                tables.append(
+                    {
+                        "page_number": el.page_number,
+                        "html": el.text_as_html,
+                        "element_id": el.element_id,
+                    }
+                )
+            if t in _IMAGE_TYPES:
+                rank += 1
+                row: dict[str, Any] = {
+                    "page_number": el.page_number,
+                    "rank": rank,
+                    "element_id": el.element_id,
+                    "summary": (el.text or "")[:500],
+                }
+                if png_cycle:
+                    p = png_cycle[(rank - 1) % len(png_cycle)]
+                    row["image_path"] = str(p.resolve())
+                    b64 = _load_png_b64(p)
+                    if b64:
+                        row["image_base64"] = b64
+                images.append(row)
+
+    return {"text": text, "tables": tables, "images": images}
+
+
+def build_document_index(multimodal_out: dict[str, Any]) -> DocumentIndex:
+    """
+    Build searchable index from :func:`build_multimodal_index` / ``ChunkSmith_MultimodalIndexer_PDF`` output.
+    """
+    structure = multimodal_out.get("structure")
+    if not isinstance(structure, list):
+        structure = []
+
+    bundle = multimodal_out.get("canonical_bundle")
+    if not isinstance(bundle, dict):
+        bundle = None
+
+    image_dir: Path | None = None
+    if bundle:
+        pi = bundle.get("path_image")
+        if isinstance(pi, str) and pi.strip():
+            image_dir = Path(pi)
+    if image_dir is None:
+        ar = multimodal_out.get("artifact_root")
+        if isinstance(ar, str):
+            cand = Path(ar) / "image"
+            if cand.is_dir():
+                image_dir = cand
+
+    png_cycle: list[Path] = []
+    if image_dir and image_dir.is_dir():
+        png_cycle = sorted(image_dir.glob("*.png"))
+
+    elements = _bundle_elements(bundle) if bundle else []
+
+    media_by_node: dict[str, dict[str, Any]] = {}
+    for node in flatten_structure(structure):
+        nid = node.get("node_id")
+        if nid is None:
+            continue
+        key = str(nid).strip()
+        if not key:
+            continue
+        media_by_node[key] = _media_for_node(
+            node,
+            structure=structure,
+            elements=elements,
+            image_dir=image_dir,
+            png_cycle=png_cycle,
+        )
+
+    cf = None
+    if bundle and isinstance(bundle.get("coded_formate"), str):
+        cf = bundle["coded_formate"]
+    elif isinstance(multimodal_out.get("compressed_tree_string"), str):
+        cf = multimodal_out.get("compressed_tree_string")
+
+    return DocumentIndex(
+        doc_name=str(multimodal_out.get("doc_name") or "document"),
+        structure=structure,
+        media_by_node=media_by_node,
+        canonical_bundle=bundle,
+        coded_formate=cf,
+        image_dir=str(image_dir.resolve()) if image_dir else None,
+    )
+
+
+def _page_outline_summary(bundle: dict[str, Any], page: int) -> str:
+    """Short preview for page-level nodes so the picker can match topics."""
+    parts: list[str] = []
+    has_table = False
+    has_figure = False
+    for el in bundle.get("elements") or []:
+        if not isinstance(el, dict):
+            continue
+        if el.get("page_number") is None or int(el["page_number"]) != page:
+            continue
+        kind = str(el.get("element_type") or "").strip().lower()
+        if kind == "table":
+            has_table = True
+        if kind in _IMAGE_TYPES:
+            has_figure = True
+        text = str(el.get("text") or "").strip().replace("\n", " ")
+        if not text or len(text) < 4:
+            continue
+        if kind in ("title", "header", "subtitle", "narrativetext", "listitem") and len(parts) < 4:
+            snippet = text[:140] + ("…" if len(text) > 140 else "")
+            if snippet not in parts:
+                parts.append(snippet)
+    preview = " · ".join(parts)[:420].strip()
+    flags: list[str] = []
+    if has_table:
+        flags.append("table")
+    if has_figure:
+        flags.append("figure")
+    if flags:
+        preview = f"{preview} [{' '.join(flags)}]" if preview else f"[{' '.join(flags)}]"
+    return preview or f"Page {page}"
+
+
+def structure_from_canonical_bundle(bundle: dict[str, Any]) -> list[dict[str, Any]]:
+    """Build a page-level outline when no ``*_pageindex.json`` exists."""
+    pages: set[int] = set()
+    for el in bundle.get("elements") or []:
+        if not isinstance(el, dict):
+            continue
+        pn = el.get("page_number")
+        if pn is not None:
+            pages.add(max(1, int(pn)))
+    if not pages:
+        return [
+            {
+                "node_id": "0000",
+                "title": "Document",
+                "start_index": 1,
+                "end_index": 1,
+                "summary": "Full document (no page numbers in bundle).",
+            }
+        ]
+    return [
+        {
+            "node_id": f"{i:04d}",
+            "title": f"Page {p}",
+            "start_index": p,
+            "end_index": p,
+            "summary": _page_outline_summary(bundle, p),
+        }
+        for i, p in enumerate(sorted(pages))
+    ]
+
+
+def _artifact_root_for_saved_bundle(
+    canonical_bundle_path: Path,
+    bundle: dict[str, Any],
+) -> Path | None:
+    ar = bundle.get("artifact_root")
+    if isinstance(ar, str) and ar.strip():
+        root = Path(ar)
+        if root.is_dir():
+            return root.resolve()
+    parent = canonical_bundle_path.parent
+    if parent.name == "json":
+        root = parent.parent
+        if (root / "image").is_dir() or (root / "json").is_dir():
+            return root.resolve()
+    return None
+
+
+def _bundle_with_resolved_image_dir(
+    bundle: dict[str, Any],
+    artifact_root: Path | None,
+) -> dict[str, Any]:
+    if artifact_root is None:
+        return bundle
+    image_dir = artifact_root / "image"
+    if not image_dir.is_dir():
+        return bundle
+    pi = bundle.get("path_image")
+    if isinstance(pi, str) and pi.strip() and Path(pi).is_dir():
+        return bundle
+    out = dict(bundle)
+    out["path_image"] = str(image_dir.resolve())
+    out.setdefault("artifact_root", str(artifact_root.resolve()))
+    return out
+
+
+def build_document_index_from_saved(
+    *,
+    pageindex_path: Path | None = None,
+    canonical_bundle_path: Path | None = None,
+    artifact_root: Path | None = None,
+) -> DocumentIndex:
+    """Load index from saved JSON (pageindex, outline JSON, or canonical bundle only)."""
+    import json
+
+    if pageindex_path is None and canonical_bundle_path is None:
+        raise ValueError("pageindex_path or canonical_bundle_path is required")
+
+    if pageindex_path is not None and pageindex_path.is_file():
+        raw = json.loads(pageindex_path.read_text(encoding="utf-8"))
+        if isinstance(raw.get("elements"), list) and not isinstance(raw.get("structure"), list):
+            return build_document_index(raw)
+        if canonical_bundle_path and canonical_bundle_path.is_file():
+            cb = json.loads(canonical_bundle_path.read_text(encoding="utf-8"))
+            if isinstance(cb, dict):
+                ar = artifact_root or _artifact_root_for_saved_bundle(canonical_bundle_path, cb)
+                cb = _bundle_with_resolved_image_dir(cb, ar)
+                raw["canonical_bundle"] = cb
+                if ar is not None:
+                    raw["artifact_root"] = str(ar)
+        return build_document_index(raw)
+
+    if canonical_bundle_path is None or not canonical_bundle_path.is_file():
+        raise FileNotFoundError(f"Index file not found: {pageindex_path}")
+
+    bundle = json.loads(canonical_bundle_path.read_text(encoding="utf-8"))
+    if not isinstance(bundle, dict) or not isinstance(bundle.get("elements"), list):
+        raise ValueError(f"Not a canonical bundle: {canonical_bundle_path}")
+
+    doc_name = str(bundle.get("doc_name") or canonical_bundle_path.stem)
+    if doc_name.endswith("_canonical_bundle"):
+        doc_name = doc_name[: -len("_canonical_bundle")]
+    ar = artifact_root or _artifact_root_for_saved_bundle(canonical_bundle_path, bundle)
+    bundle = _bundle_with_resolved_image_dir(bundle, ar)
+    raw: dict[str, Any] = {
+        "doc_name": doc_name,
+        "structure": structure_from_canonical_bundle(bundle),
+        "canonical_bundle": bundle,
+    }
+    if ar is not None:
+        raw["artifact_root"] = str(ar)
+    return build_document_index(raw)