PageIndex/pageindex/tokens.py
mountain b2756c73d2 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
2026-07-07 12:15:34 +08:00

11 lines
352 B
Python

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