mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-24 21:41:04 +02:00
42d57b6f4c
The 0.2.x surface is the mainstream cloud API for now — Collection is an additive layer, and chat_completions has capabilities query() does not cover (temperature, citations, message history). PEP 702 deprecated + PendingDeprecationWarning told users these methods are going away (IDE strikethrough, pytest warning spam) when there is no removal plan. Docstrings keep neutral Collection API cross-references.
299 lines
12 KiB
Python
299 lines
12 KiB
Python
# pageindex/client.py
|
|
from __future__ import annotations
|
|
from pathlib import Path
|
|
from typing import Any, Iterator
|
|
|
|
from .cloud_api import API_BASE
|
|
from .collection import Collection
|
|
from .config import IndexConfig
|
|
from .errors import PageIndexAPIError
|
|
from .parser.protocol import DocumentParser
|
|
|
|
def _normalize_retrieve_model(model: str) -> str:
|
|
"""Preserve supported Agents SDK prefixes and route other provider paths via LiteLLM."""
|
|
passthrough_prefixes = ("litellm/", "openai/")
|
|
if not model or "/" not in model:
|
|
return model
|
|
if model.startswith(passthrough_prefixes):
|
|
return model
|
|
return f"litellm/{model}"
|
|
|
|
|
|
class PageIndexClient:
|
|
"""PageIndex client — supports both local and cloud modes.
|
|
|
|
Args:
|
|
api_key: PageIndex cloud API key. When provided, cloud mode is used
|
|
and local-only params (model, storage_path, index_config, …) are ignored.
|
|
model: LLM model for indexing (local mode only, default: gpt-4o-2024-11-20).
|
|
retrieve_model: LLM model for agent QA (local mode only, default: gpt-5.4).
|
|
storage_path: Directory for SQLite DB and files (local mode only, default: ./.pageindex).
|
|
storage: Custom StorageEngine instance (local mode only).
|
|
index_config: Advanced indexing parameters (local mode only, optional).
|
|
Pass an IndexConfig instance or a dict. Defaults are sensible for most use cases.
|
|
|
|
Usage:
|
|
# Local mode (auto-detected when no api_key)
|
|
client = PageIndexClient(model="gpt-5.4")
|
|
|
|
# Cloud mode (auto-detected when api_key provided)
|
|
client = PageIndexClient(api_key="your-api-key")
|
|
|
|
# Or use LocalClient / CloudClient for explicit mode selection
|
|
"""
|
|
|
|
BASE_URL = API_BASE # single source of truth lives in cloud_api
|
|
|
|
def __init__(self, api_key: str | None = None, model: str = None,
|
|
retrieve_model: str = None, storage_path: str = None,
|
|
storage=None, index_config: IndexConfig | dict = None):
|
|
# Track whether api_key was passed as empty string vs None — only
|
|
# affects the error message when legacy cloud methods are then called.
|
|
self._empty_api_key = api_key == ""
|
|
if self._empty_api_key:
|
|
import logging
|
|
logging.getLogger(__name__).warning(
|
|
"PageIndexClient received an empty api_key; falling back to local mode. "
|
|
"Pass api_key=None to silence this warning, or provide a real key for cloud mode."
|
|
)
|
|
api_key = None
|
|
if api_key is not None:
|
|
self._init_cloud(api_key)
|
|
else:
|
|
self._init_local(model, retrieve_model, storage_path, storage, index_config)
|
|
|
|
def _init_cloud(self, api_key: str):
|
|
from .backend.cloud import CloudBackend
|
|
from .cloud_api import LegacyCloudAPI
|
|
# Callables: re-read per request so post-construction BASE_URL /
|
|
# api_key reassignment (0.2.x patterns) still applies.
|
|
self.api_key = api_key
|
|
base_url = lambda: self.BASE_URL
|
|
api_key_ref = lambda: self.api_key
|
|
self._backend = CloudBackend(api_key=api_key_ref, base_url=base_url)
|
|
self._legacy_cloud_api = LegacyCloudAPI(api_key=api_key_ref, base_url=base_url)
|
|
|
|
def _init_local(self, model: str = None, retrieve_model: str = None,
|
|
storage_path: str = None, storage=None,
|
|
index_config: IndexConfig | dict = None):
|
|
self._legacy_cloud_api = None
|
|
|
|
# Build IndexConfig: merge model/retrieve_model with index_config
|
|
overrides = {}
|
|
if model:
|
|
overrides["model"] = model
|
|
if retrieve_model:
|
|
overrides["retrieve_model"] = retrieve_model
|
|
if isinstance(index_config, IndexConfig):
|
|
opt = index_config.model_copy(update=overrides)
|
|
elif isinstance(index_config, dict):
|
|
merged = {**index_config, **overrides} # explicit model/retrieve_model win
|
|
opt = IndexConfig(**merged)
|
|
else:
|
|
opt = IndexConfig(**overrides) if overrides else IndexConfig()
|
|
|
|
self.model = opt.model
|
|
self.retrieve_model = _normalize_retrieve_model(opt.retrieve_model or self.model)
|
|
|
|
storage_path = Path(storage_path or ".pageindex").resolve()
|
|
storage_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
from .storage.sqlite import SQLiteStorage
|
|
from .backend.local import LocalBackend
|
|
storage_engine = storage or SQLiteStorage(str(storage_path / "pageindex.db"))
|
|
self._backend = LocalBackend(
|
|
storage=storage_engine,
|
|
files_dir=str(storage_path / "files"),
|
|
model=opt.model,
|
|
retrieve_model=self.retrieve_model,
|
|
index_config=opt,
|
|
)
|
|
|
|
def collection(self, name: str = "default") -> Collection:
|
|
"""Get or create a collection. Defaults to 'default'."""
|
|
self._backend.get_or_create_collection(name)
|
|
return Collection(name=name, backend=self._backend)
|
|
|
|
def list_collections(self) -> list[str]:
|
|
return self._backend.list_collections()
|
|
|
|
def delete_collection(self, name: str) -> None:
|
|
self._backend.delete_collection(name)
|
|
|
|
def register_parser(self, parser: DocumentParser) -> None:
|
|
"""Register a custom document parser. Only available in local mode."""
|
|
from .backend.protocol import SupportsParserRegistration
|
|
if not isinstance(self._backend, SupportsParserRegistration):
|
|
from .errors import PageIndexError
|
|
raise PageIndexError("Custom parsers are not supported in cloud mode")
|
|
self._backend.register_parser(parser)
|
|
|
|
def _require_cloud_api(self):
|
|
if self._legacy_cloud_api is None:
|
|
from .errors import PageIndexAPIError
|
|
if getattr(self, "_empty_api_key", False):
|
|
raise PageIndexAPIError(
|
|
"Cannot call legacy SDK methods: api_key was an empty string, "
|
|
"so PageIndexClient fell back to local mode. Pass a real "
|
|
"PageIndex cloud API key, or migrate to the Collection API "
|
|
"(client.collection(...)) for local mode."
|
|
)
|
|
raise PageIndexAPIError(
|
|
"This method calls the PageIndex cloud API — create the client "
|
|
"with an api_key (get one at https://dash.pageindex.ai)."
|
|
)
|
|
return self._legacy_cloud_api
|
|
|
|
# ── pageindex 0.2.x cloud SDK surface (cloud mode only) ──
|
|
def submit_document(
|
|
self,
|
|
file_path: str,
|
|
mode: str | None = None,
|
|
beta_headers: list[str] | None = None,
|
|
folder_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``client.collection(...).add(path)``."""
|
|
return self._require_cloud_api().submit_document(
|
|
file_path=file_path,
|
|
mode=mode,
|
|
beta_headers=beta_headers,
|
|
folder_id=folder_id,
|
|
)
|
|
|
|
def get_ocr(self, doc_id: str, format: str = "page") -> dict[str, Any]:
|
|
"""Collection API equivalent: ``collection.get_page_content(doc_id, pages)``."""
|
|
return self._require_cloud_api().get_ocr(doc_id=doc_id, format=format)
|
|
|
|
def get_tree(self, doc_id: str, node_summary: bool = False) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``collection.get_document_structure(doc_id)``."""
|
|
return self._require_cloud_api().get_tree(doc_id=doc_id, node_summary=node_summary)
|
|
|
|
def is_retrieval_ready(self, doc_id: str) -> bool:
|
|
"""The Collection API (``collection.add``) handles readiness internally."""
|
|
return self._require_cloud_api().is_retrieval_ready(doc_id=doc_id)
|
|
|
|
def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``collection.query(question, doc_ids=[doc_id])``."""
|
|
return self._require_cloud_api().submit_query(
|
|
doc_id=doc_id,
|
|
query=query,
|
|
thinking=thinking,
|
|
)
|
|
|
|
def get_retrieval(self, retrieval_id: str) -> dict[str, Any]:
|
|
"""The Collection API (``collection.query``) returns answers synchronously."""
|
|
return self._require_cloud_api().get_retrieval(retrieval_id=retrieval_id)
|
|
|
|
def chat_completions(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
stream: bool = False,
|
|
doc_id: str | list[str] | None = None,
|
|
temperature: float | None = None,
|
|
stream_metadata: bool = False,
|
|
enable_citations: bool = False,
|
|
) -> dict[str, Any] | Iterator[str] | Iterator[dict[str, Any]]:
|
|
"""Collection API equivalent: ``collection.query(...)`` (fewer knobs — no temperature/citations/history)."""
|
|
return self._require_cloud_api().chat_completions(
|
|
messages=messages,
|
|
stream=stream,
|
|
doc_id=doc_id,
|
|
temperature=temperature,
|
|
stream_metadata=stream_metadata,
|
|
enable_citations=enable_citations,
|
|
)
|
|
|
|
def get_document(self, doc_id: str) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``collection.get_document(doc_id)``."""
|
|
return self._require_cloud_api().get_document(doc_id=doc_id)
|
|
|
|
def delete_document(self, doc_id: str) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``collection.delete_document(doc_id)``."""
|
|
return self._require_cloud_api().delete_document(doc_id=doc_id)
|
|
|
|
def list_documents(
|
|
self,
|
|
limit: int = 50,
|
|
offset: int = 0,
|
|
folder_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``collection.list_documents()``.
|
|
|
|
Note the return shape differs between the two APIs:
|
|
|
|
- This legacy method returns the raw API envelope
|
|
``{"documents": [...], "total": int, "limit": int, "offset": int}``
|
|
where each document carries keys ``id`` / ``name`` / ``description``.
|
|
- ``collection.list_documents()`` returns a plain ``list[dict]`` where
|
|
each entry uses keys ``doc_id`` / ``doc_name`` / ``doc_description``
|
|
/ ``doc_type`` and is not paginated.
|
|
|
|
Code that migrates by a simple name swap will silently break — update
|
|
callers to the new key names and dropped pagination envelope.
|
|
"""
|
|
return self._require_cloud_api().list_documents(
|
|
limit=limit,
|
|
offset=offset,
|
|
folder_id=folder_id,
|
|
)
|
|
|
|
def create_folder(
|
|
self,
|
|
name: str,
|
|
description: str | None = None,
|
|
parent_folder_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``client.collection(name)`` (auto-creates)."""
|
|
return self._require_cloud_api().create_folder(
|
|
name=name,
|
|
description=description,
|
|
parent_folder_id=parent_folder_id,
|
|
)
|
|
|
|
def list_folders(self, parent_folder_id: str | None = None) -> dict[str, Any]:
|
|
"""Collection API equivalent: ``client.list_collections()``."""
|
|
return self._require_cloud_api().list_folders(parent_folder_id=parent_folder_id)
|
|
|
|
|
|
class LocalClient(PageIndexClient):
|
|
"""Local mode — indexes and queries documents on your machine.
|
|
|
|
Args:
|
|
model: LLM model for indexing (default: gpt-4o-2024-11-20)
|
|
retrieve_model: LLM model for agent QA (default: gpt-5.4)
|
|
storage_path: Directory for SQLite DB and files (default: ./.pageindex)
|
|
storage: Custom StorageEngine instance (default: SQLiteStorage)
|
|
index_config: Advanced indexing parameters. Pass an IndexConfig instance
|
|
or a dict. All fields have sensible defaults — most users don't need this.
|
|
|
|
Example::
|
|
|
|
# Simple — defaults are fine
|
|
client = LocalClient(model="gpt-5.4")
|
|
|
|
# Advanced — tune indexing parameters
|
|
from pageindex.config import IndexConfig
|
|
client = LocalClient(
|
|
model="gpt-5.4",
|
|
index_config=IndexConfig(toc_check_page_num=30),
|
|
)
|
|
"""
|
|
|
|
def __init__(self, model: str = None, retrieve_model: str = None,
|
|
storage_path: str = None, storage=None,
|
|
index_config: IndexConfig | dict = None):
|
|
self._empty_api_key = False
|
|
self._init_local(model, retrieve_model, storage_path, storage, index_config)
|
|
|
|
|
|
class CloudClient(PageIndexClient):
|
|
"""Cloud mode — fully managed by PageIndex cloud service. No LLM key needed."""
|
|
|
|
def __init__(self, api_key: str):
|
|
if not api_key:
|
|
raise PageIndexAPIError(
|
|
"CloudClient requires a PageIndex API key — get one at "
|
|
"https://dash.pageindex.ai."
|
|
)
|
|
self._empty_api_key = False
|
|
self._init_cloud(api_key)
|