From 4456b929680a2792e341e339044514de6e5cd315 Mon Sep 17 00:00:00 2001
From: mountain <kose2livs@gmail.com>
Date: Thu, 9 Jul 2026 19:26:57 +0800
Subject: [PATCH] fix: preserve image metadata in cloud get_page_content

Cloud OCR page results carry an `images` list per page, but the page
reconstruction only kept `page` and `content`, dropping images for cloud
callers of collection.get_page_content(). The local backend preserves them
and the PageContent contract / SDK prompts expect them (so the downstream
UI can render figures). Pass `images` through, omitting it when empty to
mirror the local backend's shape. Verified against the real OCR endpoint:
per-page keys are page_index/markdown/images.
---
 pageindex/backend/cloud.py  |  7 ++++++-
 tests/test_cloud_backend.py | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py
index 69ce703..c57734c 100644
--- a/pageindex/backend/cloud.py
+++ b/pageindex/backend/cloud.py
@@ -247,7 +247,12 @@ class CloudBackend:
         if isinstance(all_pages, list):
             return [
                 {"page": p.get("page", p.get("page_index")),
-                 "content": p.get("content", p.get("markdown", ""))}
+                 "content": p.get("content", p.get("markdown", "")),
+                 # Cloud OCR pages carry an `images` list (empty on text-only
+                 # pages). Preserve it — omitting when empty, mirroring the local
+                 # backend — so cloud callers get the same PageContent shape and
+                 # the SDK-prompted UI can render figures.
+                 **({"images": p["images"]} if p.get("images") else {})}
                 for p in all_pages
                 if p.get("page", p.get("page_index")) in page_nums
             ]
diff --git a/tests/test_cloud_backend.py b/tests/test_cloud_backend.py
index 35b8629..e68fcf3 100644
--- a/tests/test_cloud_backend.py
+++ b/tests/test_cloud_backend.py
@@ -142,6 +142,24 @@ def test_get_document_include_text_warns(monkeypatch):
         backend.get_document("col", "d1", include_text=True)
 
 
+def test_get_page_content_preserves_images(monkeypatch):
+    # Cloud OCR pages carry an `images` list; get_page_content must pass it
+    # through (parity with the local backend and the documented PageContent
+    # shape) so the SDK-prompted UI can render figures — omitting it only when
+    # empty. Real API uses page_index/markdown/images keys.
+    backend = CloudBackend(api_key="pi-test")
+    imgs = [{"path": "fig1.png", "width": 640, "height": 480}]
+    monkeypatch.setattr(backend, "_doc_request", lambda *a, **k: {"result": [
+        {"page_index": 1, "markdown": "page one", "images": imgs},
+        {"page_index": 2, "markdown": "page two", "images": []},  # empty -> omitted
+    ]})
+    out = backend.get_page_content("col", "d1", "1,2")
+    assert out == [
+        {"page": 1, "content": "page one", "images": imgs},
+        {"page": 2, "content": "page two"},
+    ]
+
+
 # ── query_stream: terminal contract and error propagation ───────────────────
 
 def _collect_events(backend, **kwargs):