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
mountain fe36e25773 fix: six P2 correctness/robustness cleanups from the SDK review 
- legacy call_llm: open AsyncOpenAI via `async with` so the client (and
  its HTTP connection pool) is closed instead of leaked.
- LocalBackend.add_document: fail fast with CollectionNotFoundError when
  the collection doesn't exist, before the expensive parse + LLM index
  (previously the missing FK only tripped at save time, after paying for
  the LLM work). Also raise builtin FileNotFoundError for a missing path
  instead of FileTypeError (which now means only "unsupported extension").
- Collection.query(doc_ids=None): the empty-collection guard now always
  runs — previously it was skipped once PAGEINDEX_EXPERIMENTAL_MULTIDOC
  was set. A single list_documents call serves both the guard and the
  multi-doc warning (no separate call just to decide whether to warn).
- CloudBackend.query_stream: on early consumer break / GeneratorExit,
  signal the background SSE thread to stop and force-close the response,
  so it no longer drains the whole stream in the background.

Adds regression tests for each (client closed, fail-fast on unknown
collection, FileNotFoundError, empty-check under the multidoc env flag,
single list call, early-break thread stop).

Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
2026-07-07 11:34:13 +08:00

474 lines
20 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 ..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)}/")
# ── 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)
# 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)
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:
pass
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