# 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 ..errors import CloudAPIError, DocumentNotFoundError, PageIndexError from ..events import QueryEvent logger = logging.getLogger(__name__) API_BASE = "https://api.pageindex.ai" _INTERNAL_TOOLS = frozenset({"ToolSearch", "Read", "Grep", "Glob", "Bash", "Edit", "Write"}) 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 "" raise CloudAPIError(f"Cloud API error {resp.status_code}: {body}", 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: if not re.match(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_collection(self, name: str) -> None: self._validate_collection_name(name) try: resp = self._request("POST", "/folder/", json={"name": name}) self._folder_id_cache[name] = resp.get("folder", {}).get("id") 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", []): if folder.get("name") == name: self._folder_id_cache[name] = folder["id"] return resp = self._request("POST", "/folder/", json={"name": name}) self._folder_id_cache[name] = resp.get("folder", {}).get("id") 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", []): if folder.get("name") == name: self._folder_id_cache[name] = folder["id"] return folder["id"] self._folder_id_cache[name] = None return None def list_collections(self) -> list[str]: data = self._request("GET", "/folders/") return [f["name"] for f in data.get("folders", [])] 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. 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 isinstance(all_pages, list): return [ {"page": p.get("page", p.get("page_index")), "content": p.get("content", p.get("markdown", ""))} for p in all_pages if p.get("page", p.get("page_index")) in page_nums ] return [] @staticmethod def _normalize_tree(nodes: list) -> list: """Normalize cloud tree nodes to match local schema.""" 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", []) 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 if resp.status_code != 200: body = resp.text[:500] if resp.text else "" raise CloudAPIError( f"Cloud streaming error {resp.status_code}: {body}", 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]