From b2756c73d2552dade52da2339b8355cc8c055483 Mon Sep 17 00:00:00 2001 From: mountain Date: Tue, 7 Jul 2026 12:15:34 +0800 Subject: [PATCH] refactor(sdk): typed returns, protocol contract, parser layering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Engineering-quality cleanups from the SDK review (no behavior change): - Return-type discoverability: add pageindex/types.py with TypedDicts (DocumentInfo, DocumentDetail, PageContent) and annotate Collection / Backend methods with them; add docstrings to every public Collection method (including the get_page_content `pages` spec). Exported from the package. Zero runtime cost — these are plain dicts. - Backend protocol as a real contract: * query_stream is an async generator, so the protocol now declares it as `def ... -> AsyncIterator[QueryEvent]` (not `async def`, which typed it as a coroutine and never matched the implementations). * custom-parser support is expressed as a runtime_checkable SupportsParserRegistration capability protocol; the client uses isinstance(...) instead of hasattr(...) duck-typing. - Parser layering: move count_tokens into a leaf module pageindex/tokens.py so parser/* imports it from there instead of reaching back into pageindex.index (a reverse dependency). index.utils re-exports it for backward compatibility. Adds tests/test_architecture.py enforcing: parser never imports index, count_tokens is a single shared leaf, the capability protocol works, both backends satisfy Backend, and the TypedDicts are exported. Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS --- pageindex/__init__.py | 4 +++ pageindex/backend/protocol.py | 21 +++++++++---- pageindex/client.py | 3 +- pageindex/collection.py | 34 +++++++++++++++++++-- pageindex/index/utils.py | 7 +---- pageindex/parser/markdown.py | 2 +- pageindex/parser/pdf.py | 2 +- pageindex/tokens.py | 11 +++++++ pageindex/types.py | 32 ++++++++++++++++++++ tests/test_architecture.py | 56 +++++++++++++++++++++++++++++++++++ 10 files changed, 155 insertions(+), 17 deletions(-) create mode 100644 pageindex/tokens.py create mode 100644 pageindex/types.py create mode 100644 tests/test_architecture.py diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 166c33a..b649cb2 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -10,6 +10,7 @@ from .retrieve import get_document, get_document_structure, get_page_content from .client import PageIndexClient, LocalClient, CloudClient from .config import IndexConfig, set_llm_params from .collection import Collection +from .types import DocumentInfo, DocumentDetail, PageContent from .parser.protocol import ContentNode, ParsedDocument, DocumentParser from .storage.protocol import StorageEngine from .events import QueryEvent @@ -30,6 +31,9 @@ __all__ = [ "IndexConfig", "set_llm_params", "Collection", + "DocumentInfo", + "DocumentDetail", + "PageContent", "ContentNode", "ParsedDocument", "DocumentParser", diff --git a/pageindex/backend/protocol.py b/pageindex/backend/protocol.py index 214aff4..c5b390a 100644 --- a/pageindex/backend/protocol.py +++ b/pageindex/backend/protocol.py @@ -3,6 +3,7 @@ from dataclasses import dataclass, field from typing import Protocol, Any, AsyncIterator, runtime_checkable from ..events import QueryEvent +from ..types import DocumentInfo, DocumentDetail, PageContent @dataclass @@ -22,15 +23,25 @@ class Backend(Protocol): # Document management def add_document(self, collection: str, file_path: str) -> str: ... - def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> dict: ... + def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> DocumentDetail: ... def get_document_structure(self, collection: str, doc_id: str) -> list: ... - def get_page_content(self, collection: str, doc_id: str, pages: str) -> list: ... - def list_documents(self, collection: str) -> list[dict]: ... + def get_page_content(self, collection: str, doc_id: str, pages: str) -> list[PageContent]: ... + def list_documents(self, collection: str) -> list[DocumentInfo]: ... def delete_document(self, collection: str, doc_id: str) -> None: ... # Query — doc_ids accepts a single id or a list; implementations should # normalize internally (a bare str is treated as a single-element list). def query(self, collection: str, question: str, doc_ids: str | list[str] | None = None) -> str: ... - async def query_stream(self, collection: str, question: str, - doc_ids: str | list[str] | None = None) -> AsyncIterator[QueryEvent]: ... + # query_stream is an async generator: calling it returns an async iterator + # WITHOUT awaiting, so it is declared as a plain def returning + # AsyncIterator (not `async def`, which would be a coroutine). + def query_stream(self, collection: str, question: str, + doc_ids: str | list[str] | None = None) -> AsyncIterator[QueryEvent]: ... + + +@runtime_checkable +class SupportsParserRegistration(Protocol): + """Capability protocol: a backend that accepts custom document parsers + (local mode). Cloud backends don't implement this.""" + def register_parser(self, parser: Any) -> None: ... diff --git a/pageindex/client.py b/pageindex/client.py index fdbacbc..a143ca8 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -145,7 +145,8 @@ class PageIndexClient: def register_parser(self, parser: DocumentParser) -> None: """Register a custom document parser. Only available in local mode.""" - if not hasattr(self._backend, 'register_parser'): + 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) diff --git a/pageindex/collection.py b/pageindex/collection.py index 4cf2f64..fe0110d 100644 --- a/pageindex/collection.py +++ b/pageindex/collection.py @@ -5,6 +5,7 @@ import warnings from typing import AsyncIterator from .events import QueryEvent from .backend.protocol import Backend +from .types import DocumentInfo, DocumentDetail, PageContent def _multidoc_acked() -> bool: @@ -50,21 +51,48 @@ class Collection: return self._name def add(self, file_path: str) -> str: + """Index a document (PDF or Markdown) into this collection. + + Returns the ``doc_id``. Re-adding byte-identical content returns the + existing doc_id (content-hash dedup); change ``IndexConfig`` won't + force a re-index — delete the doc first if you need a fresh tree. + """ return self._backend.add_document(self._name, file_path) - def list_documents(self) -> list[dict]: + def list_documents(self) -> list[DocumentInfo]: + """List every document in this collection. + + Each item has ``doc_id``, ``doc_name``, ``doc_description``, ``doc_type``. + """ return self._backend.list_documents(self._name) - def get_document(self, doc_id: str, include_text: bool = False) -> dict: + def get_document(self, doc_id: str, include_text: bool = False) -> DocumentDetail: + """Return a document's metadata plus its tree under ``structure``. + + ``include_text=True`` fills each node's text from cached pages (local + backend only; can be large — avoid for LLM contexts). Raises + ``DocumentNotFoundError`` if the doc_id is unknown. + """ return self._backend.get_document(self._name, doc_id, include_text=include_text) def get_document_structure(self, doc_id: str) -> list: + """Return the document's hierarchical tree (a list of node dicts).""" return self._backend.get_document_structure(self._name, doc_id) - def get_page_content(self, doc_id: str, pages: str) -> list: + def get_page_content(self, doc_id: str, pages: str) -> list[PageContent]: + """Return content for specific pages. + + ``pages`` is a range/list spec: ``"5-7"``, ``"3,8"``, or ``"12"``. + Each returned item has ``page`` and ``content`` (and ``images`` when + present). For Markdown docs, "page" numbers map to line ranges. + """ return self._backend.get_page_content(self._name, doc_id, pages) def delete_document(self, doc_id: str) -> None: + """Delete a document and its stored files/artifacts. + + Raises ``DocumentNotFoundError`` if the doc_id is unknown. + """ self._backend.delete_document(self._name, doc_id) def query(self, question: str, diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index 60cb770..a2701e4 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -17,16 +17,11 @@ from pprint import pprint from types import SimpleNamespace as config from ..config import get_llm_params +from ..tokens import count_tokens # re-exported for backward compat logger = logging.getLogger(__name__) -def count_tokens(text, model=None): - if not text: - return 0 - return litellm.token_counter(model=model, text=text) - - def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): if model: model = model.removeprefix("litellm/") diff --git a/pageindex/parser/markdown.py b/pageindex/parser/markdown.py index 0eba5bb..7c843e8 100644 --- a/pageindex/parser/markdown.py +++ b/pageindex/parser/markdown.py @@ -1,7 +1,7 @@ import re from pathlib import Path from .protocol import ContentNode, ParsedDocument -from ..index.utils import count_tokens +from ..tokens import count_tokens class MarkdownParser: diff --git a/pageindex/parser/pdf.py b/pageindex/parser/pdf.py index 83106da..a2739a5 100644 --- a/pageindex/parser/pdf.py +++ b/pageindex/parser/pdf.py @@ -1,7 +1,7 @@ import pymupdf from pathlib import Path from .protocol import ContentNode, ParsedDocument -from ..index.utils import count_tokens +from ..tokens import count_tokens # Minimum image dimension to keep (skip icons/artifacts) _MIN_IMAGE_SIZE = 32 diff --git a/pageindex/tokens.py b/pageindex/tokens.py new file mode 100644 index 0000000..8d5e9a3 --- /dev/null +++ b/pageindex/tokens.py @@ -0,0 +1,11 @@ +# pageindex/tokens.py +# Leaf utility so both the parser and index layers can count tokens without +# the parser reaching back into pageindex.index (a reverse/horizontal +# dependency). Depends only on litellm. +import litellm + + +def count_tokens(text, model=None): + if not text: + return 0 + return litellm.token_counter(model=model, text=text) diff --git a/pageindex/types.py b/pageindex/types.py new file mode 100644 index 0000000..c997497 --- /dev/null +++ b/pageindex/types.py @@ -0,0 +1,32 @@ +# pageindex/types.py +# TypedDicts describing the plain-dict shapes the SDK returns, so callers get +# key/field discovery in their IDE without any runtime cost (these are dicts). +from __future__ import annotations + +from typing import Any, TypedDict + + +class DocumentInfo(TypedDict): + """A document as returned by ``list_documents()``.""" + doc_id: str + doc_name: str + doc_description: str + doc_type: str + + +class DocumentDetail(DocumentInfo, total=False): + """A document with its tree, as returned by ``get_document()``. + + ``structure`` is always present; ``file_path`` is local-only and + ``status`` is cloud-only, hence total=False. + """ + structure: list[dict[str, Any]] + file_path: str # local backend only + status: str # cloud backend only + + +class PageContent(TypedDict, total=False): + """One page of content, as returned by ``get_page_content()``.""" + page: int + content: str + images: list[dict[str, Any]] diff --git a/tests/test_architecture.py b/tests/test_architecture.py new file mode 100644 index 0000000..b52f51c --- /dev/null +++ b/tests/test_architecture.py @@ -0,0 +1,56 @@ +"""Layering / protocol-contract guards for the SDK.""" +import ast +from pathlib import Path + +import pytest + +from pageindex.backend.protocol import Backend, SupportsParserRegistration +from pageindex.backend.cloud import CloudBackend + + +def test_parser_layer_does_not_import_index(): + """parser/* must not depend on the index package (reverse dependency).""" + parser_dir = Path(__file__).parent.parent / "pageindex" / "parser" + offenders = [] + for py in parser_dir.glob("*.py"): + tree = ast.parse(py.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module and "index" in node.module.split("."): + offenders.append(f"{py.name}: from {node.module}") + assert not offenders, f"parser imports index: {offenders}" + + +def test_count_tokens_lives_in_leaf_module(): + from pageindex.tokens import count_tokens + from pageindex.index.utils import count_tokens as reexport + from pageindex.parser.pdf import count_tokens as parser_ct + # index re-exports the leaf function; parser imports the same leaf. + assert reexport is count_tokens is parser_ct + + +def test_parser_registration_is_a_capability_protocol(): + from unittest.mock import MagicMock + from pageindex.backend.local import LocalBackend + lb = LocalBackend(storage=MagicMock(), files_dir="/tmp/x", model="m") + assert isinstance(lb, SupportsParserRegistration) + assert not isinstance(CloudBackend(api_key="pi-test"), SupportsParserRegistration) + + +def test_register_parser_rejected_in_cloud_mode(): + from pageindex import CloudClient + from pageindex.errors import PageIndexError + client = CloudClient(api_key="pi-test") + with pytest.raises(PageIndexError, match="not supported in cloud mode"): + client.register_parser(object()) + + +def test_both_backends_satisfy_backend_protocol(): + from unittest.mock import MagicMock + from pageindex.backend.local import LocalBackend + assert isinstance(CloudBackend(api_key="pi-test"), Backend) + assert isinstance(LocalBackend(storage=MagicMock(), files_dir="/tmp/x", model="m"), Backend) + + +def test_typed_dicts_are_exported(): + import pageindex.types as t + assert {"DocumentInfo", "DocumentDetail", "PageContent"} <= set(dir(t))