Merge 6209e763d2 into a6245c1bbb

litellm-rust/crates/token-counter/src/types.rs

@@ -49,8 +49,9 @@ fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<Tex
 }
 
 /// Free-form JSON the host counts as text: strings and integers via `str()`,
-/// objects via `json.dumps()`, lists flattened. Objects keep document order so
-/// the dumped text matches Python byte for byte.
+/// objects via `json.dumps()`, lists flattened, base64 `data:` URL payloads
+/// dropped. Objects keep document order so the dumped text matches Python byte
+/// for byte.
 #[derive(Clone, Debug, PartialEq, Serialize)]
 #[serde(untagged)]
 pub(crate) enum TextValue {
@@ -107,10 +108,11 @@ impl<'de> Visitor<'de> for TextValueVisitor {
     }
 
     fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
-        Ok(TextValue::Text(value.to_owned()))
+        Ok(TextValue::Text(without_inline_data(value).to_owned()))
     }
 
-    fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
+    fn visit_string<E>(self, mut value: String) -> Result<Self::Value, E> {
+        value.truncate(without_inline_data(&value).len());
         Ok(TextValue::Text(value))
     }
 
@@ -137,6 +139,21 @@ impl<'de> Visitor<'de> for TextValueVisitor {
     }
 }
 
+const DATA_URL_PREFIX: &str = "data:";
+const BASE64_MARKER: &str = ";base64,";
+
+/// `value` cut after `;base64,` when it is a base64 `data:` URL. Inline images
+/// and files are priced per image or file, not per character, so their payload
+/// is not counted as text; Python's `_without_inline_data` cuts the same strings.
+fn without_inline_data(value: &str) -> &str {
+    if !value.starts_with(DATA_URL_PREFIX) {
+        return value;
+    }
+    value
+        .find(BASE64_MARKER)
+        .map_or(value, |marker| &value[..marker + BASE64_MARKER.len()])
+}
+
 /// Python counts every string-valued key of a message, so any key beyond these
 /// makes the shape unsupported rather than silently uncounted.
 #[derive(Clone, Debug, Deserialize, PartialEq)]

litellm-rust/crates/token-counter/tests/token_counter.rs

@@ -96,6 +96,26 @@ mod json {
         assert_eq!(count.input_tokens, expected);
     }
 
