diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs index 6370d9d74f8..d14bc852de0 100644 --- a/litellm-rust/crates/token-counter/src/counter.rs +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -124,7 +124,31 @@ impl TokenCounter { .sum::>()?, 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::>()?, + 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 { diff --git a/litellm-rust/crates/token-counter/src/types.rs b/litellm-rust/crates/token-counter/src/types.rs index c1236f94d9a..f71ef79b32b 100644 --- a/litellm-rust/crates/token-counter/src/types.rs +++ b/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, D::Error> { + match Option::::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, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct LegacyFunctionCall { + #[serde(default, deserialize_with = "explicit_null_arguments_as_none")] + pub(crate) arguments: Option, +} + /// 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, pub(crate) name: Option, pub(crate) content: Option, + pub(crate) tool_call_id: Option, + pub(crate) tool_calls: Option>, + pub(crate) function_call: Option, } #[derive(Clone, Debug, Deserialize, PartialEq)] diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 542bd4a1fc4..77dc618db38 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/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"}}]}"#