mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-24 21:41:04 +02:00
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:
parent
360ab662a6
commit
0a78f1d85d
3 changed files with 151 additions and 80 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue