fix(model-catalog): declare the tool and Responses token count framing keys (#49683)

* fix(model-catalog): declare the tool and Responses token count framing keys

* fix(model-catalog): read token count framings as strict integers

The cost map guard read the framing keys as lax counts, so a text "214"
or a true passed while litellm's reader drops both and falls back to the
plain framing. Read all four framing keys as plain u64, matching the reader

* fix(model-catalog): accept whole-float token count framings the way the Python reader does

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm-rust/crates/model-catalog/src/wire/attributes.rs

@@ -11,7 +11,7 @@ use crate::error::{DateCause, Error, Repr};
 use crate::vocabulary::{
     AudioFormat, EffortLevel, Endpoint, InputModality, Mode, OutputModality, VertexAiAudioApi,
 };
-use crate::wire::lax::Count;
+use crate::wire::lax::{Count, NumericCount};
 
 macro_rules! attributes {
     ($($(#[$doc:meta])* $name:ident: $ty:ty $(as $via:ty)?),* $(,)?) => {
@@ -104,9 +104,12 @@ attributes! {
     /// Cloud regions the model is available in ('global' or region ids).
     supported_regions: Vec<String>,
     thinking_always_on: bool,
-    token_count_framing: u64 as Count,
+    token_count_framing: u64 as NumericCount,
     token_count_model: TokenCountModel,
+    token_count_responses_framing: u64 as NumericCount,
+    token_count_responses_tool_framing: u64 as NumericCount,
     token_count_tokenizer: TokenCountTokenizer,
+    token_count_tool_framing: u64 as NumericCount,
     /// Provider default tokens-per-minute limit.
     tpm: u64 as Count,
     use_openai_responses_path: bool,
@@ -302,6 +305,18 @@ mod tests {
     #[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)]
     #[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)]
     #[case::reasoning_mode_not_a_string("reasoning_modes", json!([1]), false)]
+    #[case::framing("token_count_framing", json!(12), true)]
+    #[case::framing_written_as_text("token_count_framing", json!("12"), false)]
+    #[case::framing_written_as_whole_float("token_count_framing", json!(12.0), true)]
+    #[case::fractional_tool_framing("token_count_tool_framing", json!(1.5), false)]
+    #[case::tool_framing("token_count_tool_framing", json!(214), true)]
+    #[case::tool_framing_written_as_text("token_count_tool_framing", json!("214"), false)]
+    #[case::tool_framing_as_a_flag("token_count_tool_framing", json!(true), false)]
+    #[case::negative_tool_framing("token_count_tool_framing", json!(-1), false)]
+    #[case::responses_framing("token_count_responses_framing", json!(38), true)]
+    #[case::responses_framing_as_a_flag("token_count_responses_framing", json!(false), false)]
+    #[case::responses_tool_framing("token_count_responses_tool_framing", json!(267), true)]
+    #[case::responses_tool_framing_written_as_text("token_count_responses_tool_framing", json!("267"), false)]
     fn each_attribute_reads_its_own_type(
         #[case] key: &str,
         #[case] value: Value,

litellm-rust/crates/model-catalog/src/wire/lax.rs

@@ -8,6 +8,8 @@ pub(crate) struct LaxI64;
 
 pub(crate) struct Count;
 
+pub(crate) struct NumericCount;
+
 impl<'de> DeserializeAs<'de, i64> for LaxI64 {
     fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
         deserializer.deserialize_any(Self)
@@ -52,6 +54,34 @@ impl<'de> DeserializeAs<'de, u64> for Count {
     }
 }
 
+impl<'de> DeserializeAs<'de, u64> for NumericCount {
+    fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<u64, D::Error> {
+        deserializer.deserialize_any(Self)
+    }
+}
+
+impl Visitor<'_> for NumericCount {
+    type Value = u64;
+
+    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str("a whole number of zero or more")
+    }
+
+    fn visit_u64<E: de::Error>(self, value: u64) -> Result<u64, E> {
+        Ok(value)
+    }
+
+    fn visit_i64<E: de::Error>(self, value: i64) -> Result<u64, E> {
+        u64::try_from(value).map_err(|_| E::invalid_value(de::Unexpected::Signed(value), &self))
+    }
+
+    fn visit_f64<E: de::Error>(self, value: f64) -> Result<u64, E> {
+        integral_float(value)
+            .and_then(|whole| u64::try_from(whole).ok())
+            .ok_or_else(|| E::invalid_value(de::Unexpected::Float(value), &self))
+    }
+}
+
 fn integer_string(value: &str) -> Option<i64> {
     let integer = match value.split_once('.') {
         Some((integer, fraction)) => {
@@ -99,6 +129,12 @@ mod tests {
             .map(DeserializeAsWrap::into_inner)
     }
 
+    fn numeric_count(value: Value) -> Option<u64> {
+        DeserializeAsWrap::<u64, NumericCount>::deserialize(value)
+            .ok()
+            .map(DeserializeAsWrap::into_inner)
+    }
+
     use serde::Deserialize;
 
     #[rstest]
@@ -125,4 +161,23 @@ mod tests {
     ) {
         assert_eq!(count(written), expected);
     }
+
+    #[rstest]
+    #[case::integer(json!(214), Some(214))]
+    #[case::zero(json!(0), Some(0))]
+    #[case::whole_float(json!(214.0), Some(214))]
+    #[case::negative_zero(json!(-0.0), Some(0))]
+    #[case::above_i64(json!(u64::MAX), Some(u64::MAX))]
+    #[case::negative(json!(-1), None)]
+    #[case::negative_whole_float(json!(-1.0), None)]
+    #[case::fraction(json!(1.5), None)]
+    #[case::string(json!("214"), None)]
+    #[case::boolean(json!(true), None)]
+    #[case::list(json!([214]), None)]
+    fn numeric_counts_accept_only_whole_json_numbers(
+        #[case] written: Value,
+        #[case] expected: Option<u64>,
+    ) {
+        assert_eq!(numeric_count(written), expected);
+    }
 }