From dd7b805846174bd8672708ca830136deca8b9d41 Mon Sep 17 00:00:00 2001 From: Chaitany Patel Date: Thu, 10 Sep 2026 11:59:37 +0530 Subject: [PATCH 1/2] feat(context): Add TokenBudget and pluggable tokenizers Create the feast.context module with the token accounting that prompt assembly in an OnDemandFeatureView builds on. Tokenizers resolve the way online store types do in repo_config: a built-in name maps to a class path in TOKENIZER_CLASS_FOR_TYPE, and anything else is itself the path of a Tokenizer subclass with a no-argument constructor, loaded through import_class. Adding a tokenizer is therefore what adding a vector store is - write the class, pass its path, no registration call and no mutable global state. Feast ships cl100k_base, o200k_base and a dependency-free character estimate. TokenBudget is a frozen dataclass holding a resolved Tokenizer: consume() returns a new budget rather than mutating, and try_consume() returns None instead of raising, which is the primitive priority_select() will use to fill a budget greedily. TokenBudget.of() resolves a name, class path or instance and can refuse the fallback; is_approximate reports whether counts are exact. Tokenizers compare by value, so budgets built from the same name are interchangeable as dict keys and set members. When tiktoken cannot be loaded, get_tokenizer() degrades to the character estimate with a logged warning; fallback=False makes it fatal for paths that need exact counts. TiktokenTokenizer counts with disallowed_special=(), so a feature value containing a literal "<|endoftext|>" is treated as ordinary text instead of raising. tiktoken is not declared as a dependency yet, so the tests that assert exact encodings skip when it is absent. Part of RHOAIENG-80116 (PR 1/9). Signed-off-by: Chaitany Patel --- sdk/python/feast/context/__init__.py | 39 +++ sdk/python/feast/context/errors.py | 53 ++++ sdk/python/feast/context/token_budget.py | 163 ++++++++++++ sdk/python/feast/context/tokenizer.py | 237 ++++++++++++++++++ sdk/python/tests/unit/context/__init__.py | 0 sdk/python/tests/unit/context/conftest.py | 22 ++ .../tests/unit/context/test_token_budget.py | 200 +++++++++++++++ .../tests/unit/context/test_tokenizer.py | 141 +++++++++++ 8 files changed, 855 insertions(+) create mode 100644 sdk/python/feast/context/__init__.py create mode 100644 sdk/python/feast/context/errors.py create mode 100644 sdk/python/feast/context/token_budget.py create mode 100644 sdk/python/feast/context/tokenizer.py create mode 100644 sdk/python/tests/unit/context/__init__.py create mode 100644 sdk/python/tests/unit/context/conftest.py create mode 100644 sdk/python/tests/unit/context/test_token_budget.py create mode 100644 sdk/python/tests/unit/context/test_tokenizer.py diff --git a/sdk/python/feast/context/__init__.py b/sdk/python/feast/context/__init__.py new file mode 100644 index 00000000000..62de13ebead --- /dev/null +++ b/sdk/python/feast/context/__init__.py @@ -0,0 +1,39 @@ +"""Token-budget-aware helpers for assembling LLM context in feature views. + +Plain functions over immutable values, so an OnDemandFeatureView calling them +gets the same prompt offline and online. +""" + +from feast.context.errors import ( + ContextError, + TokenBudgetExceededError, + TokenizerNotFoundError, + TokenizerUnavailableError, +) +from feast.context.token_budget import TokenBudget +from feast.context.tokenizer import ( + ApproximateTokenizer, + Cl100kBaseTokenizer, + O200kBaseTokenizer, + TiktokenTokenizer, + Tokenizer, + TokenizerName, + TokenizerSpec, + get_tokenizer, +) + +__all__ = [ + "ApproximateTokenizer", + "Cl100kBaseTokenizer", + "ContextError", + "O200kBaseTokenizer", + "TiktokenTokenizer", + "TokenBudget", + "TokenBudgetExceededError", + "Tokenizer", + "TokenizerName", + "TokenizerNotFoundError", + "TokenizerSpec", + "TokenizerUnavailableError", + "get_tokenizer", +] diff --git a/sdk/python/feast/context/errors.py b/sdk/python/feast/context/errors.py new file mode 100644 index 00000000000..2504eabebd9 --- /dev/null +++ b/sdk/python/feast/context/errors.py @@ -0,0 +1,53 @@ +from typing import Iterable + +from fastapi import status as HttpStatusCode + +from feast.errors import FeastError + + +class ContextError(FeastError): + """Base class for all feast.context failures.""" + + +class TokenizerNotFoundError(ContextError): + def __init__(self, name: str, available: Iterable[str]): + super().__init__( + f"Unknown tokenizer '{name}'. Built-in tokenizers: " + f"{', '.join(sorted(available))}. For your own, pass the class " + f"path of a Tokenizer subclass, e.g. 'my_pkg.MyTokenizer'." + ) + + def http_status_code(self) -> int: + return HttpStatusCode.HTTP_400_BAD_REQUEST + + +class TokenizerUnavailableError(ContextError): + """A known tokenizer cannot be loaded. + + Covers a missing dependency, and one present but unable to load its + encoding — tiktoken fetching a BPE file on a host with no network, say. + """ + + def __init__(self, name: str, reason: str, install_hint: str): + super().__init__( + f"Tokenizer '{name}' is unavailable: {reason}. Install its " + f"dependencies with: {install_hint}" + ) + self.tokenizer_name = name + + def http_status_code(self) -> int: + return HttpStatusCode.HTTP_400_BAD_REQUEST + + +class TokenBudgetExceededError(ContextError): + def __init__(self, requested: int, remaining: int, max_tokens: int): + super().__init__( + f"Cannot consume {requested} tokens: only {remaining} of " + f"{max_tokens} tokens remain in the budget." + ) + self.requested = requested + self.remaining = remaining + self.max_tokens = max_tokens + + def http_status_code(self) -> int: + return HttpStatusCode.HTTP_400_BAD_REQUEST diff --git a/sdk/python/feast/context/token_budget.py b/sdk/python/feast/context/token_budget.py new file mode 100644 index 00000000000..2b2ccc8657d --- /dev/null +++ b/sdk/python/feast/context/token_budget.py @@ -0,0 +1,163 @@ +"""Token accounting for prompt assembly.""" + +from dataclasses import dataclass, field +from typing import Optional + +from feast.context.errors import TokenBudgetExceededError +from feast.context.tokenizer import ( + ApproximateTokenizer, + Tokenizer, + TokenizerName, + TokenizerSpec, + get_tokenizer, +) + + +@dataclass(frozen=True) +class TokenBudget: + """An immutable token allowance measured with a specific tokenizer. + + ``consume()`` returns a new budget instead of mutating this one, so a + budget is safe to share across threads and to reuse between assemblies:: + + budget = TokenBudget.of(4096, "cl100k_base") + budget.count_tokens("Hello world") # -> 2 + budget = budget.consume("Hello world") + budget.remaining # -> 4094 + + Attributes: + max_tokens: Size of the allowance. Zero admits only empty content. + tokenizer: The tokenizer doing the counting. A name is accepted too + and resolved at construction; :meth:`of` types that properly and + can forbid the fallback to the estimate. + consumed: Tokens already spent. + """ + + max_tokens: int + tokenizer: Tokenizer = field(default_factory=get_tokenizer) + consumed: int = 0 + + def __post_init__(self) -> None: + if self.max_tokens < 0: + raise ValueError(f"max_tokens must not be negative, got {self.max_tokens}.") + if self.consumed < 0: + raise ValueError(f"consumed must not be negative, got {self.consumed}.") + if self.consumed > self.max_tokens: + raise ValueError( + f"consumed ({self.consumed}) must not exceed max_tokens " + f"({self.max_tokens})." + ) + if not isinstance(self.tokenizer, Tokenizer): + # A name from config lands here; resolve it once, so counting + # never re-enters resolution and equality stays value-based. + object.__setattr__(self, "tokenizer", get_tokenizer(self.tokenizer)) + + @classmethod + def of( + cls, + max_tokens: int, + tokenizer: TokenizerSpec = TokenizerName.CL100K_BASE, + *, + consumed: int = 0, + fallback: bool = True, + ) -> "TokenBudget": + """Build a budget from a tokenizer name, instance, or class path. + + Args: + fallback: Set False to raise instead of degrading to the character + estimate when the tokenizer cannot be loaded. + """ + return cls( + max_tokens=max_tokens, + tokenizer=get_tokenizer(tokenizer, fallback=fallback), + consumed=consumed, + ) + + @property + def remaining(self) -> int: + """Tokens still available.""" + return self.max_tokens - self.consumed + + @property + def is_exhausted(self) -> bool: + """Whether the allowance is fully spent.""" + return self.remaining == 0 + + @property + def tokenizer_name(self) -> str: + """Name of the tokenizer actually counting.""" + return self.tokenizer.name + + @property + def is_approximate(self) -> bool: + """Whether counts are estimated rather than exact. + + True when the requested tokenizer could not be loaded and the budget + fell back to counting characters. + """ + return isinstance(self.tokenizer, ApproximateTokenizer) + + def count_tokens(self, text: str) -> int: + """Tokens ``text`` would occupy, regardless of what remains.""" + return self.tokenizer.count_tokens(text) + + def fits(self, text: str) -> bool: + """Whether ``text`` fits in the remaining allowance.""" + return self.count_tokens(text) <= self.remaining + + def fits_tokens(self, tokens: int) -> bool: + """Whether ``tokens`` more tokens fit in the remaining allowance.""" + _check_not_negative(tokens) + return tokens <= self.remaining + + def consume(self, text: str) -> "TokenBudget": + """Charge ``text`` against the budget, returning the new budget. + + Raises: + TokenBudgetExceededError: ``text`` does not fit; guard with + ``fits()`` or use ``try_consume()``. + """ + return self.consume_tokens(self.count_tokens(text)) + + def consume_tokens(self, tokens: int) -> "TokenBudget": + """Charge ``tokens`` against the budget, returning the new budget. + + Raises: + TokenBudgetExceededError: ``tokens`` exceeds what remains. + """ + if not self.fits_tokens(tokens): + raise TokenBudgetExceededError(tokens, self.remaining, self.max_tokens) + return self._replace_consumed(self.consumed + tokens) + + def try_consume(self, text: str) -> Optional["TokenBudget"]: + """Charge ``text`` if it fits, else return None. + + Lets a caller keep whichever candidate sections still fit, without + catching exceptions. + """ + tokens = self.count_tokens(text) + if not self.fits_tokens(tokens): + return None + return self._replace_consumed(self.consumed + tokens) + + def reset(self) -> "TokenBudget": + """A budget with the same limit and tokenizer, nothing consumed.""" + return self._replace_consumed(0) + + def _replace_consumed(self, consumed: int) -> "TokenBudget": + return TokenBudget( + max_tokens=self.max_tokens, + tokenizer=self.tokenizer, + consumed=consumed, + ) + + def __str__(self) -> str: + return ( + f"TokenBudget({self.consumed}/{self.max_tokens} tokens used, " + f"tokenizer={self.tokenizer_name})" + ) + + +def _check_not_negative(tokens: int) -> None: + if tokens < 0: + raise ValueError(f"tokens must not be negative, got {tokens}.") diff --git a/sdk/python/feast/context/tokenizer.py b/sdk/python/feast/context/tokenizer.py new file mode 100644 index 00000000000..09af77def68 --- /dev/null +++ b/sdk/python/feast/context/tokenizer.py @@ -0,0 +1,237 @@ +"""Tokenizers and how a name resolves to one. + +Resolution follows the same convention as online store types in +``repo_config``: a built-in name maps to a class path, and anything else is +itself the fully-qualified path of a :class:`Tokenizer` subclass with a +no-argument constructor:: + + TokenBudget(max_tokens=4096, tokenizer="cl100k_base") + TokenBudget(max_tokens=4096, tokenizer="my_pkg.tokenizers.LlamaTokenizer") +""" + +import logging +import math +from abc import ABC, abstractmethod +from enum import Enum +from functools import lru_cache +from typing import Any, Union + +from feast.context.errors import TokenizerNotFoundError, TokenizerUnavailableError +from feast.errors import FeastInvalidBaseClass +from feast.importer import import_class + +logger = logging.getLogger(__name__) + +#: Average characters per token for English prose. +DEFAULT_CHARS_PER_TOKEN = 4 + +#: Quoted in error messages when tiktoken is missing. +TIKTOKEN_INSTALL_HINT = "pip install tiktoken" + + +class TokenizerName(str, Enum): + """Built-in tokenizer names.""" + + #: GPT-4, GPT-3.5-turbo, text-embedding-3-*. + CL100K_BASE = "cl100k_base" + #: GPT-4o and the o-series models. + O200K_BASE = "o200k_base" + #: Character estimate; no third-party dependency. + APPROXIMATE = "approximate" + + +TOKENIZER_CLASS_FOR_TYPE = { + TokenizerName.CL100K_BASE.value: "feast.context.tokenizer.Cl100kBaseTokenizer", + TokenizerName.O200K_BASE.value: "feast.context.tokenizer.O200kBaseTokenizer", + TokenizerName.APPROXIMATE.value: "feast.context.tokenizer.ApproximateTokenizer", +} + + +class Tokenizer(ABC): + """Counts the tokens a model would consume for a piece of text. + + Implementations must be stateless and thread-safe: one instance per name is + shared across every :class:`TokenBudget`. A third-party subclass needs a + no-argument constructor and a class name ending in ``Tokenizer``, which is + what resolution by class path looks for. + + Two tokenizers that count identically compare equal, so budgets built from + the same name are interchangeable as dict keys and set members. + """ + + @property + @abstractmethod + def name(self) -> str: + """Identifier reported in budgets and logs.""" + + @abstractmethod + def count_tokens(self, text: str) -> int: + """Number of tokens in ``text``. Returns 0 for the empty string.""" + + def _identity(self) -> tuple[object, ...]: + """Whatever makes two instances count the same way.""" + return (self.name,) + + def __eq__(self, other: object) -> bool: + return isinstance(other, Tokenizer) and self._identity() == other._identity() + + def __hash__(self) -> int: + return hash(self._identity()) + + def __repr__(self) -> str: + return f"{type(self).__name__}(name={self.name!r})" + + +class ApproximateTokenizer(Tokenizer): + """Character-count estimate used when no real tokenizer is available. + + Counts ``ceil(len(text) / chars_per_token)``: close enough for budgeting, + not for exact context-window arithmetic. Code and non-Latin scripts drift + furthest from the ratio. + """ + + def __init__(self, chars_per_token: int = DEFAULT_CHARS_PER_TOKEN) -> None: + if chars_per_token <= 0: + raise ValueError( + f"chars_per_token must be positive, got {chars_per_token}." + ) + self._chars_per_token = chars_per_token + + @property + def name(self) -> str: + return TokenizerName.APPROXIMATE.value + + @property + def chars_per_token(self) -> int: + return self._chars_per_token + + def count_tokens(self, text: str) -> int: + return math.ceil(len(text) / self._chars_per_token) + + def _identity(self) -> tuple[object, ...]: + return (self.name, self._chars_per_token) + + +class TiktokenTokenizer(Tokenizer): + """Exact token counts from a tiktoken encoding. + + Raises: + TokenizerUnavailableError: tiktoken is missing or the encoding failed + to load. + TokenizerNotFoundError: tiktoken does not know ``encoding_name``. + """ + + def __init__(self, encoding_name: str) -> None: + self._encoding_name = encoding_name + self._encoding = _load_tiktoken_encoding(encoding_name) + + @property + def name(self) -> str: + return self._encoding_name + + def count_tokens(self, text: str) -> int: + if not text: + return 0 + # Feature values are arbitrary text: count a literal "<|endoftext|>" + # rather than raise, which is tiktoken's default for special tokens. + return len(self._encoding.encode(text, disallowed_special=())) + + +class Cl100kBaseTokenizer(TiktokenTokenizer): + """cl100k_base, in the no-argument form that resolution by name needs.""" + + def __init__(self) -> None: + super().__init__(TokenizerName.CL100K_BASE.value) + + +class O200kBaseTokenizer(TiktokenTokenizer): + """o200k_base, in the no-argument form that resolution by name needs.""" + + def __init__(self) -> None: + super().__init__(TokenizerName.O200K_BASE.value) + + +def _load_tiktoken_encoding(encoding_name: str) -> Any: + """Load a tiktoken encoding. tiktoken memoizes these internally.""" + try: + import tiktoken + except ImportError as e: + raise TokenizerUnavailableError( + encoding_name, f"tiktoken is not installed ({e})", TIKTOKEN_INSTALL_HINT + ) from e + + try: + return tiktoken.get_encoding(encoding_name) + except ValueError as e: + raise TokenizerNotFoundError( + encoding_name, tiktoken.list_encoding_names() + ) from e + except Exception as e: + # Encodings download on first use; an offline host fails here. + raise TokenizerUnavailableError( + encoding_name, + f"tiktoken failed to load the encoding ({e})", + TIKTOKEN_INSTALL_HINT, + ) from e + + +#: What callers may pass wherever a tokenizer is expected. +TokenizerSpec = Union[str, TokenizerName, Tokenizer] + + +def get_tokenizer( + spec: TokenizerSpec = TokenizerName.CL100K_BASE, + *, + fallback: bool = True, +) -> Tokenizer: + """Resolve ``spec`` to a tokenizer instance. + + A :class:`Tokenizer` passes through; a name is built once, then cached. + + Args: + spec: Tokenizer instance, built-in name, or the class path of a + ``Tokenizer`` subclass with a no-argument constructor. + fallback: On missing dependencies, return + :class:`ApproximateTokenizer` with a warning instead of raising. + Set False when exact counts are required. + + Raises: + TokenizerNotFoundError: ``spec`` is neither a built-in name nor a + class path ending in ``Tokenizer``. + TokenizerUnavailableError: the tokenizer cannot be built and + ``fallback`` is False. + """ + if isinstance(spec, Tokenizer): + return spec + name = spec.value if isinstance(spec, TokenizerName) else spec + if not isinstance(name, str) or not name.strip(): + raise TypeError(f"Tokenizer name must be a non-empty string, got {spec!r}.") + + try: + return _build_tokenizer(name.strip()) + except TokenizerUnavailableError as e: + if not fallback: + raise + logger.warning( + "%s Falling back to a character-based estimate; token counts will " + "be approximate.", + e, + ) + # Through the cache, so budgets that fell back still compare equal. + return _build_tokenizer(TokenizerName.APPROXIMATE.value) + + +@lru_cache(maxsize=None) +def _build_tokenizer(name: str) -> Tokenizer: + """Import and instantiate the tokenizer ``name`` refers to.""" + tokenizer_type = TOKENIZER_CLASS_FOR_TYPE.get(name, name) + if "." not in tokenizer_type or not tokenizer_type.endswith("Tokenizer"): + raise TokenizerNotFoundError(name, TOKENIZER_CLASS_FOR_TYPE.keys()) + + module_name, class_name = tokenizer_type.rsplit(".", 1) + tokenizer = import_class(module_name, class_name, "Tokenizer")() + if not isinstance(tokenizer, Tokenizer): + # import_class checks the base class by name, so an unrelated class + # called "Tokenizer" gets this far. + raise FeastInvalidBaseClass(tokenizer_type, "Tokenizer") + return tokenizer diff --git a/sdk/python/tests/unit/context/__init__.py b/sdk/python/tests/unit/context/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sdk/python/tests/unit/context/conftest.py b/sdk/python/tests/unit/context/conftest.py new file mode 100644 index 00000000000..2bed5b15656 --- /dev/null +++ b/sdk/python/tests/unit/context/conftest.py @@ -0,0 +1,22 @@ +import builtins + +import pytest + +from feast.context import tokenizer as tokenizer_module + + +@pytest.fixture +def no_tiktoken(monkeypatch): + """Make ``import tiktoken`` fail, as it does when it is not installed.""" + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "tiktoken": + raise ImportError("No module named 'tiktoken'") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + # Drop tokenizers an earlier test may have built and cached. + tokenizer_module._build_tokenizer.cache_clear() + yield + tokenizer_module._build_tokenizer.cache_clear() diff --git a/sdk/python/tests/unit/context/test_token_budget.py b/sdk/python/tests/unit/context/test_token_budget.py new file mode 100644 index 00000000000..9480c488b53 --- /dev/null +++ b/sdk/python/tests/unit/context/test_token_budget.py @@ -0,0 +1,200 @@ +import dataclasses +import inspect + +import pytest + +from feast.context import TokenBudget +from feast.context.errors import TokenBudgetExceededError, TokenizerNotFoundError +from feast.context.tokenizer import ( + ApproximateTokenizer, + Tokenizer, + TokenizerName, + get_tokenizer, +) + +# Four characters per token: "abcd" is one token under the estimate. +ESTIMATING = ApproximateTokenizer() + + +class HalvingTokenizer(Tokenizer): + """A third-party tokenizer, resolved by class path.""" + + @property + def name(self) -> str: + return "halving" + + def count_tokens(self, text: str) -> int: + return len(text) // 2 + + +def budget(max_tokens: int = 100, consumed: int = 0) -> TokenBudget: + """A budget on the deterministic estimate, so counts need no tiktoken.""" + return TokenBudget(max_tokens=max_tokens, tokenizer=ESTIMATING, consumed=consumed) + + +class TestConstruction: + def test_defaults_to_cl100k_base(self): + field = TokenBudget.__dataclass_fields__["tokenizer"] + assert field.default_factory is get_tokenizer + assert ( + inspect.signature(get_tokenizer).parameters["spec"].default + is TokenizerName.CL100K_BASE + ) + + def test_accepts_a_tokenizer_name(self): + assert TokenBudget.of(10, "approximate").tokenizer is get_tokenizer( + "approximate" + ) + + def test_accepts_a_tokenizer_instance(self): + assert TokenBudget(max_tokens=10, tokenizer=ESTIMATING).tokenizer is ESTIMATING + + def test_resolves_a_name_passed_to_the_constructor(self): + # Config supplies a string; the field still holds a Tokenizer. + resolved = TokenBudget(max_tokens=10, tokenizer="approximate").tokenizer # type: ignore[arg-type] + assert resolved is get_tokenizer("approximate") + + def test_resolves_a_custom_tokenizer_by_class_path(self): + b = TokenBudget.of(10, "tests.unit.context.test_token_budget.HalvingTokenizer") + assert b.tokenizer_name == "halving" + assert b.count_tokens("abcdefgh") == 4 + assert b.consume("abcdefgh").remaining == 6 + + def test_unknown_tokenizer_raises(self): + with pytest.raises(TokenizerNotFoundError): + TokenBudget.of(10, "gpt-42") + + def test_rejects_negative_max_tokens(self): + with pytest.raises(ValueError, match="max_tokens"): + budget(max_tokens=-1) + + def test_rejects_negative_consumed(self): + with pytest.raises(ValueError, match="consumed"): + budget(consumed=-1) + + def test_rejects_consumed_above_max_tokens(self): + with pytest.raises(ValueError, match="must not exceed"): + budget(max_tokens=10, consumed=11) + + def test_is_immutable(self): + with pytest.raises(dataclasses.FrozenInstanceError): + budget().max_tokens = 5 # type: ignore[misc] + + def test_budgets_with_the_same_state_are_equal_and_hash_alike(self): + # Default-constructed, so a tokenizer compared by identity would fail. + assert TokenBudget(max_tokens=10) == TokenBudget(max_tokens=10) + assert len({TokenBudget(max_tokens=10), TokenBudget(max_tokens=10)}) == 1 + + def test_replace_keeps_the_resolved_tokenizer(self): + replaced = dataclasses.replace(budget(max_tokens=10), consumed=4) + assert replaced.tokenizer is ESTIMATING + assert replaced.remaining == 6 + + +class TestCounting: + def test_counts_tokens(self): + assert budget().count_tokens("abcdefgh") == 2 + + def test_empty_string_costs_nothing(self): + assert budget().count_tokens("") == 0 + + def test_is_approximate_flags_the_estimate(self): + assert budget().is_approximate is True + + +class TestFits: + @pytest.mark.parametrize( + "max_tokens, text, expected", + [ + (10, "abcdefgh", True), + (2, "abcdefgh", True), # exactly filling the budget fits + (1, "a" * 40, False), + (0, "", True), # nothing always fits + ], + ) + def test_fits_compares_against_what_remains(self, max_tokens, text, expected): + assert budget(max_tokens=max_tokens).fits(text) is expected + + def test_fits_tokens_compares_against_what_remains(self): + b = budget(max_tokens=10, consumed=8) + assert b.fits_tokens(2) is True + assert b.fits_tokens(3) is False + + def test_fits_tokens_rejects_negative_counts(self): + with pytest.raises(ValueError, match="tokens"): + budget().fits_tokens(-1) + + +class TestConsume: + def test_returns_a_new_budget(self): + original = budget(max_tokens=10) + after = original.consume("abcdefgh") + assert after is not original + assert original.remaining == 10 + assert after.remaining == 8 + assert after.consumed == 2 + + def test_exhausts_the_budget_exactly(self): + after = budget(max_tokens=2).consume("abcdefgh") + assert after.remaining == 0 + assert after.is_exhausted is True + + def test_over_budget_raises_and_leaves_the_original_untouched(self): + b = budget(max_tokens=2) + with pytest.raises(TokenBudgetExceededError, match="only 2 of 2 tokens remain"): + b.consume("a" * 40) + assert b.remaining == 2 + + def test_consume_tokens_charges_a_count_directly(self): + assert budget(max_tokens=10).consume_tokens(4).remaining == 6 + + def test_consume_tokens_rejects_negative_counts(self): + with pytest.raises(ValueError, match="tokens"): + budget().consume_tokens(-1) + + +class TestTryConsume: + def test_returns_the_new_budget_when_content_fits(self): + after = budget(max_tokens=10).try_consume("abcdefgh") + assert after is not None + assert after.remaining == 8 + + def test_returns_none_when_content_does_not_fit(self): + assert budget(max_tokens=2).try_consume("a" * 40) is None + + def test_greedy_selection_keeps_what_fits(self): + b = budget(max_tokens=3) + kept = [] + for section in ["abcd", "a" * 40, "abcdefgh"]: + candidate = b.try_consume(section) + if candidate is not None: + b = candidate + kept.append(section) + assert kept == ["abcd", "abcdefgh"] + assert b.remaining == 0 + + +class TestReset: + def test_clears_consumption_and_keeps_limit_and_tokenizer(self): + after = budget(max_tokens=10).consume("abcdefgh").reset() + assert after.consumed == 0 + assert after == budget(max_tokens=10) + + +class TestTiktokenBudget: + @pytest.mark.parametrize( + "name", [TokenizerName.CL100K_BASE.value, TokenizerName.O200K_BASE.value] + ) + def test_counts_exact_tokens(self, name): + pytest.importorskip("tiktoken") + b = TokenBudget.of(4096, name) + assert b.tokenizer_name == name + assert b.is_approximate is False + assert b.count_tokens("Hello world") == 2 + assert b.consume("Hello world").remaining == 4094 + + def test_the_default_budget_counts_cl100k_base(self): + pytest.importorskip("tiktoken") + assert TokenBudget(max_tokens=4096).tokenizer_name == ( + TokenizerName.CL100K_BASE.value + ) diff --git a/sdk/python/tests/unit/context/test_tokenizer.py b/sdk/python/tests/unit/context/test_tokenizer.py new file mode 100644 index 00000000000..2b909cad65b --- /dev/null +++ b/sdk/python/tests/unit/context/test_tokenizer.py @@ -0,0 +1,141 @@ +import pytest + +from feast.context.errors import TokenizerNotFoundError, TokenizerUnavailableError +from feast.context.tokenizer import ( + ApproximateTokenizer, + Cl100kBaseTokenizer, + O200kBaseTokenizer, + TiktokenTokenizer, + TokenizerName, + get_tokenizer, +) +from feast.errors import FeastInvalidBaseClass + + +class FakeTokenizer: + """Named like a tokenizer, but not a Tokenizer subclass.""" + + +class TestApproximateTokenizer: + @pytest.mark.parametrize( + "text, expected", + [ + ("abcdefgh", 2), + ("abcde", 2), # partial tokens round up + ("", 0), + ("a" * 100_000, 25_000), + ], + ) + def test_counts_four_characters_per_token(self, text, expected): + assert ApproximateTokenizer().count_tokens(text) == expected + + def test_chars_per_token_is_configurable(self): + assert ApproximateTokenizer(chars_per_token=2).count_tokens("abcd") == 2 + + @pytest.mark.parametrize("chars_per_token", [0, -1]) + def test_rejects_non_positive_chars_per_token(self, chars_per_token): + with pytest.raises(ValueError): + ApproximateTokenizer(chars_per_token=chars_per_token) + + +class TestTiktokenTokenizer: + @pytest.mark.parametrize( + "encoding", [TokenizerName.CL100K_BASE.value, TokenizerName.O200K_BASE.value] + ) + def test_counts_tokens_for_supported_encodings(self, encoding): + pytest.importorskip("tiktoken") + tokenizer = TiktokenTokenizer(encoding) + assert tokenizer.name == encoding + assert tokenizer.count_tokens("Hello world") == 2 + assert tokenizer.count_tokens("") == 0 + # Unicode and long text: exact counts are the encoding's business, but + # neither may crash or come back empty. + assert tokenizer.count_tokens("北京の天気") > 0 + assert tokenizer.count_tokens("word " * 10_000) >= 10_000 + + def test_special_token_text_is_counted_not_rejected(self): + pytest.importorskip("tiktoken") + tokenizer = Cl100kBaseTokenizer() + assert tokenizer.count_tokens("<|endoftext|> appears in this review") > 0 + + def test_unknown_encoding_raises(self): + pytest.importorskip("tiktoken") + with pytest.raises(TokenizerNotFoundError): + TiktokenTokenizer("no_such_encoding") + + +class TestEquality: + def test_same_name_and_settings_are_equal(self): + assert ApproximateTokenizer() == ApproximateTokenizer() + assert len({ApproximateTokenizer(), ApproximateTokenizer()}) == 1 + + def test_different_settings_are_not_equal(self): + assert ApproximateTokenizer() != ApproximateTokenizer(chars_per_token=3) + + def test_same_encoding_from_different_classes_is_equal(self): + pytest.importorskip("tiktoken") + assert Cl100kBaseTokenizer() == TiktokenTokenizer( + TokenizerName.CL100K_BASE.value + ) + assert Cl100kBaseTokenizer() != O200kBaseTokenizer() + + +class TestGetTokenizer: + @pytest.mark.parametrize( + "name, expected_class", + [ + (TokenizerName.CL100K_BASE.value, Cl100kBaseTokenizer), + (TokenizerName.O200K_BASE.value, O200kBaseTokenizer), + (TokenizerName.APPROXIMATE.value, ApproximateTokenizer), + ], + ) + def test_resolves_builtin_names(self, name, expected_class): + if expected_class is not ApproximateTokenizer: + pytest.importorskip("tiktoken") + assert isinstance(get_tokenizer(name), expected_class) + + def test_resolves_a_class_path(self): + tokenizer = get_tokenizer("feast.context.tokenizer.ApproximateTokenizer") + assert isinstance(tokenizer, ApproximateTokenizer) + + def test_rejects_a_class_path_that_is_not_a_tokenizer(self): + with pytest.raises(FeastInvalidBaseClass): + get_tokenizer("tests.unit.context.test_tokenizer.FakeTokenizer") + + def test_rejects_a_path_that_is_not_a_tokenizer_class_name(self): + with pytest.raises(TokenizerNotFoundError): + get_tokenizer("feast.repo_config.RepoConfig") + + def test_unknown_name_raises_with_builtin_names(self): + with pytest.raises(TokenizerNotFoundError, match="cl100k_base"): + get_tokenizer("gpt-42") + + def test_builds_once_and_caches(self): + assert get_tokenizer("approximate") is get_tokenizer(" approximate ") + + def test_accepts_enum_member_and_equivalent_string(self): + assert get_tokenizer(TokenizerName.APPROXIMATE) is get_tokenizer("approximate") + + def test_tokenizer_instance_passes_through(self): + tokenizer = ApproximateTokenizer() + assert get_tokenizer(tokenizer) is tokenizer + + @pytest.mark.parametrize("name", [" ", None]) + def test_rejects_invalid_names(self, name): + with pytest.raises(TypeError): + get_tokenizer(name) + + +class TestFallback: + def test_falls_back_to_estimate_without_tiktoken(self, no_tiktoken, caplog): + tokenizer = get_tokenizer(TokenizerName.CL100K_BASE) + + assert isinstance(tokenizer, ApproximateTokenizer) + assert tokenizer.count_tokens("abcdefgh") == 2 + # The shared instance, so budgets that fell back stay comparable. + assert tokenizer is get_tokenizer(TokenizerName.APPROXIMATE) + assert any("tiktoken" in record.message for record in caplog.records) + + def test_fallback_disabled_raises(self, no_tiktoken): + with pytest.raises(TokenizerUnavailableError, match="pip install tiktoken"): + get_tokenizer(TokenizerName.CL100K_BASE, fallback=False) From 9bca2fba372a3b7ec7d7c080ed61b5f48254b50e Mon Sep 17 00:00:00 2001 From: Chaitany Patel Date: Thu, 10 Sep 2026 16:32:20 +0530 Subject: [PATCH 2/2] feat(context): Add priority_select for budget-aware section selection Decides which parts of a prompt survive a token budget. Sections carry a priority; selection keeps them from critical down until the allowance is spent, dropping the least important first and never dropping a critical one. The result carries the kept sections, the dropped ones and the budget left over, and renders to the assembled text. The budget is charged for the separator between kept sections, so the rendered prompt fits the allowance rather than overshooting it once the pieces are joined. Selection is a pure function of the sections and the budget, so an OnDemandFeatureView drops the same sections offline and online. Part of RHOAIENG-80116. Stacked on the TokenBudget PR. Signed-off-by: Chaitany Patel --- sdk/python/feast/context/__init__.py | 20 ++ sdk/python/feast/context/errors.py | 37 ++- sdk/python/feast/context/priority_select.py | 284 ++++++++++++++++ .../unit/context/test_priority_select.py | 307 ++++++++++++++++++ 4 files changed, 645 insertions(+), 3 deletions(-) create mode 100644 sdk/python/feast/context/priority_select.py create mode 100644 sdk/python/tests/unit/context/test_priority_select.py diff --git a/sdk/python/feast/context/__init__.py b/sdk/python/feast/context/__init__.py index 62de13ebead..71a23b5d7e7 100644 --- a/sdk/python/feast/context/__init__.py +++ b/sdk/python/feast/context/__init__.py @@ -6,10 +6,21 @@ from feast.context.errors import ( ContextError, + RequiredSectionError, TokenBudgetExceededError, TokenizerNotFoundError, TokenizerUnavailableError, ) +from feast.context.priority_select import ( + DEFAULT_SEPARATOR, + Priority, + PrioritySpec, + Section, + SectionOrder, + SectionSpec, + Selection, + priority_select, +) from feast.context.token_budget import TokenBudget from feast.context.tokenizer import ( ApproximateTokenizer, @@ -26,7 +37,15 @@ "ApproximateTokenizer", "Cl100kBaseTokenizer", "ContextError", + "DEFAULT_SEPARATOR", "O200kBaseTokenizer", + "Priority", + "PrioritySpec", + "RequiredSectionError", + "Section", + "SectionOrder", + "SectionSpec", + "Selection", "TiktokenTokenizer", "TokenBudget", "TokenBudgetExceededError", @@ -36,4 +55,5 @@ "TokenizerSpec", "TokenizerUnavailableError", "get_tokenizer", + "priority_select", ] diff --git a/sdk/python/feast/context/errors.py b/sdk/python/feast/context/errors.py index 2504eabebd9..b96ff24c82e 100644 --- a/sdk/python/feast/context/errors.py +++ b/sdk/python/feast/context/errors.py @@ -40,10 +40,20 @@ def http_status_code(self) -> int: class TokenBudgetExceededError(ContextError): - def __init__(self, requested: int, remaining: int, max_tokens: int): + def __init__( + self, + requested: int, + remaining: int, + max_tokens: int, + *, + message: str = "", + ): super().__init__( - f"Cannot consume {requested} tokens: only {remaining} of " - f"{max_tokens} tokens remain in the budget." + message + or ( + f"Cannot consume {requested} tokens: only {remaining} of " + f"{max_tokens} tokens remain in the budget." + ) ) self.requested = requested self.remaining = remaining @@ -51,3 +61,24 @@ def __init__(self, requested: int, remaining: int, max_tokens: int): def http_status_code(self) -> int: return HttpStatusCode.HTTP_400_BAD_REQUEST + + +class RequiredSectionError(TokenBudgetExceededError): + """A section that must not be dropped does not fit the budget. + + Raised instead of quietly returning a prompt missing its instructions: + the caller has to widen the budget or shorten the section. + """ + + def __init__(self, section: str, requested: int, remaining: int, max_tokens: int): + super().__init__( + requested, + remaining, + max_tokens, + message=( + f"Section {section} is critical and cannot be dropped, but it " + f"needs {requested} tokens with only {remaining} of " + f"{max_tokens} left in the budget." + ), + ) + self.section = section diff --git a/sdk/python/feast/context/priority_select.py b/sdk/python/feast/context/priority_select.py new file mode 100644 index 00000000000..5d9e3bee329 --- /dev/null +++ b/sdk/python/feast/context/priority_select.py @@ -0,0 +1,284 @@ +"""Choosing what survives when a prompt does not fit its budget. + +Prompt assembly is a packing problem: a user profile, a watch history and a +catalogue of candidates rarely fit together in a context window, and something +has to give. :func:`priority_select` decides what, keeping sections from the +most important down until the budget is spent:: + + selection = priority_select( + [ + ("critical", system_instructions), + ("high", f"User profile: {profile}"), + ("medium", f"Recent history: {history}"), + ("low", f"Device: {device}"), + ], + TokenBudget.of(4096), + ) + prompt = selection.render() + +The decision is deterministic in the sections and the budget alone, so an +OnDemandFeatureView that calls it drops the same sections during training as it +does at serving time. +""" + +from dataclasses import dataclass +from enum import Enum, IntEnum +from typing import Iterable, Iterator, Optional, Union + +from feast.context.errors import RequiredSectionError +from feast.context.token_budget import TokenBudget + +#: Placed between kept sections when a selection is rendered. +DEFAULT_SEPARATOR = "\n\n" + +#: Characters of content quoted in messages about an unnamed section. +_LABEL_PREVIEW_CHARS = 40 + + +class Priority(IntEnum): + """How readily a section may be dropped to make room for another. + + Ordered, so ``Priority.HIGH > Priority.LOW``, and spaced by ten to leave + room for levels a caller wants in between:: + + Priority.HIGH - 1 # sits between medium and high + """ + + #: Nice to have: dropped first. + LOW = 10 + #: Useful context, expendable under pressure. + MEDIUM = 20 + #: Kept unless the budget is genuinely tight. + HIGH = 30 + #: Never dropped; the selection fails instead. + CRITICAL = 40 + + @classmethod + def parse(cls, spec: "PrioritySpec") -> "Priority": + """Resolve a priority or its case-insensitive name.""" + if isinstance(spec, Priority): + return spec + if isinstance(spec, str): + try: + return cls[spec.strip().upper()] + except KeyError: + pass + raise ValueError( + f"Unknown priority {spec!r}. Use one of: " + f"{', '.join(level.name.lower() for level in sorted(cls, reverse=True))}." + ) + + def __str__(self) -> str: + return self.name.lower() + + +#: What callers may pass wherever a priority is expected. +PrioritySpec = Union[str, Priority] + + +@dataclass(frozen=True) +class Section: + """A candidate piece of a prompt, and how hard it is to lose. + + Attributes: + content: The text itself. Blank content is ignored by selection. + priority: How readily the section is dropped. + name: Optional identifier, quoted in error messages and useful for + asserting on what a selection kept. + """ + + content: str + priority: Priority = Priority.MEDIUM + name: Optional[str] = None + + def __post_init__(self) -> None: + if not isinstance(self.content, str): + raise TypeError(f"Section content must be a string, got {self.content!r}.") + if not isinstance(self.priority, Priority): + # A name from config lands here; resolve it once, so comparisons + # and sorting see a Priority. + object.__setattr__(self, "priority", Priority.parse(self.priority)) + + @classmethod + def of(cls, spec: "SectionSpec") -> "Section": + """Normalize a section or a ``(priority, content[, name])`` tuple.""" + if isinstance(spec, Section): + return spec + if isinstance(spec, tuple) and len(spec) in (2, 3): + priority, content, *rest = spec + return cls( + content=content, + priority=Priority.parse(priority), + name=rest[0] if rest else None, + ) + raise TypeError( + f"Expected a Section or a (priority, content) tuple, got {spec!r}." + ) + + @property + def is_required(self) -> bool: + """Whether the section must be kept, budget notwithstanding.""" + return self.priority is Priority.CRITICAL + + @property + def label(self) -> str: + """Quoted identifier for messages: the name, else a content preview.""" + if self.name: + return repr(self.name) + flattened = " ".join(self.content.split()) + if len(flattened) > _LABEL_PREVIEW_CHARS: + return repr(f"{flattened[: _LABEL_PREVIEW_CHARS - 3]}...") + return repr(flattened) + + +#: A section, or the tuple shorthand for one. +SectionSpec = Union[ + Section, + tuple[PrioritySpec, str], + tuple[PrioritySpec, str, Optional[str]], +] + + +class SectionOrder(str, Enum): + """The order kept sections come back in.""" + + #: As the caller listed them, so the prompt reads as it was written. + DECLARED = "declared" + #: Most important first, ties broken by declaration order. + PRIORITY = "priority" + + +@dataclass(frozen=True) +class Selection: + """The outcome of a :func:`priority_select`: what fit and what did not. + + Iterating yields the kept content, so a selection drops straight into a + join or a template:: + + for text in selection: ... + prompt = selection.render() + + Attributes: + budget: The budget after charging every kept section. + kept: Sections that made it in, ordered as requested. + dropped: Sections left out, in declaration order. + separator: What :meth:`render` joins kept sections with. + """ + + budget: TokenBudget + kept: tuple[Section, ...] = () + dropped: tuple[Section, ...] = () + separator: str = DEFAULT_SEPARATOR + + @property + def contents(self) -> tuple[str, ...]: + """The text of each kept section.""" + return tuple(section.content for section in self.kept) + + @property + def names(self) -> tuple[Optional[str], ...]: + """The name of each kept section, None where unnamed.""" + return tuple(section.name for section in self.kept) + + @property + def is_complete(self) -> bool: + """Whether everything offered was kept.""" + return not self.dropped + + def render(self) -> str: + """The kept sections joined by :attr:`separator`.""" + return self.separator.join(self.contents) + + def __iter__(self) -> Iterator[str]: + return iter(self.contents) + + def __len__(self) -> int: + return len(self.kept) + + def __str__(self) -> str: + return self.render() + + +def priority_select( + sections: Iterable[SectionSpec], + budget: TokenBudget, + *, + separator: str = DEFAULT_SEPARATOR, + order: SectionOrder = SectionOrder.DECLARED, +) -> Selection: + """Keep as much of ``sections`` as ``budget`` allows, least important out first. + + Sections are considered from the highest priority down, ties in declaration + order, and each is kept if what remains of the budget covers it. A section + too large to fit is dropped and the scan continues, so a small low-priority + section can still be kept after a large high-priority one was refused. + + Sections whose content is blank are ignored: they would only double a + separator. ``critical`` sections are never dropped. + + The budget is charged for ``separator`` between each pair of kept sections, + so :meth:`Selection.render` fits within the allowance. Tokenizers may merge + text across a section boundary, which makes the charge an upper bound on + what the joined string really costs. + + Args: + sections: Sections, or ``(priority, content[, name])`` tuples. + budget: The allowance to fill. Anything already consumed on it stays + consumed; only what remains is available here. + separator: Placed between kept sections by + :meth:`Selection.render`, and charged against the budget. + order: Whether kept sections come back in declaration or priority + order. + + Returns: + A :class:`Selection` holding the kept sections, the dropped ones, and + the budget left over. + + Raises: + RequiredSectionError: A ``critical`` section does not fit. + TypeError: A section is neither a :class:`Section` nor a tuple. + ValueError: A priority or an order is not a known name. + """ + if not isinstance(separator, str): + raise TypeError(f"separator must be a string, got {separator!r}.") + order = SectionOrder(order) + + candidates = [ + section + for section in (Section.of(spec) for spec in sections) + if section.content.strip() + ] + separator_tokens = budget.count_tokens(separator) + + remaining = budget + kept: list[tuple[int, Section]] = [] + dropped: list[tuple[int, Section]] = [] + # Most important first, and among equals whoever was listed first. + by_importance = sorted( + enumerate(candidates), key=lambda pair: (-pair[1].priority, pair[0]) + ) + + for position, section in by_importance: + cost = remaining.count_tokens(section.content) + if kept: + cost += separator_tokens + if remaining.fits_tokens(cost): + remaining = remaining.consume_tokens(cost) + kept.append((position, section)) + elif section.is_required: + raise RequiredSectionError( + section.label, cost, remaining.remaining, remaining.max_tokens + ) + else: + dropped.append((position, section)) + + if order is SectionOrder.DECLARED: + kept.sort(key=lambda pair: pair[0]) + dropped.sort(key=lambda pair: pair[0]) + + return Selection( + budget=remaining, + kept=tuple(section for _, section in kept), + dropped=tuple(section for _, section in dropped), + separator=separator, + ) diff --git a/sdk/python/tests/unit/context/test_priority_select.py b/sdk/python/tests/unit/context/test_priority_select.py new file mode 100644 index 00000000000..0ec01a2e7b2 --- /dev/null +++ b/sdk/python/tests/unit/context/test_priority_select.py @@ -0,0 +1,307 @@ +import dataclasses + +import pytest + +from feast.context import ( + Priority, + Section, + SectionOrder, + TokenBudget, + priority_select, +) +from feast.context.errors import RequiredSectionError, TokenBudgetExceededError +from feast.context.tokenizer import ApproximateTokenizer + +# Four characters per token: "abcd" is one token under the estimate. +ESTIMATING = ApproximateTokenizer() + + +def budget(max_tokens: int = 100, consumed: int = 0) -> TokenBudget: + """A budget on the deterministic estimate, so counts need no tiktoken.""" + return TokenBudget(max_tokens=max_tokens, tokenizer=ESTIMATING, consumed=consumed) + + +def text(tokens: int, letter: str = "a") -> str: + """Content costing exactly ``tokens`` tokens under the estimate.""" + return letter * (tokens * 4) + + +class TestPriority: + def test_orders_from_critical_down(self): + assert Priority.CRITICAL > Priority.HIGH > Priority.MEDIUM > Priority.LOW + + def test_parses_a_name_whatever_its_case(self): + assert Priority.parse("critical") is Priority.CRITICAL + assert Priority.parse(" High ") is Priority.HIGH + + def test_passes_a_priority_through(self): + assert Priority.parse(Priority.LOW) is Priority.LOW + + def test_unknown_priority_lists_the_known_ones(self): + with pytest.raises(ValueError, match="critical, high, medium, low"): + Priority.parse("urgent") + + def test_reads_as_its_lowercase_name(self): + assert f"{Priority.MEDIUM}" == "medium" + + def test_leaves_room_between_levels(self): + assert Priority.MEDIUM < Priority.HIGH - 1 < Priority.HIGH + + +class TestSection: + def test_defaults_to_medium(self): + assert Section("body").priority is Priority.MEDIUM + + def test_resolves_a_priority_name(self): + assert Section("body", "high").priority is Priority.HIGH # type: ignore[arg-type] + + def test_rejects_non_string_content(self): + with pytest.raises(TypeError, match="content must be a string"): + Section(42) # type: ignore[arg-type] + + def test_is_immutable(self): + with pytest.raises(dataclasses.FrozenInstanceError): + Section("body").content = "other" # type: ignore[misc] + + def test_of_builds_from_a_priority_and_content(self): + assert Section.of(("low", "body")) == Section("body", Priority.LOW) + + def test_of_builds_from_a_named_triple(self): + section = Section.of(("high", "body", "profile")) + assert section == Section("body", Priority.HIGH, "profile") + + def test_of_passes_a_section_through(self): + section = Section("body") + assert Section.of(section) is section + + def test_of_rejects_anything_else(self): + with pytest.raises(TypeError, match="Section or a"): + Section.of("body") # type: ignore[arg-type] + + def test_only_critical_sections_are_required(self): + assert Section("body", Priority.CRITICAL).is_required is True + assert Section("body", Priority.HIGH).is_required is False + + def test_label_prefers_the_name(self): + assert Section("body", name="profile").label == "'profile'" + + def test_label_falls_back_to_a_flattened_preview(self): + assert Section("two\n words").label == "'two words'" + + def test_label_truncates_a_long_preview(self): + label = Section("word " * 20).label + assert label.endswith("...'") + assert len(label) == len("''") + 40 + + +class TestSelects: + def test_keeps_everything_that_fits(self): + selection = priority_select([("high", text(2)), ("low", text(2))], budget(10)) + assert selection.contents == (text(2), text(2)) + assert selection.is_complete is True + assert selection.dropped == () + + def test_keeps_declaration_order_by_default(self): + selection = priority_select( + [("low", "third"), ("critical", "first"), ("medium", "second")], + budget(100), + ) + assert selection.contents == ("third", "first", "second") + + def test_orders_by_priority_when_asked(self): + selection = priority_select( + [("low", "third"), ("critical", "first"), ("medium", "second")], + budget(100), + order=SectionOrder.PRIORITY, + ) + assert selection.contents == ("first", "second", "third") + + def test_accepts_an_order_by_name(self): + selection = priority_select( + [("low", "second"), ("high", "first")], budget(100), order="priority" + ) + assert selection.contents == ("first", "second") + + def test_rejects_an_unknown_order(self): + with pytest.raises(ValueError, match="alphabetical"): + priority_select([], budget(), order="alphabetical") + + def test_drops_the_least_important_first(self): + selection = priority_select( + [ + ("high", text(2, "h")), + ("medium", text(2, "m")), + ("low", text(2, "l")), + ], + budget(4), + separator="", + ) + assert selection.contents == (text(2, "h"), text(2, "m")) + assert selection.dropped == (Section(text(2, "l"), Priority.LOW),) + assert selection.is_complete is False + + def test_keeps_a_critical_section_listed_last(self): + selection = priority_select( + [("high", text(2, "h")), ("critical", text(2, "c"))], + budget(2), + separator="", + ) + assert selection.contents == (text(2, "c"),) + assert selection.dropped == (Section(text(2, "h"), Priority.HIGH),) + + def test_breaks_ties_by_declaration_order(self): + selection = priority_select( + [("medium", text(2, "a")), ("medium", text(2, "b"))], + budget(2), + separator="", + ) + assert selection.contents == (text(2, "a"),) + + def test_keeps_a_smaller_section_after_refusing_a_larger_one(self): + selection = priority_select( + [("high", text(5, "h")), ("medium", text(2, "m")), ("low", text(1, "l"))], + budget(3), + separator="", + ) + assert selection.contents == (text(2, "m"), text(1, "l")) + assert selection.budget.is_exhausted is True + + def test_ignores_blank_sections(self): + selection = priority_select( + [("critical", " \n "), ("high", "body"), ("low", "")], budget(100) + ) + assert selection.contents == ("body",) + assert selection.dropped == () + + def test_selects_nothing_from_nothing(self): + selection = priority_select([], budget(10)) + assert selection.contents == () + assert selection.is_complete is True + assert selection.budget == budget(10) + + def test_accepts_sections_as_well_as_tuples(self): + selection = priority_select( + [Section("body", Priority.HIGH, "profile"), ("low", "extra")], budget(100) + ) + assert selection.names == ("profile", None) + + def test_leaves_the_original_budget_untouched(self): + original = budget(10) + priority_select([("high", text(4))], original, separator="") + assert original.consumed == 0 + + +class TestBudgetAccounting: + def test_charges_the_kept_sections(self): + selection = priority_select( + [("high", text(3)), ("low", text(2))], budget(10), separator="" + ) + assert selection.budget.consumed == 5 + assert selection.budget.remaining == 5 + + def test_spends_only_what_is_left_of_a_used_budget(self): + selection = priority_select( + [("high", text(3)), ("low", text(3))], + budget(10, consumed=6), + separator="", + ) + assert selection.contents == (text(3),) + assert selection.budget.remaining == 1 + + def test_charges_one_separator_between_kept_sections(self): + # "\n\n" is one token under the estimate: three sections cost two. + selection = priority_select( + [("high", text(3)), ("high", text(3)), ("high", text(3))], budget(100) + ) + assert selection.budget.consumed == 3 * 3 + 2 + + def test_charges_nothing_for_an_empty_separator(self): + selection = priority_select( + [("high", text(3)), ("high", text(3))], budget(100), separator="" + ) + assert selection.budget.consumed == 6 + + def test_the_rendered_prompt_fits_the_budget(self): + selection = priority_select( + [("high", text(3, "h")), ("medium", text(3, "m")), ("low", text(3, "l"))], + budget(10), + ) + rendered = selection.render() + assert selection.contents == (text(3, "h"), text(3, "m")) + assert rendered == f"{text(3, 'h')}\n\n{text(3, 'm')}" + assert ESTIMATING.count_tokens(rendered) <= 10 + assert selection.budget.consumed == ESTIMATING.count_tokens(rendered) + + def test_rejects_a_non_string_separator(self): + with pytest.raises(TypeError, match="separator must be a string"): + priority_select([], budget(), separator=None) # type: ignore[arg-type] + + +class TestRequiredSections: + def test_raises_when_a_critical_section_does_not_fit(self): + with pytest.raises(RequiredSectionError, match="'rules' is critical"): + priority_select([("critical", text(5), "rules")], budget(4), separator="") + + def test_raises_when_critical_sections_together_overflow(self): + with pytest.raises(RequiredSectionError, match="'second'"): + priority_select( + [("critical", text(3), "first"), ("critical", text(3), "second")], + budget(5), + separator="", + ) + + def test_reports_the_shortfall(self): + with pytest.raises(RequiredSectionError) as excinfo: + priority_select([("critical", text(5))], budget(4), separator="") + error = excinfo.value + assert error.requested == 5 + assert error.remaining == 4 + assert error.max_tokens == 4 + assert error.section == f"'{text(5)}'" + + def test_is_a_token_budget_error(self): + # Callers already guarding prompt assembly keep catching one type. + assert issubclass(RequiredSectionError, TokenBudgetExceededError) + + +class TestSelection: + @pytest.fixture + def selection(self): + return priority_select( + [("high", "first", "profile"), ("low", "second")], budget(100) + ) + + def test_renders_sections_joined_by_the_separator(self, selection): + assert selection.render() == "first\n\nsecond" + + def test_renders_as_its_string(self, selection): + assert f"{selection}" == "first\n\nsecond" + + def test_iterates_over_the_kept_content(self, selection): + assert list(selection) == ["first", "second"] + assert " | ".join(selection) == "first | second" + + def test_counts_the_kept_sections(self, selection): + assert len(selection) == 2 + + def test_exposes_the_kept_names(self, selection): + assert selection.names == ("profile", None) + + def test_is_immutable(self, selection): + with pytest.raises(dataclasses.FrozenInstanceError): + selection.kept = () + + +class TestTiktokenSelection: + def test_the_rendered_prompt_fits_an_exact_budget(self): + pytest.importorskip("tiktoken") + sections = [ + ("critical", "You are a helpful recommendation assistant."), + ("high", "The user enjoys slow cinema and long documentaries."), + ("medium", "Recently watched: " + "an eight part nature series, " * 20), + ("low", "Device: living room television. Time: Sunday evening."), + ] + selection = priority_select(sections, TokenBudget.of(64)) + assert selection.is_complete is False + assert selection.budget.count_tokens(selection.render()) <= 64 + assert selection.contents[0] == sections[0][1]