Skip to content

[aisdk] add tokens count feature#589

Open
thiago-figueredo wants to merge 1 commit into
jetify-com:mainfrom
thiago-figueredo:thiago-figueredo/add-tokens-count-feature
Open

[aisdk] add tokens count feature#589
thiago-figueredo wants to merge 1 commit into
jetify-com:mainfrom
thiago-figueredo:thiago-figueredo/add-tokens-count-feature

Conversation

@thiago-figueredo

Copy link
Copy Markdown

Summary

Token counting allows estimating the cost of a request and checking whether a prompt fits within the model's context window before making an API call.

  • Add TokenCounter interface that LanguageModel implementations can optionally implement
  • Add CountTokens and CountTokensStr convenience functions for counting tokens in messages
  • Implement token counting for Anthropic provider using their Beta Messages CountTokens API
  • Implement token counting for OpenAI provider using tiktoken (with proper encoding selection for different model families)
  • Add mock support for token counting in tests

How was it tested?

  • Added unit tests in count_tokens_test.go covering basic message counting, multiple messages, error handling, and unsupported models
  • Added unit tests in tokenizer_test.go for OpenAI's tiktoken-based counting including encoding selection for different models
  • Ran tests locally
  • Ran aisdk/ai/examples/count_tokens/main.go

Community Contribution License

All community contributions in this pull request are licensed to the project maintainers under the terms of the Apache 2 License.

By creating this pull request I represent that I have the right to license the contributions to the project maintainers under the Apache 2 License as stated in the Community Contribution License.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant