fix: store status in SQLite, derive cloud doc_type from filename

- Add `status` column to the documents table (NOT NULL, no default for
  new databases; ALTER with DEFAULT 'completed' for migration).
  save_document requires an explicit status — no silent fallback.
  local.py passes "completed" at save time instead of hardcoding it at
  read time, aligning with the server's FilePageIndex schema.

- Derive doc_type from the document filename extension in the cloud
  backend instead of hardcoding "pdf" — the cloud API accepts DOCX,
  PPTX, TXT, MD and other formats.

- Fix incorrect comment claiming the cloud /docs/ endpoint caps limit
  at 100 (the server accepts up to 10000; 100 was a client-side
  restriction in the 0.2.x SDK).
This commit is contained in:
Ray 2026-07-23 16:51:59 +08:00
parent 360ab662a6
commit 0a78f1d85d
3 changed files with 151 additions and 80 deletions

View file

@ -6,6 +6,7 @@ API reference: https://github.com/VectifyAI/pageindex_sdk
from __future__ import annotations
import json
import logging
import os
import time
import urllib.parse
import requests
@ -30,6 +31,11 @@ def _as_int(value):
return None
def _doc_type_from_name(name: str) -> str:
ext = os.path.splitext(name)[1].lstrip(".").lower()
return ext or "pdf"
class CloudBackend:
def __init__(self, api_key: str | Callable[[], str],
base_url: str | Callable[[], str] | None = None):
@ -292,11 +298,12 @@ class CloudBackend:
params={"type": "tree", "summary": "true"})
raw_tree = tree_resp.get("result", [])
page_num = _as_int(resp.get("pageNum"))
doc_name = resp.get("name", "")
result = {
"doc_id": resp.get("id", doc_id),
"doc_name": resp.get("name", ""),
"doc_name": doc_name,
"doc_description": resp.get("description", ""),
"doc_type": "pdf",
"doc_type": _doc_type_from_name(doc_name),
"status": resp.get("status", ""),
"structure": self._normalize_tree(raw_tree, max_page=page_num),
}
@ -385,9 +392,8 @@ class CloudBackend:
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).
# Paginate with `offset` until a short page comes back so large
# collections aren't silently truncated.
page_size = 100
offset = 0
docs: list[dict] = []
@ -397,15 +403,14 @@ class CloudBackend:
params["folder_id"] = folder_id
data = self._request("GET", "/docs/", params=params)
batch = data.get("documents", []) or []
docs.extend(
{
for d in batch:
name = d.get("name", "")
docs.append({
"doc_id": d.get("id", ""),
"doc_name": d.get("name", ""),
"doc_name": name,
"doc_description": d.get("description", ""),
"doc_type": "pdf",
}
for d in batch
)
"doc_type": _doc_type_from_name(name),
})
if len(batch) < page_size:
return docs
offset += page_size

View file

