Merge bc7cce5920 into e768ad55ce

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

@@ -124,7 +124,31 @@ impl TokenCounter {
                 .sum::<Result<usize, _>>()?,
             None => 0,
         };
-        Ok(TOKENS_PER_MESSAGE + role_tokens + name_tokens + content_tokens)
+        // Python counts only each tool call's `arguments` string
+        // (`_count_function_call_tokens`); names ride with the tool
+        // definitions and `tool_choice`.
+        let tool_call_tokens = match &message.tool_calls {
+            Some(calls) => calls
+                .iter()
+                .map(|call| self.count_text(call.function.arguments.as_deref().unwrap_or("")))
+                .sum::<Result<usize, _>>()?,
+            None => 0,
+        };
+        let tool_call_id_tokens = match &message.tool_call_id {
+            Some(id) => self.count_text(id)?,
+            None => 0,
+        };
+        let legacy_call_tokens = match &message.function_call {
+            Some(call) => self.count_text(call.arguments.as_deref().unwrap_or(""))?,
+            None => 0,
+        };
+        Ok(TOKENS_PER_MESSAGE
+            + role_tokens
+            + name_tokens
+            + content_tokens
+            + tool_call_tokens
+            + tool_call_id_tokens
+            + legacy_call_tokens)
     }
 
     fn count_content_item(&self, item: &ContentItem) -> Result<usize, Error> {

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

@@ -137,6 +137,47 @@ impl<'de> Visitor<'de> for TextValueVisitor {
     }
 }
 
+/// The parts of an assistant tool call Python counts: only the `arguments`
+/// string contributes (`_count_function_call_tokens` in
+/// `litellm_core_utils/token_counter.py`). An absent arguments key counts as the
+/// empty string, while an explicit null counts as the string "None" the way
+/// Python's `str(None)` does; any other non-string `arguments` declines so
+/// Python handles the fallback instead of miscounting.
+///
+/// `deny_unknown_fields` is deliberately absent from this shape, and from
+/// `ToolCallFunction` and `LegacyFunctionCall`: Python reads only
+/// `function.arguments` off a tool call, so the extra `id`, `type` and `name`
+/// keys the OpenAI shape carries are ignored rather than declined.
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub(crate) struct ToolCall {
+    pub(crate) function: ToolCallFunction,
+}
+
+/// Python runs `str(function.get("arguments", ""))`: a missing key counts as
+/// empty, but an explicit null counts as the string "None". `serde(default)`
+/// keeps a missing key at `None`, and the deserializer turns an explicit null
+/// into "None" so the counter sees what Python counts.
+fn explicit_null_arguments_as_none<'de, D: Deserializer<'de>>(
+    deserializer: D,
+) -> Result<Option<String>, D::Error> {
+    match Option::<String>::deserialize(deserializer)? {
+        Some(arguments) => Ok(Some(arguments)),
+        None => Ok(Some(String::from("None"))),
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub(crate) struct ToolCallFunction {
+    #[serde(default, deserialize_with = "explicit_null_arguments_as_none")]
+    pub(crate) arguments: Option<String>,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub(crate) struct LegacyFunctionCall {
+    #[serde(default, deserialize_with = "explicit_null_arguments_as_none")]
+    pub(crate) arguments: Option<String>,
+}
+
 /// 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)]
@@ -145,6 +186,9 @@ pub(crate) struct Message {
     pub(crate) role: Option<String>,
     pub(crate) name: Option<String>,
     pub(crate) content: Option<MessageContent>,
+    pub(crate) tool_call_id: Option<String>,
+    pub(crate) tool_calls: Option<Vec<ToolCall>>,
+    pub(crate) function_call: Option<LegacyFunctionCall>,
 }
 
 #[derive(Clone, Debug, Deserialize, PartialEq)]

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

@@ -70,6 +70,27 @@ mod json {
     const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour",
   "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#;
 
+    const TOOL_CALLS_AND_TOOL_RESULT: &str = r#"{"model":"claude-sonnet-4-5","messages":[
+  {"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{\"city\": \"Tokyo\"}"}}]},
+  {"role":"tool","tool_call_id":"1","content":"Sunny, 22C"}]}"#;
+
+    const LEGACY_FUNCTION_CALL: &str = r#"{"model":"claude-sonnet-4-5","messages":[
+  {"role":"assistant","function_call":{"name":"f","arguments":"{\"city\": \"Tokyo\"}"}}]}"#;
+
+    const ASSISTANT_CONTENT_WITH_TOOL_CALL: &str = r#"{"model":"claude-sonnet-4-5","messages":[
+  {"role":"assistant","content":"calling","tool_calls":[{"id":"1","type":"function","function":{"name":"get_weather","arguments":"{\"city\": \"Paris, France\"}"}}]}]}"#;
+
+    const NULL_TOOL_CALL_ARGUMENTS: &str = r#"{"model":"claude-sonnet-4-5","messages":[
+  {"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":null}}]},
+  {"role":"tool","tool_call_id":"1","content":"x"}]}"#;
+
+    const MISSING_TOOL_CALL_ARGUMENTS: &str = r#"{"model":"claude-sonnet-4-5","messages":[
+  {"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f"}}]},
+  {"role":"tool","tool_call_id":"1","content":"x"}]}"#;
+
+    const LEGACY_NULL_FUNCTION_CALL_ARGUMENTS: &str = r#"{"model":"claude-sonnet-4-5","messages":[
+  {"role":"assistant","function_call":{"name":"f","arguments":null}}]}"#;
+
     fn assert_count_request_matches_python_token_counter(
         load: JsonLoader,
         body: &str,
@@ -154,6 +175,12 @@ mod json {
             #[case::responses_input_items(super::RESPONSES_INPUT, 62)]
             #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)]
             #[case::rerank_query_and_documents(super::RERANK, 41)]
+            #[case::tool_calls_and_tool_result(super::TOOL_CALLS_AND_TOOL_RESULT, 24)]
+            #[case::legacy_function_call(super::LEGACY_FUNCTION_CALL, 14)]
+            #[case::assistant_content_with_tool_call(super::ASSISTANT_CONTENT_WITH_TOOL_CALL, 16)]
+            #[case::null_tool_call_arguments_count_as_none(super::NULL_TOOL_CALL_ARGUMENTS, 14)]
+            #[case::missing_tool_call_arguments_count_as_empty(super::MISSING_TOOL_CALL_ARGUMENTS, 13)]
+            #[case::legacy_null_function_call_arguments_count_as_none(super::LEGACY_NULL_FUNCTION_CALL_ARGUMENTS, 8)]
             fn count_request_matches_python_token_counter(
                 #[case] body: &str,
                 #[case] expected: usize,
@@ -235,8 +262,17 @@ mod json {
 #[rstest]
 #[case::not_json(b"not json" as &[u8])]
 #[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)]
-#[case::message_with_tool_calls(
-    br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#
+#[case::tool_calls_not_a_list(
+    br#"{"model":"m","messages":[{"role":"assistant","tool_calls":"nope"}]}"#
+)]
+#[case::tool_call_without_function(
+    br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function"}]}]}"#
+)]
+#[case::tool_call_arguments_is_a_number(
+    br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":42}}]}]}"#
+)]
+#[case::tool_call_arguments_is_an_object(
+    br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":{"city":"Tokyo"}}}]}]}"#
 )]
 #[case::dict_content(
     br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"#