Add token counting for OpenAI tools/tool_choice #4716
Conversation
|
The latest updates on your projects. Learn more about Vercel for Git ↗︎
|
| @@ -3,15 +3,14 @@ | |||
|
|
|||
| import os | |||
| import sys | |||
| import traceback | |||
There was a problem hiding this comment.
I removed unused imports
| diff = counted_tokens - expected_tokens | ||
| assert ( | ||
| diff >= 0 and diff <= 3 | ||
| ), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by 3 in the over-counting direction." |
There was a problem hiding this comment.
Note that this doesnt do an exact comparison, as I haven't figured out how to get the counts to match exactly, and I don't know if anyone has, since this is all based on reverse-engineering. So I allow over-counting by 3.
| @@ -401,6 +401,18 @@ class ChatCompletionToolParam(TypedDict): | |||
| function: ChatCompletionToolParamFunctionChunk | |||
|
|
|||
|
|
|||
| class Function(TypedDict, total=False): | |||
There was a problem hiding this comment.
It seems that you replicate the types from the openai package? I'm not sure why that's done (versus directly importing from openai), but I copied the types here.
|
By the way, it appears that token_counter repeats the message-token-calculation twice- once before it knows what tokenizer it will use, and then again when it decide to use openAI tokenizer. It'd be more efficient if we avoided the first calculation for openai. |
| num_tokens += 9 # Additional tokens for function definition of tools | ||
| # If there's a system message and tools are present, subtract four tokens | ||
| if tools and includes_system_message: | ||
| num_tokens -= 4 |
There was a problem hiding this comment.
interesting is there any doc / openai reference for this? @pamelafox
There was a problem hiding this comment.
Unfortunately, not as far as I know. I found this logic from OpenAI discussion forum threads (like https://community.openai.com/t/how-to-calculate-the-tokens-when-using-function-call/266573/59) and other packages, such as https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java (which I reference in a docstring in this file).
I then verified it by running my test calls against live OpenAI. For basic functions, it gets it exactly correct, but it over-counts by up to 3 tokens for fancier function definitions.
Title
This PR implements token counting for OpenAI calls that include tools and tool_choice parameter, based on the approach I use in https://github.com/pamelafox/openai-messages-token-helper
I don't know if any other LLMs support tools/tool_choice at this point, and if so, I don't know how to count tokens for other LLMs yet, so this is OpenAI-specific.
I brought in my test cases that have been verified with OpenAI, and also fixed an issue that I saw with counting tokens for images.
I can also expand on the current documentation to mention tools/tool_choice, if this approach looks good.
Relevant issues
Fixes #4664
Type
🆕 New Feature
🐛 Bug Fix
✅ Test
Changes
[REQUIRED] Testing