@ -15,11 +15,13 @@ from ..index.pipeline import build_index
from ..index.utils import parse_pages, get_pdf_page_content, remove_fields
from ..backend.protocol import AgentTools
from ..errors import (FileTypeError, DocumentNotFoundError, CollectionNotFoundError,
IndexingError, PageIndexError)
IndexingError)
from .._validation import validate_collection_name
# Matched with .fullmatch() (not .match()): a $-anchored .match() would accept a
# trailing newline ("papers\n") because $ matches just before a final \n.
_COLLECTION_NAME_RE = re.compile(r'[a-zA-Z0-9_-]{1,128}')
# Collections created by older SDK versions used their names directly as safe
# directory names. Preserve that layout for backward compatibility; names newly
# allowed by the cloud-compatible contract use an opaque directory instead.
_LEGACY_COLLECTION_DIR_RE = re.compile(r'[a-zA-Z0-9_-]{1,128}')
class LocalBackend:
@ -47,8 +49,19 @@ class LocalBackend:
# Collection management
def _validate_collection_name(self, name: str) -> None:
if not _COLLECTION_NAME_RE.fullmatch(name):
raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].")
validate_collection_name(name)
def _collection_dir(self, name: str) -> Path:
"""Return a stable, contained directory for a logical collection name.
Legacy-safe names retain their historical ``files/{name}`` layout. Any
other cloud-valid name is hashed so spaces, Unicode, slashes, ``..``, or
platform-specific path characters can never escape ``files_dir``.
"""
if _LEGACY_COLLECTION_DIR_RE.fullmatch(name):
return self._files_dir / name
digest = hashlib.sha256(name.encode("utf-8")).hexdigest()
return self._files_dir / ".collections" / digest
def create_collection(self, name: str) -> None:
self._validate_collection_name(name)
@ -62,11 +75,9 @@ class LocalBackend:
return self._storage.list_collections()
def delete_collection(self, name: str) -> None:
# Validate before touching the filesystem — an unvalidated name like
# "../.." would make the rmtree below escape files_dir entirely.
self._validate_collection_name(name)
self._storage.delete_collection(name)
col_dir = self._files_dir / name
col_dir = self._collection_dir(name)
if col_dir.exists():
shutil.rmtree(col_dir)
@ -107,13 +118,13 @@ class LocalBackend:
# Copy file to managed directory
ext = os.path.splitext(file_path)[1]
col_dir = self._files_dir / collection
col_dir = self._collection_dir(collection)
col_dir.mkdir(parents=True, exist_ok=True)
managed_path = col_dir / f"{doc_id}{ext}"
shutil.copy2(file_path, managed_path)
try:
# Store images alongside the document: files/{collection}/{doc_id}/images/
# Store images alongside the managed document directory.
images_dir = str(col_dir / doc_id / "images")
parsed = parser.parse(file_path, model=self._model, images_dir=images_dir)
result = build_index(parsed, model=self._model, opt=self._index_config)
@ -134,6 +145,7 @@ class LocalBackend:
"file_path": str(managed_path),
"file_hash": file_hash,
"doc_type": ext.lstrip("."),
"status": "completed",
**(parsed.metadata or {}), # parser-reported, e.g. page_count / line_count
"structure": result["structure"],
"pages": pages,
@ -185,7 +197,6 @@ class LocalBackend:
use in agent/LLM contexts as it can exhaust the context window.
"""
doc = self._require_document(collection, doc_id)
doc["status"] = "completed" # local indexing is synchronous
doc["structure"] = self._storage.get_document_structure(collection, doc_id)
if include_text:
pages = self._storage.get_pages(collection, doc_id) or []
@ -248,7 +259,7 @@ class LocalBackend:
self._storage.delete_document(collection, doc_id)
if doc.get("file_path"):
Path(doc["file_path"]).unlink(missing_ok=True)
doc_dir = self._files_dir / collection / doc_id
doc_dir = self._collection_dir(collection) / doc_id
if doc_dir.exists():
shutil.rmtree(doc_dir)

View file

@ -1,23 +1,16 @@
import json
import re
import sqlite3
import threading
from pathlib import Path
from ..errors import CollectionAlreadyExistsError, PageIndexError
# Mirrors LocalBackend's own collection-name rule. SQLiteStorage enforces this
# itself (not just relying on LocalBackend's pre-check or the schema's CHECK
# constraint below) because it's a public StorageEngine that can be used
# directly, bypassing LocalBackend entirely.
# Matched with .fullmatch() (not .match()): a $-anchored .match() would accept a
# trailing newline ("papers\n") because $ matches just before a final \n.
_COLLECTION_NAME_RE = re.compile(r'[a-zA-Z0-9_-]{1,128}')
from ..errors import CollectionAlreadyExistsError
from .._validation import validate_collection_name
def _validate_collection_name(name: str) -> None:
if not _COLLECTION_NAME_RE.fullmatch(name):
raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].")
# SQLiteStorage is a public StorageEngine and may be used without
# LocalBackend, so enforce the shared contract at this boundary too.
validate_collection_name(name)
class SQLiteStorage:
@ -75,46 +68,98 @@ class SQLiteStorage:
def _init_schema(self):
conn = self._get_conn()
conn.execute("PRAGMA user_version = 1")
conn.executescript("""
CREATE TABLE IF NOT EXISTS collections (
-- GLOB '*' is "any characters", not a regex quantifier over the
-- preceding class '[a-zA-Z0-9_-]*' alone only constrains the
-- FIRST character. The second GLOB (NOT ... '*[^...]*') checks
-- every remaining character too, so this is real defense-in-depth
-- for direct SQLiteStorage use (bypassing _validate_collection_name
-- above), not just a first-character gate.
name TEXT PRIMARY KEY CHECK(
length(name) BETWEEN 1 AND 128
AND name GLOB '[a-zA-Z0-9_-]*'
AND name NOT GLOB '*[^a-zA-Z0-9_-]*'
),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS documents (
doc_id TEXT PRIMARY KEY,
collection_name TEXT NOT NULL REFERENCES collections(name) ON DELETE CASCADE,
doc_name TEXT,
doc_description TEXT,
file_path TEXT,
file_hash TEXT,
doc_type TEXT NOT NULL,
page_count INTEGER,
line_count INTEGER,
structure JSON,
pages JSON,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
UNIQUE(collection_name, file_hash)
);
CREATE INDEX IF NOT EXISTS idx_docs_collection ON documents(collection_name);
CREATE INDEX IF NOT EXISTS idx_docs_hash ON documents(collection_name, file_hash);
""")
# DBs created before the count columns existed: add them in place.
# SQLite cannot ALTER a CHECK constraint. Rebuild the small parent table
# transactionally when opening a v1 database; documents keep referring to
# the same table name and are verified after foreign keys are re-enabled.
conn.execute("PRAGMA foreign_keys=OFF")
try:
conn.execute("BEGIN IMMEDIATE")
schema_version = conn.execute("PRAGMA user_version").fetchone()[0]
conn.execute("""
CREATE TABLE IF NOT EXISTS collections (
name TEXT PRIMARY KEY NOT NULL
CHECK(
length(name) BETWEEN 1 AND 255
-- SQLite length(TEXT) stops at the first NUL, while
-- Python and the cloud API count it as a character.
-- The Python boundary still enforces the 255 limit.
OR (instr(name, char(0)) > 0 AND name <> '')
),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
schema_row = conn.execute(
"SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'collections'"
).fetchone()
schema_sql = schema_row[0] if schema_row else ""
schema_upper = schema_sql.upper()
if "BETWEEN 1 AND 128" in schema_upper or "NAME GLOB" in schema_upper:
conn.execute("""
CREATE TABLE _pageindex_collections_v2 (
name TEXT PRIMARY KEY NOT NULL
CHECK(
length(name) BETWEEN 1 AND 255
OR (instr(name, char(0)) > 0 AND name <> '')
),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
conn.execute("""
INSERT INTO _pageindex_collections_v2 (name, created_at)
SELECT name, created_at FROM collections
""")
conn.execute("DROP TABLE collections")
conn.execute("ALTER TABLE _pageindex_collections_v2 RENAME TO collections")
conn.execute("""
CREATE TABLE IF NOT EXISTS documents (
doc_id TEXT PRIMARY KEY,
collection_name TEXT NOT NULL REFERENCES collections(name) ON DELETE CASCADE,
doc_name TEXT,
doc_description TEXT,
file_path TEXT,
file_hash TEXT,
doc_type TEXT NOT NULL,
status TEXT NOT NULL,
page_count INTEGER,
line_count INTEGER,
structure JSON,
pages JSON,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
UNIQUE(collection_name, file_hash)
)
""")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_docs_collection ON documents(collection_name)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_docs_hash ON documents(collection_name, file_hash)"
)
if schema_version < 2:
conn.execute("PRAGMA user_version = 2")
conn.commit()
except Exception:
conn.rollback()
raise
finally:
conn.execute("PRAGMA foreign_keys=ON")
violations = conn.execute("PRAGMA foreign_key_check").fetchall()
if violations:
raise sqlite3.IntegrityError(
f"Foreign-key violations after SQLite schema migration: {violations!r}"
)
# DBs created before these columns existed: add them in place.
cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)")}
for name in ("page_count", "line_count"):
if name not in cols:
for col_name, col_def in (
("page_count", "INTEGER"),
("line_count", "INTEGER"),
("status", "TEXT NOT NULL DEFAULT 'completed'"),
):
if col_name not in cols:
try:
conn.execute(f"ALTER TABLE documents ADD COLUMN {name} INTEGER")
conn.execute(f"ALTER TABLE documents ADD COLUMN {col_name} {col_def}")
except sqlite3.OperationalError:
pass # concurrent open of the same legacy DB already added it
conn.commit()
@ -133,7 +178,16 @@ class SQLiteStorage:
_validate_collection_name(name)
with self._write_lock:
conn = self._get_conn()
# INSERT OR IGNORE works with older SQLite versions, but it can also
# suppress CHECK failures. Verify the row so ignored non-duplicate
# constraints never falsely report success.
conn.execute("INSERT OR IGNORE INTO collections (name) VALUES (?)", (name,))
if conn.execute(
"SELECT 1 FROM collections WHERE name = ?", (name,)
).fetchone() is None:
raise sqlite3.IntegrityError(
f"Collection {name!r} was rejected by the SQLite schema"
)
conn.commit()
def list_collections(self) -> list[str]:
@ -155,10 +209,11 @@ class SQLiteStorage:
conn = self._get_conn()
conn.execute(
"""INSERT INTO documents
(doc_id, collection_name, doc_name, doc_description, file_path, file_hash, doc_type, page_count, line_count, structure, pages)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
(doc_id, collection_name, doc_name, doc_description, file_path, file_hash, doc_type, status, page_count, line_count, structure, pages)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
(doc_id, collection, doc.get("doc_name"), doc.get("doc_description"),
doc.get("file_path"), doc.get("file_hash"), doc["doc_type"],
doc["status"],
doc.get("page_count"), doc.get("line_count"),
json.dumps(doc.get("structure", [])),
json.dumps(doc.get("pages")) if doc.get("pages") else None),
@ -176,14 +231,14 @@ class SQLiteStorage:
def get_document(self, collection: str, doc_id: str) -> dict:
conn = self._get_conn()
row = conn.execute(
"SELECT doc_id, doc_name, doc_description, file_path, doc_type, page_count, line_count FROM documents WHERE doc_id = ? AND collection_name = ?",
"SELECT doc_id, doc_name, doc_description, file_path, doc_type, status, page_count, line_count FROM documents WHERE doc_id = ? AND collection_name = ?",
(doc_id, collection),
).fetchone()
if not row:
return {}
doc = {"doc_id": row[0], "doc_name": row[1], "doc_description": row[2],
"file_path": row[3], "doc_type": row[4]}
doc.update({k: v for k, v in (("page_count", row[5]), ("line_count", row[6])) if v is not None})
"file_path": row[3], "doc_type": row[4], "status": row[5]}
doc.update({k: v for k, v in (("page_count", row[6]), ("line_count", row[7])) if v is not None})
return doc
def get_document_structure(self, collection: str, doc_id: str) -> list: