From c6a4437e980e1a498066399f69666a14aa618d05 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 16 Oct 2026 00:30:12 +0000
Subject: [PATCH] ci(cost map): classify token_count_framing and
 token_count_fallback_model in the schema generator (#46784)

* ci(cost map): classify token_count_framing and token_count_fallback_model in the schema generator

* ci(cost map): bound token_count_framing to non-negative and token_count_fallback_model to non-empty

* test(cost map): reject a fractional token_count_framing in the generated schema

* test(cost map): pin that a token_count_framing of 0 passes the generated schema
---
 ci_cd/generate_model_prices_schema.py  |  9 +++++++
 tests/unit/test_model_prices_schema.py | 33 ++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py
index 8eec07dadda..e429aae36b9 100644
--- a/ci_cd/generate_model_prices_schema.py
+++ b/ci_cd/generate_model_prices_schema.py
@@ -189,6 +189,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = {
         **NONNEG_INTEGER,
         "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
     },
+    "token_count_framing": {
+        **NONNEG_INTEGER,
+        "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
+    },
     "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
     "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
 }
@@ -283,6 +287,11 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
             "enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
         },
         "comment": STRING,
+        "token_count_fallback_model": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Model that counts tokens for this one when the provider cannot count it.",
+        },
         "audio_transcription_config": STRING,
         "vertex_ai_audio_api": {
             "type": "string",
diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py
index 314a401e88e..b12d3b34c8b 100644
--- a/tests/unit/test_model_prices_schema.py
+++ b/tests/unit/test_model_prices_schema.py
@@ -299,6 +299,39 @@ def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict):
     assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}})
 
 
+TOKEN_COUNT_ENTRY: Final = MappingProxyType(
+    {
+        "litellm_provider": "bedrock_converse",
+        "mode": "chat",
+        "token_count_framing": 78,
+        "token_count_fallback_model": "us.anthropic.claude-opus-5-5",
+    }
+)
+
+
+@pytest.mark.parametrize("framing", [0, 78])
+def test_generator_classifies_token_count_settings(framing: int):
+    generator = load_generator()
+    schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)})))
+    assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, "token_count_framing": framing}})
+
+
+@pytest.mark.parametrize(
+    "setting",
+    [
+        {"token_count_framing": -1},
+        {"token_count_framing": "78"},
+        {"token_count_framing": 1.5},
+        {"token_count_fallback_model": ""},
+        {"token_count_fallback_model": 5},
+    ],
+)
+def test_generated_token_count_schema_rejects_malformed_settings(setting: dict):
+    generator = load_generator()
+    schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)})))
+    assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}})
+
+
 def find_duplicate_keys(path: Path) -> list[str]:
     duplicates: list[str] = []