Merge c1c3d5e052 into 82ca10cc2d

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,54 @@ 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 _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
@@ -827,12 +876,10 @@ 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
         )
+        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

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
 
@@ -27,10 +28,14 @@ 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,
+    _price_pdf_pages,
     _get_exact_count_function,
     _get_extrapolating_count_function,
     _get_tiktoken_count_function,
@@ -1608,6 +1613,57 @@ 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():
+    pages: Final = tuple(_fake_pdf_page(content_bytes=1_000) for _ in range(40))
+
+    _price_pdf_pages(pages, len)
+
+    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():
+    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.",)))