apunkt/PageIndex
1
0
Fork 0
mirror of https://github.com/VectifyAI/PageIndex.git synced 2026-07-15 21:11:05 +02:00
Code Issues Projects Releases Packages Wiki Activity
PageIndex/pageindex/backend/cloud.py
Ray 3ff04d501a fix: address max-effort review findings 
- clamp LLM-derived page indices in _get_text_of_pages and
  get_text_of_pdf_pages_with_labels; dedupe get_text_of_pdf_pages
- guard _normalize_tree and folders/documents iterations against
  explicit nulls in cloud API responses
- coerce cloud OCR page numbers to int before filtering in
  get_page_content
- folder cache: raise on missing folder id instead of caching None;
  stop caching name-not-found so later lookups can succeed
- defang doc_id in agent doc-context prompt
- append api-key hint to 401 errors (request, legacy and streaming
  paths)
- remove stale legacy JSON workspace sample data unreadable by the
  SQLite storage
2026-07-13 18:49:09 +08:00

523 lines
23 KiB
Python
Raw Blame History

# pageindex/backend/cloud.py
"""CloudBackend — connects to PageIndex cloud service (api.pageindex.ai).
API reference: https://github.com/VectifyAI/pageindex_sdk
"""
from __future__ import annotations
import json
import logging
import os
import re
import time
import urllib.parse
import requests
from typing import AsyncIterator
from ..cloud_api import API_BASE # single source of truth for the cloud base URL
from ..errors import AUTH_HINT, CloudAPIError, DocumentNotFoundError, PageIndexError
from ..events import QueryEvent
logger = logging.getLogger(__name__)
_INTERNAL_TOOLS = frozenset({"ToolSearch", "Read", "Grep", "Glob", "Bash", "Edit", "Write"})
def _as_int(value):
try:
return int(value)
except (TypeError, ValueError):
return None
class CloudBackend:
def __init__(self, api_key: str):
self._api_key = api_key
self._headers = {"api_key": api_key}
self._folder_id_cache: dict[str, str | None] = {}
self._folder_warning_shown = False
# ── HTTP helpers ──────────────────────────────────────────────────────
# Folder API statuses meaning "folders are not available on this account"
# (403: requires Max plan; 404: endpoint not exposed). Anything else is a
# real error and must propagate rather than silently degrade.
_FOLDER_UNAVAILABLE = (403, 404)
def _warn_folder_upgrade(self) -> None:
if not self._folder_warning_shown:
import warnings
warnings.warn(
"Folders (collections) are not available on this plan. "
"All documents are stored in a single global space — collection names are ignored. "
"Upgrade at https://dash.pageindex.ai/subscription",
UserWarning,
stacklevel=4,
)
self._folder_warning_shown = True
def _request(self, method: str, path: str, retries: int = 3, **kwargs) -> dict:
"""HTTP helper. ``retries`` caps total attempts — pass 1 for
non-idempotent, expensive calls (e.g. chat completions) where a
retry would redo the full server-side work."""
url = f"{API_BASE}{path}"
kwargs.setdefault("timeout", 30)
last_status: int | None = None
for attempt in range(retries):
if attempt and "files" in kwargs:
# Rewind file objects before a retry — the previous attempt
# consumed them, and re-sending without seek(0) would upload
# an empty multipart body.
for value in kwargs["files"].values():
fobj = value[1] if isinstance(value, tuple) else value
if hasattr(fobj, "seek"):
fobj.seek(0)
try:
resp = requests.request(method, url, headers=self._headers, **kwargs)
if resp.status_code in (429, 500, 502, 503):
last_status = resp.status_code
if attempt == retries - 1:
break
logger.warning("Cloud API %s %s returned %d, retrying...", method, path, resp.status_code)
time.sleep(2 ** attempt)
continue
if resp.status_code != 200:
body = resp.text[:500] if resp.text else ""
msg = f"Cloud API error {resp.status_code}: {body}"
if resp.status_code == 401:
msg += f" — {AUTH_HINT}"
raise CloudAPIError(msg, status_code=resp.status_code)
return resp.json() if resp.content else {}
except requests.RequestException as e:
if attempt == retries - 1:
raise CloudAPIError(f"Cloud API request failed: {e}") from e
time.sleep(2 ** attempt)
raise CloudAPIError(f"Cloud API {method} {path} failed after retries"
+ (f" (last status {last_status})" if last_status else ""),
status_code=last_status)
@staticmethod
def _validate_collection_name(name: str) -> None:
# .fullmatch() (not .match()): a $-anchored .match() would accept a
# trailing newline ("papers\n") because $ matches just before a final \n.
if not re.fullmatch(r'[a-zA-Z0-9_-]{1,128}', name):
raise PageIndexError(
f"Invalid collection name: {name!r}. "
"Must be 1-128 chars of [a-zA-Z0-9_-]."
)
@staticmethod
def _enc(value: str) -> str:
return urllib.parse.quote(value, safe="")
# ── Collection management (mapped to folders) ─────────────────────────
def _create_folder(self, name: str) -> str:
"""POST /folder/ and return the new folder id, never a falsy value."""
resp = self._request("POST", "/folder/", json={"name": name})
folder_id = resp.get("folder", {}).get("id")
if not folder_id:
raise PageIndexError(
f"Cloud API returned no folder id when creating {name!r} "
f"(response keys: {list(resp)})"
)
return folder_id
def create_collection(self, name: str) -> None:
self._validate_collection_name(name)
try:
self._folder_id_cache[name] = self._create_folder(name)
except CloudAPIError as e:
if e.status_code in self._FOLDER_UNAVAILABLE:
self._warn_folder_upgrade()
self._folder_id_cache[name] = None
else:
raise
def get_or_create_collection(self, name: str) -> None:
self._validate_collection_name(name)
try:
data = self._request("GET", "/folders/")
for folder in data.get("folders", []) or []:
if folder.get("name") == name:
self._folder_id_cache[name] = folder["id"]
return
self._folder_id_cache[name] = self._create_folder(name)
except CloudAPIError as e:
if e.status_code in self._FOLDER_UNAVAILABLE:
self._warn_folder_upgrade()
self._folder_id_cache[name] = None
else:
raise
def _get_folder_id(self, name: str) -> str | None:
"""Resolve collection name to folder ID. Returns None if folders not available.
Only "folders unavailable on this plan" (403/404) is cached as None —
transient errors (network, 5xx) propagate so a blip can't silently
drop documents into the global space forever.
"""
if name in self._folder_id_cache:
return self._folder_id_cache.get(name)
try:
data = self._request("GET", "/folders/")
except CloudAPIError as e:
if e.status_code in self._FOLDER_UNAVAILABLE:
self._warn_folder_upgrade()
self._folder_id_cache[name] = None
return None
raise
for folder in data.get("folders", []) or []:
if folder.get("name") == name:
self._folder_id_cache[name] = folder["id"]
return folder["id"]
return None
def list_collections(self) -> list[str]:
data = self._request("GET", "/folders/")
return [f["name"] for f in data.get("folders", []) or []]
def delete_collection(self, name: str) -> None:
folder_id = self._get_folder_id(name)
if folder_id:
self._request("DELETE", f"/folder/{self._enc(folder_id)}/")
# Drop the cached id so a later same-name op re-resolves instead of
# reusing the now-deleted folder_id. Only when it was a REAL id —
# if folder_id was falsy, the cache holds the "folders unavailable
# on this plan" None sentinel, which must survive so we don't
# re-issue a doomed GET /folders/ on the next call.
self._folder_id_cache.pop(name, None)
# ── Document management ───────────────────────────────────────────────
def add_document(self, collection: str, file_path: str) -> str:
folder_id = self._get_folder_id(collection)
data = {"if_retrieval": "true"}
if folder_id:
data["folder_id"] = folder_id
with open(file_path, "rb") as f:
resp = self._request("POST", "/doc/", files={"file": f}, data=data)
doc_id = resp["doc_id"]
# Poll until indexing completes. The cloud API signals readiness via
# status == "completed"; retrieval_ready is not a reliable indicator.
for _ in range(120): # 10 min max
tree_resp = self._request("GET", f"/doc/{self._enc(doc_id)}/", params={"type": "tree"})
status = tree_resp.get("status", "")
if status == "completed":
return doc_id
if status == "failed":
raise CloudAPIError(f"Document {doc_id} indexing failed")
time.sleep(5)
raise CloudAPIError(f"Document {doc_id} indexing timed out")
def _doc_request(self, doc_id: str, method: str, path: str, **kwargs) -> dict:
"""Doc-scoped request: maps HTTP 404 to DocumentNotFoundError for
parity with the local backend's error taxonomy."""
try:
return self._request(method, path, **kwargs)
except CloudAPIError as e:
if e.status_code == 404:
raise DocumentNotFoundError(f"Document {doc_id} not found") from e
raise
def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> dict:
if include_text:
import warnings
warnings.warn(
"include_text is not supported by the cloud backend; "
"returning the structure without node text. "
"Use get_page_content(doc_id, pages) to fetch content.",
UserWarning,
stacklevel=3,
)
resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/metadata/")
# Fetch structure in the same call via tree endpoint
tree_resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/",
params={"type": "tree", "summary": "true"})
raw_tree = tree_resp.get("tree", tree_resp.get("structure", tree_resp.get("result", [])))
return {
"doc_id": resp.get("id", doc_id),
"doc_name": resp.get("name", ""),
"doc_description": resp.get("description", ""),
"doc_type": "pdf",
"status": resp.get("status", ""),
"structure": self._normalize_tree(raw_tree),
}
def get_document_structure(self, collection: str, doc_id: str) -> list:
resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/",
params={"type": "tree", "summary": "true"})
raw_tree = resp.get("tree", resp.get("structure", resp.get("result", [])))
return self._normalize_tree(raw_tree)
def get_page_content(self, collection: str, doc_id: str, pages: str) -> list:
resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/",
params={"type": "ocr", "format": "page"})
# Filter to requested pages
from ..index.utils import parse_pages
page_nums = set(parse_pages(pages))
all_pages = resp.get("pages", resp.get("ocr", resp.get("result", [])))
if not isinstance(all_pages, list):
return []
result = []
for p in all_pages:
page = _as_int(p.get("page", p.get("page_index")))
if page not in page_nums:
continue
entry = {"page": page,
"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.
if p.get("images"):
entry["images"] = p["images"]
result.append(entry)
return result
@staticmethod
def _normalize_tree(nodes: list | None) -> list:
"""Normalize cloud tree nodes to match local schema."""
if not nodes:
return []
result = []
for node in nodes:
normalized = {
"title": node.get("title", ""),
"node_id": node.get("node_id", ""),
"summary": node.get("summary", node.get("prefix_summary", "")),
"start_index": node.get("start_index", node.get("page_index")),
"end_index": node.get("end_index", node.get("page_index")),
}
if "text" in node:
normalized["text"] = node["text"]
children = node.get("nodes", [])
if children:
normalized["nodes"] = CloudBackend._normalize_tree(children)
result.append(normalized)
return result
def list_documents(self, collection: str) -> list[dict]:
folder_id = self._get_folder_id(collection)
# The API caps `limit` at 100; paginate with `offset` until a short
# page comes back so collections with >100 docs aren't silently
# truncated (queries over the whole collection rely on this list).
page_size = 100
offset = 0
docs: list[dict] = []
while True:
params = {"limit": page_size, "offset": offset}
if folder_id:
params["folder_id"] = folder_id
data = self._request("GET", "/docs/", params=params)
batch = data.get("documents", []) or []
docs.extend(
{
"doc_id": d.get("id", ""),
"doc_name": d.get("name", ""),
"doc_description": d.get("description", ""),
"doc_type": "pdf",
}
for d in batch
)
if len(batch) < page_size:
return docs
offset += page_size
def delete_document(self, collection: str, doc_id: str) -> None:
self._doc_request(doc_id, "DELETE", f"/doc/{self._enc(doc_id)}/")
# ── Query (uses cloud chat/completions, no LLM key needed) ────────────
def query(self, collection: str, question: str,
doc_ids: str | list[str] | None = None) -> str:
"""Non-streaming query via cloud chat/completions."""
if isinstance(doc_ids, str):
doc_ids = [doc_ids]
elif doc_ids == []:
raise ValueError(
"doc_ids cannot be empty; pass None to query the whole collection"
)
doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection)
if not doc_id:
raise ValueError("collection has no documents to query")
# A non-streaming completion returns nothing until generation
# finishes, so it needs far more than the default 30s. retries=1:
# retrying this non-idempotent call would redo the full server-side
# retrieval + generation (and bill it) on every attempt.
resp = self._request("POST", "/chat/completions/", retries=1, timeout=300, json={
"messages": [{"role": "user", "content": question}],
"doc_id": doc_id,
"stream": False,
})
# Extract answer from response
choices = resp.get("choices", [])
if choices:
return choices[0].get("message", {}).get("content", "")
return resp.get("content", resp.get("answer", ""))
async def query_stream(self, collection: str, question: str,
doc_ids: str | list[str] | None = None) -> AsyncIterator[QueryEvent]:
"""Streaming query via cloud chat/completions SSE.
Events are yielded in real-time as they arrive from the server.
A background thread handles the blocking HTTP stream and pushes
events through an asyncio.Queue for true async streaming.
"""
import asyncio
import threading
if isinstance(doc_ids, str):
doc_ids = [doc_ids]
elif doc_ids == []:
raise ValueError(
"doc_ids cannot be empty; pass None to query the whole collection"
)
doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection)
if not doc_id:
raise ValueError("collection has no documents to query")
headers = self._headers
# Queue carries QueryEvent, an Exception to re-raise, or None (end).
queue: asyncio.Queue[QueryEvent | Exception | None] = asyncio.Queue()
loop = asyncio.get_running_loop()
# Set when the consumer stops early (break / GeneratorExit) so the
# background thread stops draining the SSE stream instead of pulling
# it to completion in the background.
stop = threading.Event()
resp_holder: dict[str, requests.Response] = {}
def _put(item: QueryEvent | Exception | None) -> None:
try:
loop.call_soon_threadsafe(queue.put_nowait, item)
except RuntimeError:
pass # event loop already closed; consumer is gone
def _stream():
"""Background thread: read SSE and push events to queue.
Everything — including the initial connect — runs inside try so a
failure can never die silently and leave the consumer awaiting a
sentinel that never arrives. Errors are forwarded as exceptions
(raised in the consumer), never disguised as answer events.
"""
resp = None
answer_parts: list[str] = []
try:
resp = requests.post(
f"{API_BASE}/chat/completions/",
headers=headers,
json={
"messages": [{"role": "user", "content": question}],
"doc_id": doc_id,
"stream": True,
"stream_metadata": True,
},
stream=True,
timeout=120,
)
resp_holder["resp"] = resp
# The consumer may have abandoned the stream while we were still
# blocked in requests.post() (its connect phase, before resp
# existed to close). Now that resp exists, bail immediately
# rather than reading/draining a stream nobody is listening to;
# the finally block closes resp and pushes the sentinel.
if stop.is_set():
return
if resp.status_code != 200:
body = resp.text[:500] if resp.text else ""
msg = f"Cloud streaming error {resp.status_code}: {body}"
if resp.status_code == 401:
msg += f" — {AUTH_HINT}"
raise CloudAPIError(msg, status_code=resp.status_code)
current_tool_name = None
current_tool_args: list[str] = []
for line in resp.iter_lines(decode_unicode=True):
if stop.is_set():
return # consumer abandoned the stream
if not line or not line.startswith("data: "):
continue
data_str = line[6:]
if data_str.strip() == "[DONE]":
break
try:
chunk = json.loads(data_str)
except json.JSONDecodeError:
continue
meta = chunk.get("block_metadata", {})
block_type = meta.get("type", "")
choices = chunk.get("choices", [])
delta = choices[0].get("delta", {}) if choices else {}
content = delta.get("content", "")
if block_type == "mcp_tool_use_start":
current_tool_name = meta.get("tool_name", "")
current_tool_args = []
elif block_type == "tool_use":
if content:
current_tool_args.append(content)
elif block_type == "tool_use_stop":
if current_tool_name and current_tool_name not in _INTERNAL_TOOLS:
args_str = "".join(current_tool_args)
_put(QueryEvent(type="tool_call", data={
"name": current_tool_name,
"args": args_str,
}))
current_tool_name = None
current_tool_args = []
elif block_type == "text" and content:
answer_parts.append(content)
_put(QueryEvent(type="answer_delta", data=content))
# Same terminal contract as the local backend: a final
# answer_done event carrying the full answer text.
_put(QueryEvent(type="answer_done", data="".join(answer_parts)))
except requests.RequestException as e:
_put(CloudAPIError(f"Cloud streaming request failed: {e}"))
except Exception as e:
_put(e)
finally:
if resp is not None:
resp.close()
_put(None) # sentinel
thread = threading.Thread(target=_stream, daemon=True)
thread.start()
try:
while True:
item = await queue.get()
if item is None:
break
if isinstance(item, Exception):
raise item
yield item
finally:
# On early break / GeneratorExit / raised error: tell the thread to
# stop and force-close the response so a read blocked mid-stream
# unblocks instead of draining the rest in the background.
stop.set()
resp = resp_holder.get("resp")
if resp is not None:
try:
resp.close()
except Exception:
# Best-effort: the response may already be closed/invalid
# during teardown; closing is just to unblock the thread.
logger.debug("Ignoring error closing streaming response during cleanup",
exc_info=True)
thread.join(timeout=5)
def _get_all_doc_ids(self, collection: str) -> list[str]:
"""Get all document IDs in a collection."""
docs = self.list_documents(collection)
return [d["doc_id"] for d in docs]
Reference in a new issue View git blame Copy permalink