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 91e016d2e4 fix: reject trailing newline in collection-name validation 
Python's $ anchor matches just before a final newline, so a $-anchored
re.match(r'^[a-zA-Z0-9_-]{1,128}$', name) accepted "papers\n". In local
mode get_or_create_collection() then hit SQLite's CHECK via
INSERT OR IGNORE, silently created no row, and returned a Collection that
failed later on add(). Switch all three duplicated validators (local,
cloud, sqlite backends) to re.fullmatch() so the whole string must match.
2026-07-09 19:12:56 +08:00

489 lines
21 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:
# .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_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. 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 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]
Reference in a new issue View git blame Copy permalink