codekingpro/portable-devtools

1from __future__ import annotations
2
3import functools
4import importlib
5import pkgutil
6import threading
7from typing import Any, Callable, Sequence
8
9import tiktoken_ext
10
11import tiktoken
12from tiktoken.core import Encoding
13
14_lock = threading.RLock()
15ENCODINGS: dict[str, Encoding] = {}
16ENCODING_CONSTRUCTORS: dict[str, Callable[[], dict[str, Any]]] | None = None
17
18
19@functools.lru_cache
20def _available_plugin_modules() -> Sequence[str]:
21    # tiktoken_ext is a namespace package
22    # submodules inside tiktoken_ext will be inspected for ENCODING_CONSTRUCTORS attributes
23    # - we use namespace package pattern so `pkgutil.iter_modules` is fast
24    # - it's a separate top-level package because namespace subpackages of non-namespace
25    #   packages don't quite do what you want with editable installs
26    mods = []
27    plugin_mods = pkgutil.iter_modules(tiktoken_ext.__path__, tiktoken_ext.__name__ + ".")
28    for _, mod_name, _ in plugin_mods:
29        mods.append(mod_name)
30    return mods
31
32
33def _find_constructors() -> None:
34    global ENCODING_CONSTRUCTORS
35    with _lock:
36        if ENCODING_CONSTRUCTORS is not None:
37            return
38        ENCODING_CONSTRUCTORS = {}
39
40        try:
41            for mod_name in _available_plugin_modules():
42                mod = importlib.import_module(mod_name)
43                try:
44                    constructors = mod.ENCODING_CONSTRUCTORS
45                except AttributeError as e:
46                    raise ValueError(
47                        f"tiktoken plugin {mod_name} does not define ENCODING_CONSTRUCTORS"
48                    ) from e
49                for enc_name, constructor in constructors.items():
50                    if enc_name in ENCODING_CONSTRUCTORS:
51                        raise ValueError(
52                            f"Duplicate encoding name {enc_name} in tiktoken plugin {mod_name}"
53                        )
54                    ENCODING_CONSTRUCTORS[enc_name] = constructor
55        except Exception:
56            # Ensure we idempotently raise errors
57            ENCODING_CONSTRUCTORS = None
58            raise
59
60
61
62
63def get_encoding(encoding_name: str) -> Encoding:
64    if not isinstance(encoding_name, str):
65        raise ValueError(f"Expected a string in get_encoding, got {type(encoding_name)}")
66
67    if encoding_name in ENCODINGS:
68        return ENCODINGS[encoding_name]
69
70    with _lock:
71        if encoding_name in ENCODINGS:
72            return ENCODINGS[encoding_name]
73
74        if ENCODING_CONSTRUCTORS is None:
75            _find_constructors()
76            assert ENCODING_CONSTRUCTORS is not None
77
78        if encoding_name not in ENCODING_CONSTRUCTORS:
79            raise ValueError(
80                f"Unknown encoding {encoding_name}.\n"
81                f"Plugins found: {_available_plugin_modules()}\n"
82                f"tiktoken version: {tiktoken.__version__} (are you on latest?)"
83            )
84
85        constructor = ENCODING_CONSTRUCTORS[encoding_name]
86        enc = Encoding(**constructor())
87        ENCODINGS[encoding_name] = enc
88        return enc
89
90
91def list_encoding_names() -> list[str]:
92    with _lock:
93        if ENCODING_CONSTRUCTORS is None:
94            _find_constructors()
95            assert ENCODING_CONSTRUCTORS is not None
96        return list(ENCODING_CONSTRUCTORS)
97