refactor(sdk): typed returns, protocol contract, parser layering

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

@@ -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",

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: ...

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)

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,

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/")

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:

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

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)

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]]

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))