fix(token_counter): align shared stream handling

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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, ...]:

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)