REVIEW: Citation display helpers (#520)

pageindex/__init__.py

@@ -9,6 +9,7 @@
 
 if _TYPE_CHECKING:
     from .flash import page_index_flash
+    from .imaging import highlight_region
     from .page_index_classic import page_index, page_index_main
     from .page_index_md import md_to_tree
     from .tree_optimize import optimize_tree
@@ -19,16 +20,17 @@
     "IndexConfig", "CloudIndexConfig", "LocalIndexConfig", "ChatConfig",
     "ChatProcessOptions", "ChatStream",
     "page_index", "page_index_main", "page_index_flash",
-    "optimize_tree", "md_to_tree",
+    "optimize_tree", "md_to_tree", "highlight_region",
 ]
 
 _LAZY = {
+    "highlight_region": ".imaging",
     "page_index_flash": ".flash",
     "optimize_tree": ".tree_optimize",
     "md_to_tree": ".page_index_md",
 }
 _SUBMODULES = {"agent_tools", "chat_stream", "client", "cloud_api", "errors",
-               "flash", "integrations", "local_api", "local_chat",
+               "flash", "imaging", "integrations", "local_api", "local_chat",
                "local_store", "mcp_bridge", "page_index_classic",
                "page_index_md", "tree_optimize", "types", "utils"}
 

pageindex/client.py

@@ -48,38 +48,43 @@ def _parse_pages(pages: str) -> list[int]:
 # The two citation tag formats PageIndex chat writes and renders.
 _OLD_CITATION_RE = re.compile(
     r"<doc=([^;<>]+);page=(\d+)(?:;block(?:_id)?=([^;<>]+))?>")
-_CITE_TAG_RE = re.compile(r"<cite\s([^<>]*)>")
+_CITE_TAG_RE = re.compile(r"<cite\s([^<>]*)>(?:(?P<inner>[^<>]*)</cite>)?")
 _CITE_ATTR_RE = re.compile(r"""\b(\w+)=(["'])(.*?)\2""", re.S)
 
 
+def _citation_key(m: re.Match) -> Optional[tuple[str, int, Optional[str]]]:
+    """(document, page, block_id) of one matched tag, or None when it names
+    no document or no positive page."""
+    if m.re is _OLD_CITATION_RE:
+        doc, page_str, block_id = m.group(1), m.group(2), m.group(3)
+    else:
+        attrs = {name: value for name, _, value in
+                 _CITE_ATTR_RE.findall(m.group(1))}
+        doc, page_str, block_id = (attrs.get("doc", ""), attrs.get("page", ""),
+                                   attrs.get("block"))
+    doc = doc.strip()
+    block_id = (block_id or "").strip() or None
+    try:
+        page = int(page_str.split("-")[0])
+    except ValueError:
+        return None
+    return (doc, page, block_id) if doc and page > 0 else None
+
+
 def _parse_citations(text: str) -> list[dict[str, Any]]:
     """``<doc=…;page=…;block=…>`` tags (the managed chat's format), then
     ``<cite doc= page= block=/>`` tags; deduplicated, ``block_id`` only
     when the tag carries one."""
     found: list[dict[str, Any]] = []
     seen: set[tuple[str, int, Optional[str]]] = set()
-
-    def add(doc: str, page_str: str, block_id: Optional[str]) -> None:
-        doc = doc.strip()
-        block_id = (block_id or "").strip() or None
-        try:
-            page = int(page_str.split("-")[0])
-        except ValueError:
-            return
-        key = (doc, page, block_id)
-        if doc and page > 0 and key not in seen:
+    for m in [*_OLD_CITATION_RE.finditer(text), *_CITE_TAG_RE.finditer(text)]:
+        key = _citation_key(m)
+        if key and key not in seen:
             seen.add(key)
-            entry: dict[str, Any] = {"document": doc, "page": page}
-            if block_id:
-                entry["block_id"] = block_id
+            entry: dict[str, Any] = {"document": key[0], "page": key[1]}
+            if key[2]:
+                entry["block_id"] = key[2]
             found.append(entry)
