diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 83c82dc9c..26128127c 100644 --- a/pageindex/__init__.py +++ b/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"} diff --git a/pageindex/client.py b/pageindex/client.py index 1eef62487..1d64c622b 100644 --- a/pageindex/client.py +++ b/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"]+);page=(\d+)(?:;block(?:_id)?=([^;<>]+))?>") -_CITE_TAG_RE = re.compile(r"]*)>") +_CITE_TAG_RE = re.compile(r"]*)>(?:(?P[^<>]*))?") _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]]: """```` tags (the managed chat's format), then ```` 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 ````) 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 (```` and ````) + 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): diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index 3f3f2696e..d515b3177 100644 --- a/pageindex/cloud_api.py +++ b/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, diff --git a/pageindex/imaging.py b/pageindex/imaging.py new file mode 100644 index 000000000..0f0a9e7de --- /dev/null +++ b/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") diff --git a/pyproject.toml b/pyproject.toml index 646668ac5..f1f66e1cd 100644 --- a/pyproject.toml +++ b/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. diff --git a/requirements.txt b/requirements.txt index a066eac05..35d3da3fb 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/test_client.py b/tests/test_client.py index 87344ada1..a91a74a67 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -610,23 +610,50 @@ def test_get_page_content_span_bomb_rejected(local_client, indexed_doc): assert local_client.get_page_content(indexed_doc, "5-10004") == [] -def test_local_citations_resolve_to_pages_only(local_client, indexed_doc): +def test_local_get_citations(local_client, indexed_doc): """Local page content has no blocks: page citations resolve to the document id, get_block names the cloud-only exit, and a block citation keeps its entry without a bbox.""" answer = 'Apples and none .' - assert local_client.resolve_citations(answer) == [ + assert local_client.get_citations(answer) == [ {"document": "sample.pdf", "doc_id": indexed_doc, "page": 1}, {"document": "other.pdf", "doc_id": None, "page": 2}, ] with pytest.raises(PageIndexAPIError, match="get_block is cloud-only"): local_client.get_block(indexed_doc, "p1_text_1") - assert local_client.resolve_citations( + assert local_client.get_citations( '') == [ {"document": "sample.pdf", "doc_id": indexed_doc, "page": 1, "block_id": "p1_text_1"}] +def test_local_resolve_citations(local_client, indexed_doc): + """resolve_citations rewrites tags and enriches entries on local docs.""" + answer = 'Apples and none .' + assert local_client.resolve_citations(answer) == { + "answer": "Apples [[1]](#pageindex-citation-01) and none " + "[[2]](#pageindex-citation-02).", + "citations": [ + {"anchor": "pageindex-citation-01", "index": 1, + "document": "sample.pdf", "doc_id": indexed_doc, "page": 1}, + {"anchor": "pageindex-citation-02", "index": 2, + "document": "other.pdf", "doc_id": None, "page": 2}, + ], + } + + +def test_local_paths_and_images(local_client, indexed_doc): + """Local documents have no folders and no stored images: a document's + path is its name, and the folder and image lookups refuse.""" + assert local_client.get_document_path(indexed_doc) == "sample.pdf" + for call in (lambda: local_client.get_folder_path("f-1"), + lambda: local_client.get_folder_id("Research"), + lambda: local_client.get_page_image(indexed_doc, 1), + lambda: local_client.get_document_image(indexed_doc, "img-0.jpeg")): + with pytest.raises(PageIndexAPIError, match="is cloud-only"): + call() + + def test_submit_does_not_create_cwd_logs(local_client, sample_pdf, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) def fake_page_index_main(doc, opt=None, logger=None, page_list=None): @@ -1493,6 +1520,8 @@ def test_cloud_errors_carry_status_code(cloud, monkeypatch, sample_pdf): messages=[{"role": "user", "content": "q"}]), lambda: client.get_document("pi-1"), lambda: client.get_block("pi-1", "p1_text_1"), + lambda: client.get_page_image("pi-1", 1), + lambda: client.get_document_image("pi-1", "img-0.jpeg"), lambda: client.delete_document("pi-1"), lambda: client.list_documents(), lambda: client.create_folder("f"), @@ -1568,7 +1597,7 @@ def handler(method, url, kw): return handler -def test_resolve_citations_cloud(cloud, monkeypatch): +def test_get_citations_cloud(cloud, monkeypatch): """Names become ids from the library, block citations pick up the block's fields, a block the document lacks keeps a bbox-less entry, an unknown document keeps doc_id None.""" @@ -1583,14 +1612,14 @@ def test_resolve_citations_cloud(cloud, monkeypatch): answer = ('X ' 'Y ' 'Z W ') - assert client.resolve_citations(answer) == [ + assert client.get_citations(answer) == [ {"document": "a.pdf", **block}, {"document": "a.pdf", "doc_id": "pi-a", "page": 9, "block_id": "p9_text_9"}, {"document": "b.pdf", "doc_id": "pi-b", "page": 2}, {"document": "c.pdf", "doc_id": None, "page": 1}, ] - assert client.resolve_citations("no citations here") == [] + assert client.get_citations("no citations here") == [] # doc_id= skips the library listing: one metadata call per id. library = _library({"pi-a": ("a.pdf", {"p3_text_5": block})}) @@ -1599,20 +1628,20 @@ def no_listing(method, url, kw): return library(method, url, kw) _patch_requests(monkeypatch, no_listing) for scope in ("pi-a", ["pi-a", "pi-a"]): # a repeated id is not a collision - assert client.resolve_citations( + assert client.get_citations( '', doc_id=scope, ) == [{"document": "a.pdf", **block}] for not_text in (None, ['']): with pytest.raises(PageIndexAPIError, match="answer must be a str"): - client.resolve_citations(not_text) + client.get_citations(not_text) with pytest.raises(PageIndexAPIError, match="doc_id must be a string or a list"): - client.resolve_citations(answer, doc_id=5) + client.get_citations(answer, doc_id=5) with pytest.raises(PageIndexAPIError, match="doc_id is empty"): - client.resolve_citations(answer, doc_id=[]) + client.get_citations(answer, doc_id=[]) -def test_resolve_citations_quotes_and_tag_edges(cloud, monkeypatch): +def test_get_citations_quotes_and_tag_edges(cloud, monkeypatch): """A quoted name keeps its apostrophe, and neither field of a managed-chat tag runs past the tag: an unterminated name and an unterminated block both stop at the next tag instead of eating it.""" @@ -1621,17 +1650,17 @@ def test_resolve_citations_quotes_and_tag_edges(cloud, monkeypatch): "pi-b": ("b.pdf", {})})) answer = ("""Held and overall """ " is long, but revenue rose .") - assert client.resolve_citations(answer) == [ + assert client.get_citations(answer) == [ {"document": "b.pdf", "doc_id": "pi-b", "page": 7}, {"document": "Moody's Outlook.pdf", "doc_id": "pi-m", "page": 3}, ] # A block= that never closes must not swallow the citation after it. - assert client.resolve_citations( + assert client.get_citations( '' ) == [{"document": "b.pdf", "doc_id": "pi-b", "page": 9}] -def test_resolve_citations_lists_the_whole_library(cloud, monkeypatch): +def test_get_citations_lists_the_whole_library(cloud, monkeypatch): client, calls, fake = cloud seen_offsets = [] library = _library({f"pi-{i}": (f"{i}.pdf", {}) for i in range(150)}) @@ -1640,12 +1669,12 @@ def handler(method, url, kw): seen_offsets.append(kw["params"]["offset"]) return library(method, url, kw) _patch_requests(monkeypatch, handler) - assert client.resolve_citations('') == [ + assert client.get_citations('') == [ {"document": "149.pdf", "doc_id": "pi-149", "page": 1}] assert seen_offsets == [0, 100] -def test_resolve_citations_listing_survives_a_shifting_library(cloud, monkeypatch): +def test_get_citations_listing_survives_a_shifting_library(cloud, monkeypatch): """A listing without 'total' ends on the empty page, a document re-served after an upload shifted the window is still one document, and an entry missing 'name' or 'id' is skipped, not a KeyError.""" @@ -1660,24 +1689,24 @@ def handler(method, url, kw): entries.insert(0, {"id": "pi-new", "name": "new.pdf"}) return FakeResponse({"documents": entries[offset:offset + 100]}) _patch_requests(monkeypatch, handler) - assert client.resolve_citations('') == [ + assert client.get_citations('') == [ {"document": "99.pdf", "doc_id": "pi-99", "page": 1}] -def test_resolve_citations_refuses_a_shared_name(cloud, monkeypatch): +def test_get_citations_refuses_a_shared_name(cloud, monkeypatch): """Two documents under one cited name would silently pick a bbox from the wrong one: raise, name both ids, point at doc_id=.""" client, calls, fake = cloud _patch_requests(monkeypatch, _library({"pi-1": ("a.pdf", {}), "pi-2": ("a.pdf", {})})) with pytest.raises(PageIndexAPIError, match=r"pi-1, pi-2.*doc_id="): - client.resolve_citations('') - assert client.resolve_citations('', + client.get_citations('') + assert client.get_citations('', doc_id=["pi-2"]) == [ {"document": "a.pdf", "doc_id": "pi-2", "page": 1}] -def test_resolve_citations_keeps_blocks_it_cannot_read(cloud, monkeypatch): +def test_get_citations_keeps_blocks_it_cannot_read(cloud, monkeypatch): """A denied block (403) keeps its bbox-less entry like a missing one; a transport failure still propagates with its status.""" client, calls, fake = cloud @@ -1688,14 +1717,54 @@ def handler(method, url, kw): return _library({"pi-a": ("a.pdf", {})})(method, url, kw) _patch_requests(monkeypatch, handler) answer = '' - assert client.resolve_citations(answer) == [ + assert client.get_citations(answer) == [ {"document": "a.pdf", "doc_id": "pi-a", "page": 1, "block_id": "p1_text_1"}] status["code"] = 500 with pytest.raises(PageIndexAPIError, match="Failed to get block: boom") as err: - client.resolve_citations(answer) + client.get_citations(answer) assert err.value.status_code == 500 +def test_resolve_citations_rewrites_tags(cloud, monkeypatch): + """Each tag becomes the numbered link of its get_citations() entry, a + repeated citation reuses its number, and entries lead with anchor and + index. A block cited under the wrong page still gets its link: the + entry carries the block's real page, the number follows the tag.""" + client, calls, fake = cloud + block = {"doc_id": "pi-a", "page": 3, "block_id": "p3_text_5", + "bbox": [163, 398, 842, 589], "block_type": "text", + "text": "Apples."} + _patch_requests(monkeypatch, _library({ + "pi-a": ("a.pdf", {"p3_text_5": block}), + "pi-b": ("b.pdf", {}), + })) + answer = ('X ' + 'Y ' + 'Z ' + 'W ') + result = client.resolve_citations(answer) + assert result["answer"] == ( + "X [[1]](#pageindex-citation-01) Y [[2]](#pageindex-citation-02) " + "Z [[1]](#pageindex-citation-01) W [[3]](#pageindex-citation-03)") + assert result["citations"] == [ + {"anchor": "pageindex-citation-01", "index": 1, "document": "a.pdf", + **block}, + {"anchor": "pageindex-citation-02", "index": 2, "document": "b.pdf", + "doc_id": "pi-b", "page": 2}, + {"anchor": "pageindex-citation-03", "index": 3, "document": "a.pdf", + **block}, + ] + assert all(list(c)[:2] == ["anchor", "index"] for c in result["citations"]) + paired = client.resolve_citations( + "P quoted " + 'Q .') + assert paired["answer"] == ( + "P [[1]](#pageindex-citation-01)quoted " + "Q [[1]](#pageindex-citation-01).") + assert client.resolve_citations("no citations") == { + "answer": "no citations", "citations": []} + + def test_get_block_request_wiring(cloud): client, calls, fake = cloud fake.payload = {"doc_id": "pi/1", "page": 3, "block_id": "p3_text_5", @@ -1706,6 +1775,75 @@ def test_get_block_request_wiring(cloud): assert calls[-1]["timeout"] == 30 +def test_image_url_request_wiring(cloud): + client, calls, fake = cloud + fake.payload = {"doc_id": "pi/1", "images": [ + {"page": 3, "mime_type": "image/jpeg", "object_key": "k3", + "url": "https://signed/page-3.jpg", "expires_in": 3600}]} + assert client.get_page_image("pi/1", 3) == "https://signed/page-3.jpg" + assert calls[-1]["url"] == "https://api.pageindex.ai/doc/s3/pi%2F1/images" + assert calls[-1]["params"] == {"start": 3, "end": 3} + assert calls[-1]["headers"] == {"api_key": "secret"} + fake.payload["images"][0]["url"] = None + with pytest.raises(PageIndexAPIError, match="No image URL returned for page 3"): + client.get_page_image("pi/1", 3) + + fake.payload = {"doc_id": "pi/1", "img_id": "img-7.jpeg", + "url": "https://signed/img-7.jpeg", "expires_in": 3600} + assert client.get_document_image("pi/1", "img-7.jpeg") == "https://signed/img-7.jpeg" + assert calls[-1]["url"] == "https://api.pageindex.ai/doc/pi%2F1/image/img-7.jpeg/" + assert calls[-1]["headers"] == {"api_key": "secret"} + fake.payload = {} + with pytest.raises(PageIndexAPIError, match="No image URL returned for 'img-7.jpeg'"): + client.get_document_image("pi/1", "img-7.jpeg") + + +def test_folder_and_document_paths(cloud, monkeypatch): + """Paths are folder names root first; get_folder_id reads one back, a + document outside the API's folders is just its name, a path two + folders share raises instead of picking one, and a parent cycle ends + the walk instead of hanging it.""" + client, calls, fake = cloud + folders = [{"id": "f-r", "name": "Research", "parent_folder_id": None}, + {"id": "f-p", "name": "Papers", "parent_folder_id": "f-r"}] + doc_folders = {"pi-root": None, "pi-nested": "f-p", "pi-library": "f-lib"} + + def handler(method, url, kw): + if url.endswith("/folders/"): + return FakeResponse({"folders": folders, "total": len(folders)}) + m = re.fullmatch(r".*/doc/([^/]+)/metadata/", url) + if m: + doc_id = m.group(1) + return FakeResponse({"id": doc_id, "name": f"{doc_id}.pdf", + "folderId": doc_folders[doc_id]}) + if url.endswith("/docs/"): + name = kw.get("params", {}).get("name", "") + docs = [{"id": did, "name": f"{did}.pdf"} + for did in doc_folders if f"{did}.pdf" == name] + return FakeResponse({"documents": docs}) + return FakeResponse({}) + _patch_requests(monkeypatch, handler) + + assert client.get_folder_path("f-p") == "Research/Papers" + assert client.get_folder_id("Research/Papers") == "f-p" + assert client.get_folder_id("/Research/") == "f-r" + assert client.get_document_path("pi-nested") == "Research/Papers/pi-nested.pdf" + assert client.get_document_path("pi-root") == "pi-root.pdf" + assert client.get_document_path("pi-library") == "pi-library.pdf" + assert client.get_document_id("Research/Papers/pi-nested.pdf") == "pi-nested" + assert client.get_document_id("pi-root.pdf") == "pi-root" + with pytest.raises(PageIndexAPIError, match="Folder 'f-none' not found"): + client.get_folder_path("f-none") + with pytest.raises(PageIndexAPIError, match="No folder at path 'Research/X'"): + client.get_folder_id("Research/X") + folders.append({"id": "f-x", "name": "Research/Papers", "parent_folder_id": None}) + with pytest.raises(PageIndexAPIError, match=r"names 2 folders \(f-p, f-x\)"): + client.get_folder_id("Research/Papers") + folders[:] = [{"id": "a", "name": "A", "parent_folder_id": "b"}, + {"id": "b", "name": "B", "parent_folder_id": "a"}] + assert client.get_folder_path("a") == "B/A" + + def test_cloud_chat_stream_parsing(cloud, monkeypatch): client, calls, fake = cloud lines = [ diff --git a/tests/test_imaging.py b/tests/test_imaging.py new file mode 100644 index 000000000..1b98a531f --- /dev/null +++ b/tests/test_imaging.py @@ -0,0 +1,32 @@ +"""highlight_region self-check.""" +import pytest +from PIL import Image +from pageindex.imaging import highlight_region + + +def test_highlight_region_draws_on_image(): + img = Image.new("RGB", (500, 2000), "white") + result = highlight_region(img, [100, 200, 300, 400]) + assert isinstance(result, Image.Image) + assert result.size == (500, 2000) + assert result.getpixel((75, 700)) != (255, 255, 255) + assert result.getpixel((300, 300)) == (255, 255, 255) + + +def test_highlight_region_accepts_bytes(): + from io import BytesIO + img = Image.new("RGB", (500, 500), "white") + buf = BytesIO() + img.save(buf, "PNG") + result = highlight_region(buf.getvalue(), [0, 0, 500, 500], scale=500) + assert isinstance(result, Image.Image) + + +def test_highlight_region_invalid_bbox(): + img = Image.new("RGB", (100, 100), "white") + with pytest.raises(ValueError, match="Invalid bbox"): + highlight_region(img, [300, 200, 100, 400]) + with pytest.raises(ValueError, match="Invalid bbox"): + highlight_region(img, [0, 0, 1001, 500]) + with pytest.raises(ValueError, match="Invalid bbox"): + highlight_region(img, [0, 0, 500, 500], scale=0)