From 3feb5f27dfb5e082e52d4ed7b60ea8e0079592d7 Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Fri, 9 Oct 2026 17:23:14 -0400 Subject: [PATCH 1/2] fix(token_counter): bound the cost of pricing a base64 PDF Pricing an inline PDF parsed and text-extracted every page on the token count path, so its cost grew with page count and, superlinearly, with a single page's content stream. A small compressed upload could hold one of the four token-count slots for seconds to minutes, stalling every other request that needs a count on that worker Text is now extracted from at most EXTRAPOLATION_SAMPLES evenly spaced pages and scaled to the page count, and a page whose content stream exceeds PDF_PAGE_MAX_EXACT_CONTENT_BYTES is priced at one token per content byte instead of being extracted. Every page's rendering tokens are still counted --- litellm/litellm_core_utils/token_counter.py | 55 +++++++++++++++++-- .../litellm_core_utils/test_token_counter.py | 40 +++++++++++++- 2 files changed, 88 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 7ac3e9c80c5..cd737e8e326 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -6,7 +6,7 @@ import math import struct from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from itertools import accumulate -from typing import TYPE_CHECKING, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, Protocol, cast import anyio import anyio.lowlevel @@ -344,6 +344,7 @@ Type for a function that counts tokens in a string. """ EXTRAPOLATION_SAMPLES: Final = 16 +PDF_PAGE_MAX_EXACT_CONTENT_BYTES: Final = 64 * 1024 T_ParamSpec: Final = ParamSpec("T_ParamSpec") T_Retval = TypeVar("T_Retval") _COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") @@ -818,6 +819,43 @@ def _anthropic_rendered_page_image_tokens(width: float, height: float) -> int: return math.ceil(min(float(ANTHROPIC_IMAGE_MAX_PIXELS), area_at_max_edge) / ANTHROPIC_IMAGE_PIXELS_PER_TOKEN) +class _PdfContentStream(Protocol): + def get_data(self) -> bytes: ... + + +class _PdfMediaBox(Protocol): + @property + def width(self) -> float: ... + + @property + def height(self) -> float: ... + + +class _PdfPage(Protocol): + @property + def mediabox(self) -> _PdfMediaBox: ... + + def get_contents(self) -> _PdfContentStream | None: ... + + def extract_text(self) -> str: ... + + +def _evenly_spaced_page_indices(page_count: int) -> tuple[int, ...]: + if page_count <= EXTRAPOLATION_SAMPLES: + return tuple(range(page_count)) + return tuple((page_count - 1) * index // (EXTRAPOLATION_SAMPLES - 1) for index in range(EXTRAPOLATION_SAMPLES)) + + +def _pdf_page_text_tokens(page: _PdfPage, count_function: TokenCounterFunction) -> int: + """A page whose content stream is too large to extract cheaply is priced at one token per content byte, + which no byte-level tokenizer undercuts.""" + contents: Final = page.get_contents() + content_bytes: Final = len(contents.get_data()) if contents is not None else 0 + if content_bytes > PDF_PAGE_MAX_EXACT_CONTENT_BYTES: + return content_bytes + return count_function(page.extract_text() or "") + + def _count_inline_pdf_tokens(data_url: str, count_function: TokenCounterFunction) -> int | None: if not data_url.startswith(PDF_DATA_URL_PREFIX): return None @@ -827,12 +865,17 @@ def _count_inline_pdf_tokens(data_url: str, count_function: TokenCounterFunction verbose_logger.debug("pypdf is not installed, so the PDF document is priced like one image") return None try: - reader: Final = PdfReader(io.BytesIO(base64.b64decode(data_url[len(PDF_DATA_URL_PREFIX) :]))) - return sum( - count_function(page.extract_text() or "") - + _anthropic_rendered_page_image_tokens(float(page.mediabox.width), float(page.mediabox.height)) - for page in reader.pages + pages: Final = cast( # cast-ok: pypdf is an optional extra the type-check env does not install + Sequence[_PdfPage], PdfReader(io.BytesIO(base64.b64decode(data_url[len(PDF_DATA_URL_PREFIX) :]))).pages ) + rendering_tokens: Final = sum( + _anthropic_rendered_page_image_tokens(page.mediabox.width, page.mediabox.height) for page in pages + ) + sampled_pages: Final = tuple(pages[index] for index in _evenly_spaced_page_indices(len(pages))) + if not sampled_pages: + return rendering_tokens + sampled_text_tokens: Final = sum(_pdf_page_text_tokens(page, count_function) for page in sampled_pages) + return rendering_tokens + round(sampled_text_tokens * len(pages) / len(sampled_pages)) except Exception as e: verbose_logger.debug("Could not read the PDF document's pages (%s), so it is priced like one image", e) return None diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index a646135efc3..e18acc00db4 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -27,9 +27,12 @@ import litellm from litellm import decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +from litellm.constants import PDF_DATA_URL_PREFIX, TOKEN_COUNTER_MAX_CONCURRENT_COUNTS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( + EXTRAPOLATION_SAMPLES, + PDF_PAGE_MAX_EXACT_CONTENT_BYTES, + _count_inline_pdf_tokens, _encoding_count, _get_exact_count_function, _get_extrapolating_count_function, @@ -1608,6 +1611,41 @@ def test_pdf_page_rendering_cost_follows_anthropic_image_scaling(fields: dict[st assert strip == 328 +def test_long_pdf_reads_text_from_a_bounded_number_of_pages(): + count_function: Final = MagicMock(side_effect=len) + forty_pages: Final = _pdf_base64(("Revenue grew eleven percent while churn fell to two percent.",) * 40) + + _count_inline_pdf_tokens(PDF_DATA_URL_PREFIX + forty_pages, count_function) + + assert count_function.call_count <= EXTRAPOLATION_SAMPLES + + +def test_dense_pdf_page_is_priced_without_extracting_its_full_text(): + count_function: Final = MagicMock(side_effect=len) + sentence: Final = "Revenue grew eleven percent while churn fell to two percent. " + page_text: Final = sentence * (2 * PDF_PAGE_MAX_EXACT_CONTENT_BYTES // len(sentence)) + + priced: Final = _count_inline_pdf_tokens(PDF_DATA_URL_PREFIX + _pdf_base64((page_text,)), count_function) + + assert priced is not None and priced >= len(page_text) + assert max((len(call.args[0]) for call in count_function.call_args_list), default=0) <= ( + PDF_PAGE_MAX_EXACT_CONTENT_BYTES + ) + + +def test_long_pdf_of_identical_pages_prices_each_page_like_a_one_page_pdf(): + prompt: Final = {"type": "text", "text": "Summarize this file."} + page_text: Final = "Revenue grew eleven percent while churn fell to two percent." + page_count: Final = 3 * EXTRAPOLATION_SAMPLES + 1 + base: Final = _count_user_content([prompt]) + + def pdf_cost(pages: int) -> int: + document: Final = {"type": "document", "source": _pdf_source(_pdf_base64((page_text,) * pages))} + return _count_user_content([prompt, document]) - base + + assert pdf_cost(page_count) == page_count * pdf_cost(1) + + def test_pdf_document_without_pypdf_is_priced_like_an_image(monkeypatch: pytest.MonkeyPatch): prompt: Final = {"type": "text", "text": "Summarize this file."} source: Final = _pdf_source(_pdf_base64(("Revenue grew eleven percent while churn fell to two percent.",))) From c1c3d5e05207665543d4c3806df7986998d656fb Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Fri, 9 Oct 2026 18:15:48 -0400 Subject: [PATCH 2/2] test(token_counter): count pdf text extraction directly with injected pages Pricing a page list now lives in _price_pdf_pages, which takes any object shaped like a pypdf page. The bounded-work tests pass fake pages and count extract_text calls, so extracting every page and sampling afterwards fails them, as does extracting a page before checking its content size --- litellm/litellm_core_utils/token_counter.py | 20 ++++++++------ .../litellm_core_utils/test_token_counter.py | 26 ++++++++++++++++--- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index cd737e8e326..da21e45d67c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -856,6 +856,17 @@ def _pdf_page_text_tokens(page: _PdfPage, count_function: TokenCounterFunction) return count_function(page.extract_text() or "") +def _price_pdf_pages(pages: Sequence[_PdfPage], count_function: TokenCounterFunction) -> int: + rendering_tokens: Final = sum( + _anthropic_rendered_page_image_tokens(page.mediabox.width, page.mediabox.height) for page in pages + ) + sampled_pages: Final = tuple(pages[index] for index in _evenly_spaced_page_indices(len(pages))) + if not sampled_pages: + return rendering_tokens + sampled_text_tokens: Final = sum(_pdf_page_text_tokens(page, count_function) for page in sampled_pages) + return rendering_tokens + round(sampled_text_tokens * len(pages) / len(sampled_pages)) + + def _count_inline_pdf_tokens(data_url: str, count_function: TokenCounterFunction) -> int | None: if not data_url.startswith(PDF_DATA_URL_PREFIX): return None @@ -868,14 +879,7 @@ def _count_inline_pdf_tokens(data_url: str, count_function: TokenCounterFunction pages: Final = cast( # cast-ok: pypdf is an optional extra the type-check env does not install Sequence[_PdfPage], PdfReader(io.BytesIO(base64.b64decode(data_url[len(PDF_DATA_URL_PREFIX) :]))).pages ) - rendering_tokens: Final = sum( - _anthropic_rendered_page_image_tokens(page.mediabox.width, page.mediabox.height) for page in pages - ) - sampled_pages: Final = tuple(pages[index] for index in _evenly_spaced_page_indices(len(pages))) - if not sampled_pages: - return rendering_tokens - sampled_text_tokens: Final = sum(_pdf_page_text_tokens(page, count_function) for page in sampled_pages) - return rendering_tokens + round(sampled_text_tokens * len(pages) / len(sampled_pages)) + return _price_pdf_pages(pages, count_function) except Exception as e: verbose_logger.debug("Could not read the PDF document's pages (%s), so it is priced like one image", e) return None diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index e18acc00db4..a5f19698cca 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -13,6 +13,7 @@ from collections.abc import Mapping from concurrent.futures import Future, wait from itertools import accumulate, chain from pathlib import Path +from types import SimpleNamespace from typing import Final from unittest.mock import MagicMock @@ -34,6 +35,7 @@ from litellm.litellm_core_utils.token_counter import ( PDF_PAGE_MAX_EXACT_CONTENT_BYTES, _count_inline_pdf_tokens, _encoding_count, + _price_pdf_pages, _get_exact_count_function, _get_extrapolating_count_function, _get_tiktoken_count_function, @@ -1611,13 +1613,29 @@ def test_pdf_page_rendering_cost_follows_anthropic_image_scaling(fields: dict[st assert strip == 328 +def _fake_pdf_page(content_bytes: int) -> SimpleNamespace: + return SimpleNamespace( + mediabox=SimpleNamespace(width=612.0, height=792.0), + get_contents=lambda: SimpleNamespace(get_data=lambda: b"x" * content_bytes), + extract_text=MagicMock(return_value="Revenue grew eleven percent while churn fell to two percent."), + ) + + def test_long_pdf_reads_text_from_a_bounded_number_of_pages(): - count_function: Final = MagicMock(side_effect=len) - forty_pages: Final = _pdf_base64(("Revenue grew eleven percent while churn fell to two percent.",) * 40) + pages: Final = tuple(_fake_pdf_page(content_bytes=1_000) for _ in range(40)) - _count_inline_pdf_tokens(PDF_DATA_URL_PREFIX + forty_pages, count_function) + _price_pdf_pages(pages, len) - assert count_function.call_count <= EXTRAPOLATION_SAMPLES + assert sum(page.extract_text.call_count for page in pages) <= EXTRAPOLATION_SAMPLES + + +def test_dense_pdf_page_is_priced_from_its_size_without_extracting_its_text(): + page: Final = _fake_pdf_page(content_bytes=2 * PDF_PAGE_MAX_EXACT_CONTENT_BYTES) + + priced: Final = _price_pdf_pages((page,), len) + + page.extract_text.assert_not_called() + assert priced >= 2 * PDF_PAGE_MAX_EXACT_CONTENT_BYTES def test_dense_pdf_page_is_priced_without_extracting_its_full_text():