-
-    for m in _OLD_CITATION_RE.finditer(text):
-        add(m.group(1), m.group(2), m.group(3))
-    for m in _CITE_TAG_RE.finditer(text):
-        attrs = {name: value for name, _, value in
-                 _CITE_ATTR_RE.findall(m.group(1))}
-        add(attrs.get("doc", ""), attrs.get("page", ""), attrs.get("block"))
     return found
 
 
@@ -781,6 +786,42 @@ def get_block(self, doc_id: str, block_id: str) -> dict[str, Any]:
             "blocks. Create the client with an api_key to look up blocks."
         ).get_block(doc_id=doc_id, block_id=block_id)
 
+    def get_page_image(self, doc_id: str, page: int) -> str:
+        """
+        A short-lived URL to one page, rendered as a JPEG. Cloud-only:
+        local mode renders no page images.
+
+        Args:
+            doc_id (str): Document ID.
+            page (int): 1-based page number.
+
+        Returns:
+            str: The URL. Fetch the bytes with ``requests.get(url).content``,
+            or pass it to a vision model that takes image URLs.
+        """
+        return self._require_cloud(
+            "get_page_image is cloud-only — local mode has no page-image "
+            "rendering. Create the client with an api_key to get page images."
+        ).get_page_image(doc_id=doc_id, page=page)
+
+    def get_document_image(self, doc_id: str, img_id: str) -> str:
+        """
+        A short-lived URL to an image OCR extracted from the document.
+        Cloud-only: local mode stores no images.
+
+        Args:
+            doc_id (str): Document ID.
+            img_id (str): Image ID as page content carries it,
+                e.g. ``"img-7.jpeg"``.
+
+        Returns:
+            str: The URL, as ``get_page_image`` returns one.
+        """
+        return self._require_cloud(
+            "get_document_image is cloud-only — local mode has no embedded "
+            "image storage. Create the client with an api_key."
+        ).get_document_image(doc_id=doc_id, img_id=img_id)
+
     # ---------- TREE GENERATION ----------
 
     def get_tree(self, doc_id: str, node_summary: bool = False,
@@ -1720,11 +1761,14 @@ def get_document(self, doc_id: str) -> dict[str, Any]:
 
     def get_document_id(self, name: str) -> str:
         """
-        Look up a document's ID by its name. Useful for resolving
-        citation doc names (from ``<cite doc="…">``) to IDs.
+        Look up a document's ID by its name or path. A path like
+        ``"Research/Papers/attention.pdf"`` is accepted: the folder
+        part is stripped because document names are unique across
+        the library.
 
         Raises PageIndexAPIError if no document with that name exists.
         """
+        name = name.rsplit("/", 1)[-1] if "/" in name else name
         result = self._api.list_documents(limit=1, name=name)
         docs = result.get("documents", [])
         if docs:
@@ -2177,7 +2221,7 @@ def citation_prompt(self, format: str = "cite") -> str:
         from .agent_tools import fetch_citation_prompt
         return fetch_citation_prompt(self, format or "cite")
 
-    def resolve_citations(
+    def get_citations(
         self,
         answer: str,
         doc_id: Optional[Union[str, list[str]]] = None,
@@ -2259,6 +2303,51 @@ def resolve_citations(
             resolved.append(entry)
         return resolved
 
+    def resolve_citations(
+        self,
+        answer: str,
+        doc_id: Optional[Union[str, list[str]]] = None,
+    ) -> dict[str, Any]:
+        """
+        Display-ready citations: the answer text with citation tags
+        replaced by numbered markdown links, and each citation's full
+        data from ``get_citations()`` plus an anchor and index.
+
+        Tags (``<cite doc= page= block=/>`` and ``<doc=…;page=…>``)
+        become ``[[1]](#pageindex-citation-01)``, one number per distinct
+        citation, so a repeated citation reuses its number. The host
+        renders the anchor targets from the ``anchor`` field.
+
+        Args:
+            answer (str): The answer text, tags included.
+            doc_id (str | list[str], optional): As in ``get_citations()``.
+
+        Returns:
+            dict: ``{'answer': str, 'citations': list}`` where each
+            citation carries ``'anchor'``, ``'index'`` and the fields
+            ``get_citations()`` returns (``'document'``, ``'doc_id'``,
+            ``'page'``, and for block-level citations ``'block_id'`` plus,
+            when the block could be read, ``'bbox'``, ``'block_type'``,
+            ``'text'``).
+        """
+        entries = self.get_citations(answer, doc_id=doc_id)
+        index: dict[Any, int] = {
+            (c["document"], c["page"], c.get("block_id")): i
+            for i, c in enumerate(_parse_citations(answer), 1)}
+
+        def link(m: re.Match) -> str:
+            i = index.get(_citation_key(m))
+            if not i:
+                return m.group(0)
+            return (f"[[{i}]](#pageindex-citation-{i:02d})"
+                    f"{m.groupdict().get('inner') or ''}")
+
+        return {
+            "answer": _CITE_TAG_RE.sub(link, _OLD_CITATION_RE.sub(link, answer)),
+            "citations": [{"anchor": f"pageindex-citation-{i:02d}", "index": i,
+                           **entry} for i, entry in enumerate(entries, 1)],
+        }
+
     def folder_context(self, folder_id: str) -> str:
         """
         Folder targeting text for the first user message, placed as
@@ -2306,6 +2395,62 @@ def list_folders(self, parent_folder_id: Optional[str] = None) -> dict[str, Any]
             parent_folder_id=parent_folder_id,
         )
 
+    # ---------- PATH HELPERS ----------
+
+    def _folder_paths(self) -> dict[str, str]:
+        folders = {f["id"]: f for f in self.list_folders()["folders"]}
+        paths = {}
+        for folder_id in folders:
+            names, current = [], folder_id
+            while current in folders and len(names) < len(folders):
+                names.append(folders[current]["name"])
+                current = folders[current].get("parent_folder_id")
+            paths[folder_id] = "/".join(reversed(names))
+        return paths
+
+    def get_document_path(self, doc_id: str) -> str:
+        """
+        A document's path: its folder's path and its name, e.g.
+        ``"Research/Papers/attention.pdf"``. Just the name when the
+        document sits outside the API's folders, as at the library root
+        and for every local document.
+        """
+        doc = self.get_document(doc_id)
+        folder = doc.get("folderId") and self._folder_paths().get(doc["folderId"])
+        return f"{folder}/{doc['name']}" if folder else doc["name"]
+
+    def get_folder_path(self, folder_id: str) -> str:
+        """
+        A cloud folder's path: its ancestors' names and its own, root
+        first, e.g. ``"Research/Papers"``. Cloud-only. Raises
+        PageIndexAPIError if the folder does not exist.
+        """
+        self._require_cloud(
+            "get_folder_path is cloud-only — folders are not supported in "
+            "local mode. Create the client with an api_key.")
+        path = self._folder_paths().get(folder_id)
+        if path is None:
+            raise PageIndexAPIError(f"Folder {folder_id!r} not found.")
+        return path
+
+    def get_folder_id(self, path: str) -> str:
+        """
+        The ID of the cloud folder at ``path``, written as
+        ``get_folder_path`` writes it, e.g. ``"Research/Papers"``.
+        Cloud-only. Raises PageIndexAPIError if no folder, or more than
+        one, has that path.
+        """
+        self._require_cloud(
+            "get_folder_id is cloud-only — folders are not supported in "
+            "local mode. Create the client with an api_key.")
+        wanted = path.strip("/") if isinstance(path, str) else None
+        ids = [fid for fid, p in self._folder_paths().items() if p == wanted]
+        if len(ids) != 1:
+            raise PageIndexAPIError(
+                f"{path!r} names {len(ids)} folders ({', '.join(ids)})."
+                if ids else f"No folder at path {path!r}.")
+        return ids[0]
+
     def _require_cloud(self, message: str):
         from .cloud_api import CloudAPI
         if not isinstance(self._api, CloudAPI):

pageindex/cloud_api.py

@@ -137,6 +137,56 @@ def get_block(self, doc_id: str, block_id: str) -> Dict[str, Any]:
                 status_code=response.status_code)
         return response.json()
 
+    def get_page_image(self, doc_id: str, page: int) -> str:
+        """Presigned URL for a rendered page image.
+
+        Args:
+            doc_id (str): Document ID.
+            page (int): 1-based page number.
+
+        Returns:
+            str: A presigned URL to the page image (JPEG).
+        """
+        response = requests.get(
+            f"{self.BASE_URL}/doc/s3/{_enc(doc_id)}/images",
+            headers=self._headers(),
+            params={"start": page, "end": page},
+            timeout=30
+        )
+        if response.status_code != 200:
+            raise PageIndexAPIError(
+                f"Failed to get page image: {response.text}",
+                status_code=response.status_code)
+        for img in response.json().get("images") or []:
+            if img.get("page") == page and img.get("url"):
+                return img["url"]
+        raise PageIndexAPIError(f"No image URL returned for page {page}.")
+
+    def get_document_image(self, doc_id: str, img_id: str) -> str:
+        """Presigned URL for an embedded image extracted during OCR.
+
+        Args:
+            doc_id (str): Document ID.
+            img_id (str): Image ID as page content carries it,
+                e.g. ``"img-7.jpeg"``.
+
+        Returns:
+            str: A presigned URL to the image.
+        """
+        response = requests.get(
+            f"{self.BASE_URL}/doc/{_enc(doc_id)}/image/{_enc(img_id)}/",
+            headers=self._headers(),
+            timeout=30
+        )
+        if response.status_code != 200:
+            raise PageIndexAPIError(
+                f"Failed to get document image: {response.text}",
+                status_code=response.status_code)
+        url = response.json().get("url")
+        if not url:
+            raise PageIndexAPIError(f"No image URL returned for {img_id!r}.")
+        return url
+
     # ---------- TREE GENERATION ----------
 
     def get_tree(self, doc_id: str, node_summary: bool = False,

pageindex/imaging.py

@@ -0,0 +1,42 @@
+"""Image helpers for citation highlighting."""
+from __future__ import annotations
+
+from io import BytesIO
+from typing import Union
+
+from PIL import Image as PILImage, ImageDraw
+
+
+def highlight_region(
+    image: Union[PILImage.Image, bytes],
+    bbox: list,
+    scale: int = 1000,
+) -> PILImage.Image:
+    """Draw a translucent highlight over a bounding-box region.
+
+    Args:
+        image: A PIL Image or raw image bytes (JPEG/PNG).
+        bbox: ``[x0, y0, x1, y1]`` in units of ``scale`` from the
+            top-left corner, as ``get_block()`` returns it.
+        scale: The coordinate space ``bbox`` lives in (default 1000,
+            matching the PageIndex block coordinate system).
+
+    Returns:
+        A new PIL Image with a yellow highlight and orange outline
+        over the region.
+    """
+    if isinstance(image, bytes):
+        image = PILImage.open(BytesIO(image))
+    x0, y0, x1, y1 = map(float, bbox)
+    if scale <= 0 or not (0 <= x0 < x1 <= scale and 0 <= y0 < y1 <= scale):
+        raise ValueError(f"Invalid bbox for scale {scale}: {bbox}")
+    page = image.convert("RGBA")
+    width, height = page.size
+    rect = (x0 / scale * width, y0 / scale * height,
+            x1 / scale * width, y1 / scale * height)
+    overlay = PILImage.new("RGBA", page.size, (0, 0, 0, 0))
+    ImageDraw.Draw(overlay).rectangle(
+        rect, fill=(255, 210, 0, 65), outline=(255, 140, 0, 255),
+        width=max(2, round(width / 400)),
+    )
+    return PILImage.alpha_composite(page, overlay).convert("RGB")

pyproject.toml

@@ -40,6 +40,7 @@ sortedcontainers = ">=2.4.0"
 regex = ">=2024.0.0"
 python-dotenv = ">=1.0.0"
 pyyaml = ">=6.0"
+Pillow = ">=9.0"
 # Older releases break string prompts with SDK MCP servers (#597, #780).
 claude-agent-sdk = { version = ">=0.1.53", optional = true }
 # Older releases execute a refusal turn's tool_use blocks.

requirements.txt

@@ -12,3 +12,4 @@ python-dotenv==1.2.2
 pyyaml==6.0.2
 regex>=2024.0.0
 sortedcontainers==2.4.0
+Pillow>=9.0