+    fn assert_inline_base64_payload_is_not_counted_as_text(load: JsonLoader, prefix: &str) {
+        let counter = counter(load);
+        let count = |image_url: &str| {
+            let body = serde_json::json!({
+                "model": "m",
+                "input": [{"role": "user", "content": [
+                    {"type": "input_text", "text": "What is in this screenshot?"},
+                    {"type": "input_image", "image_url": image_url}]}, image_url]
+            })
+            .to_string();
+            counter
+                .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
+                .expect("counts")
+                .input_tokens
+        };
+        let with_payload = count(&format!("{prefix}{}", "iVBORw0KGgo".repeat(100_000)));
+        assert_eq!(with_payload, count(prefix));
+        assert!(with_payload < 100, "{with_payload}");
+    }
+
     fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) {
         let request = CountableRequest::parse(body).expect("shape parses");
         assert!(matches!(
@@ -173,6 +193,13 @@ mod json {
                 super::assert_key_presence_follows_python($loader, body, expected);
             }
 
+            #[rstest]
+            #[case::png("data:image/png;base64,")]
+            #[case::pdf("data:application/pdf;base64,")]
+            fn inline_base64_payload_is_not_counted_as_text(#[case] prefix: &str) {
+                super::assert_inline_base64_payload_is_not_counted_as_text($loader, prefix);
+            }
+
             #[rstest]
             #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])]
             #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)]

litellm/proxy/spend_tracking/input_tokens.py

@@ -167,7 +167,48 @@ def _count_text_tokens(model: str, text: object) -> int:
             stack.extend(item)
             continue
         if isinstance(item, dict):
-            token_count += litellm.token_counter(model=model, text=json.dumps(item))
+            token_count += litellm.token_counter(model=model, text=json.dumps(_without_inline_data(item)))
             continue
-        token_count += litellm.token_counter(model=model, text=str(item))
+        token_count += litellm.token_counter(model=model, text=str(_without_inline_data(item)))
     return token_count
+
+
+_DATA_URL_PREFIX: Final = "data:"
+_BASE64_MARKER: Final = ";base64,"
+
+
+def _without_inline_data(value: object) -> object:
+    """`value` with the payload of every base64 `data:` URL cut after `;base64,`.
+
+    Images and files a Responses `input` carries inline are priced by the
+    provider per image or file, not per character, so tokenizing their base64
+    as text inflates the count by orders of magnitude: one screenshot alone is
+    hundreds of thousands of "tokens". The Rust counter cuts the same strings.
+    Walks with an explicit stack: a request body's nesting is client-controlled."""
+    if not isinstance(value, (dict, list)):
+        return _without_inline_payload(value)
+    root: Final[dict | list] = {} if isinstance(value, dict) else []
+    stack: Final[list[tuple[dict | list, dict | list]]] = [(value, root)]
+    while stack:
+        source, target = stack.pop()
+        for key, item in source.items() if isinstance(source, dict) else enumerate(source):
+            copied: object
+            if isinstance(item, (dict, list)):
+                child: dict | list = {} if isinstance(item, dict) else []
+                stack.append((item, child))
+                copied = child
+            else:
+                copied = _without_inline_payload(item)
+            if isinstance(target, dict):
+                target[key] = copied
+            else:
+                target.append(copied)
+    return root
+
+
+def _without_inline_payload(value: object) -> object:
+    if isinstance(value, str) and value.startswith(_DATA_URL_PREFIX):
+        marker: Final = value.find(_BASE64_MARKER)
+        if marker != -1:
+            return value[: marker + len(_BASE64_MARKER)]
+    return value

tests/unit/proxy/spend_tracking/test_input_tokens.py

@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+import json
 from types import MappingProxyType
 from typing import Final
 
 import pytest
 
+import litellm
 from litellm.proxy.spend_tracking.input_tokens import (
     TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS,
     count_input_tokens,
@@ -27,3 +29,55 @@ async def test_large_input_is_still_counted() -> None:
 
     assert counts[CL100K_MODEL] == count_input_tokens_for_model(request_body=request_body, model=CL100K_MODEL)
     assert isinstance(counts, MappingProxyType)
+
+
+def _responses_body(image_url: str) -> dict:
+    return {
+        "model": CL100K_MODEL,
+        "input": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "input_text", "text": "What is in this screenshot?"},
+                    {"type": "input_image", "image_url": image_url},
+                ],
+            }
+        ],
+    }
+
+
+@pytest.mark.parametrize(
+    "image_url",
+    (
+        "data:image/png;base64," + "iVBORw0KGgo" * 100_000,
+        "data:image/jpeg;base64," + "/9j/4AAQSkZJRg" * 100_000,
+    ),
+    ids=("png", "jpeg"),
+)
+def test_inline_base64_payload_is_not_counted_as_text(image_url: str) -> None:
+    """A 1 MB inline image adds no more to the reservation than its `data:` prefix."""
+    prefix: Final = image_url[: image_url.index(",") + 1]
+
+    counted: Final = count_input_tokens_for_model(request_body=_responses_body(image_url), model=CL100K_MODEL)
+
+    assert counted == count_input_tokens_for_model(request_body=_responses_body(prefix), model=CL100K_MODEL)
+    assert counted is not None and counted < 100
+
+
+@pytest.mark.parametrize(
+    "text",
+    (
+        "https://example.com/cat.png",
+        "data:text/plain,not base64 so it is text",
+        "see data:image/png;base64,AAAA in the middle",
+    ),
+    ids=("remote_url", "data_url_without_base64", "data_url_not_at_start"),
+)
+def test_strings_that_are_not_inline_base64_are_counted_in_full(text: str) -> None:
+    item: Final = {"text": text}
+
+    counted: Final = count_input_tokens_for_model(request_body={"input": [item, text]}, model=CL100K_MODEL)
+
+    assert counted == litellm.token_counter(model=CL100K_MODEL, text=json.dumps(item)) + litellm.token_counter(
+        model=CL100K_MODEL, text=text
+    )

tests/unit/rust_bridge/test_token_counter.py

@@ -291,6 +291,19 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
             {"role": "assistant", "content": "Sure."},
         ],
     },
+    {
+        "model": MODEL,
+        "input": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "input_text", "text": "What is in this screenshot?"},
+                    {"type": "input_image", "image_url": "data:image/png;base64," + "iVBORw0KGgo" * 1_000},
+                ],
+            },
+            "data:application/pdf;base64,JVBERi0xLjQK",
+        ],
+    },
     {"model": MODEL, "input": "a single embedding string"},
     {"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"},
     {"model": MODEL, "query": "best harbour", "documents": ["doc one", {"text": "doc two", "title": "T", "n": 3}]},