diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 1c7b5b43fa7..b557eaa13bb 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -1623,6 +1623,13 @@ class _PdfModule(Protocol): PageObject: _PdfPageFactory +@dataclass(frozen=True, slots=True) +class _PdfStreamScan: + content: bytes + operand_bytes: int + invoked_names: tuple[str, ...] + + @dataclass(frozen=True, slots=True) class _PdfStreamMetrics: content: bytes @@ -1677,37 +1684,42 @@ def _pdf_strip_inline_images(content: bytes) -> bytes: return PDF_INLINE_IMAGE.sub(_pdf_inline_image_replacement, content) if b"BI" in content else content +def _pdf_scan_stream(content: bytes) -> _PdfStreamScan: + return _PdfStreamScan( + content=content, + operand_bytes=_pdf_string_operand_bytes(content), + invoked_names=tuple( + name + for invocation in PDF_XOBJECT_INVOCATION.finditer(content) + if (name := _pdf_invocation_name(invocation)) is not None + ), + ) + + def _pdf_page_content_metrics(page: _PdfPage, page_content: bytes) -> _PdfStreamMetrics: form_metrics_memo: Final[dict[tuple[int, int], _PdfStreamMetrics]] = {} # mutable-ok: page-local walk cache + stream_scans: Final[dict[int, _PdfStreamScan]] = {} # mutable-ok: scan each stream identity once per page def stream_metrics( - content: bytes, + scan: _PdfStreamScan, resources: _PdfResourceDictionary, ancestors: frozenset[int], ) -> _PdfStreamMetrics: xobjects: Final = _pdf_resolve_object(resources.get("/XObject")) xobject_resources: Final = xobjects if isinstance(xobjects, _PdfResourceDictionary) else _PDF_EMPTY_RESOURCES - invoked_names: Final = ( - tuple( - name - for invocation in PDF_XOBJECT_INVOCATION.finditer(content) - if (name := _pdf_invocation_name(invocation)) is not None - ) - if xobject_resources - else () - ) form_metrics: Final = tuple( form_metrics_for_xobject( _pdf_resolve_object(xobject_resources.get(name)), resources, ancestors, ) - for name in invoked_names + for name in scan.invoked_names + if xobject_resources ) return _PdfStreamMetrics( - content=content, - content_bytes=len(content) + sum(metrics.content_bytes for metrics in form_metrics), - operand_bytes=_pdf_string_operand_bytes(content) + sum(metrics.operand_bytes for metrics in form_metrics), + content=scan.content, + content_bytes=len(scan.content) + sum(metrics.content_bytes for metrics in form_metrics), + operand_bytes=scan.operand_bytes + sum(metrics.operand_bytes for metrics in form_metrics), resources=resources, forms=form_metrics, extraction_limited=any(metrics.extraction_limited for metrics in form_metrics), @@ -1730,14 +1742,15 @@ def _pdf_page_content_metrics(page: _PdfPage, page_content: bytes) -> _PdfStream cached: Final = form_metrics_memo.get(cache_key) if cached is not None: return cached - content: Final = _pdf_strip_inline_images(resolved.get_data()) - metrics: Final = stream_metrics(content, resources, ancestors | {identity}) + if identity not in stream_scans: + stream_scans[identity] = _pdf_scan_stream(_pdf_strip_inline_images(resolved.get_data())) + metrics: Final = stream_metrics(stream_scans[identity], resources, ancestors | {identity}) form_metrics_memo[cache_key] = metrics return metrics resources: Final = _pdf_resolve_object(page.get("/Resources")) page_resources: Final = resources if isinstance(resources, _PdfResourceDictionary) else _PDF_EMPTY_RESOURCES - return stream_metrics(page_content, page_resources, frozenset()) + return stream_metrics(_pdf_scan_stream(page_content), page_resources, frozenset()) def _evenly_spaced_page_indices(page_count: int) -> tuple[int, ...]: diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 9192cbe2600..2b4509a5c7e 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -46,6 +46,7 @@ from litellm.litellm_core_utils.token_counter import ( _get_extrapolating_count_function, _get_tiktoken_count_function, _price_pdf_pages, + _pdf_string_operand_bytes, calculate_img_tokens, estimate_image_tokens, high_detail_image_token_upper_bound, @@ -3616,6 +3617,46 @@ def test_shared_pdf_form_uses_each_parents_resources( assert distinct_price - shared_price == second_count - 1 +@pytest.mark.parametrize("context_count", [0, 2, 8]) +def test_shared_pdf_stream_is_scanned_once_across_inherited_contexts(context_count: int): + attributes: Final = b"/Type /XObject /Subtype /Form /BBox [0 0 612 792]" + shared_number: Final = 6 + context_count + shared_content: Final = b"/Leaf Do" + parent_content: Final = b"/Shared Do" + parents: Final = tuple( + _pdf_stream( + parent_content, + attributes + + f" /Resources << /XObject << /Shared {shared_number} 0 R /Leaf {shared_number + index + 1} 0 R >> >>".encode(), + ) + for index in range(context_count) + ) + leaves: Final = tuple( + _pdf_stream( + b"BT /F1 12 Tf (" + b"x" * (index + 1) + b") Tj ET", + attributes + b" /Resources << /Font << /F1 3 0 R >> >>", + ) + for index in range(context_count) + ) + path: Final = b"10 10 m 20 20 l S\n" + padding: Final = path * (PDF_PAGE_MAX_EXACT_CONTENT_BYTES // len(path) + 1) + data_url: Final = _pdf_data_url_with_xobjects( + padding + b"".join(f"/Parent{index} Do\n".encode() for index in range(context_count)), + b"<< " + b" ".join(f"/Parent{index} {6 + index} 0 R".encode() for index in range(context_count)) + b" >>", + (*parents, _pdf_stream(shared_content, attributes), *leaves), + ) + blank_price: Final = _count_inline_pdf_tokens(_pdf_data_url_with_xobjects(b"", b"<< >>", ()), len) + with patch( + "litellm.litellm_core_utils.token_counter._pdf_string_operand_bytes", wraps=_pdf_string_operand_bytes + ) as scan: + price: Final = _count_inline_pdf_tokens(data_url, len) + assert price is not None + assert blank_price is not None + assert price - blank_price == sum(range(1, context_count + 1)) + assert sum(call.args[0] == shared_content for call in scan.call_args_list) == int(context_count > 0) + assert sum(call.args[0] == parent_content for call in scan.call_args_list) == context_count + + def test_repeated_identity_h_form_draws_count_each_text_sample(): path_operator: Final = b"10 10 m 20 20 l S\n" path_padding: Final = path_operator * (PDF_PAGE_MAX_EXACT_CONTENT_BYTES // len(path_operator) + 1)