tokens

package
v1.3.3-0...-f31c912 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Jan 9, 2024 License: AGPL-3.0 Imports: 4 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func CountTokens

func CountTokens(tc *TokenCount) (int, int)

CountTokens is a helper that calls tc.CountAll() on a TokenCount pointer, but also return 0, 0 when receiving a nil pointer. This makes token counting less awkward in cases where we don't know whether a completion happened or not.

Types

type AsynchronousTokenCounter

type AsynchronousTokenCounter struct {
	// contains filtered or unexported fields
}

AsynchronousTokenCounter counts completion tokens that are used by a streamed completion request. When creating a AsynchronousTokenCounter, the streaming might not be finished, and we can't evaluate how many tokens will be used. In this case, the streaming routine must add streamed completion result with the Add() method and call Finish() once the completion is finished. TokenCount() will hang until either Finish() is called or the context is Done.

func NewAsynchronousTokenCounter

func NewAsynchronousTokenCounter(completionStart string) (*AsynchronousTokenCounter, error)

NewAsynchronousTokenCounter takes the partial completion request output and creates a token counter that can be already returned even if not all the content has been streamed yet. Streamed content can be added a posteriori with Add(). Once all the content is streamed, Finish() must be called.

func (*AsynchronousTokenCounter) Add

func (tc *AsynchronousTokenCounter) Add() error

Add a streamed token to the count.

func (*AsynchronousTokenCounter) TokenCount

func (tc *AsynchronousTokenCounter) TokenCount() int

TokenCount implements the TokenCounter interface. It returns how many tokens have been counted. It also marks the counter as finished. Once a counter is finished, tokens cannot be added anymore.

type StaticTokenCounter

type StaticTokenCounter int

StaticTokenCounter is a token counter whose count has already been evaluated. This can be used to count prompt tokens (we already know the exact count), or to count how many tokens were used by an already finished completion request.

func NewPromptTokenCounter

func NewPromptTokenCounter(prompt []openai.ChatCompletionMessage) (*StaticTokenCounter, error)

NewPromptTokenCounter takes a list of openai.ChatCompletionMessage and computes how many tokens are used by sending those messages to the model.

func NewSynchronousTokenCounter

func NewSynchronousTokenCounter(completion string) (*StaticTokenCounter, error)

NewSynchronousTokenCounter takes the completion request output and computes how many tokens were used by the model to generate this result.

func (*StaticTokenCounter) TokenCount

func (tc *StaticTokenCounter) TokenCount() int

TokenCount implements the TokenCounter interface.

type TokenCount

type TokenCount struct {
	Prompt     TokenCounters
	Completion TokenCounters
}

TokenCount holds TokenCounters for both Prompt and Completion tokens. As the agent performs multiple calls to the model, each call creates its own prompt and completion TokenCounter.

Prompt TokenCounters can be created before doing the call as we know the full prompt and can tokenize it. This is the PromptTokenCounter purpose.

Completion TokenCounters can be created after receiving the model response. Depending on the response type, we might have the full result already or get a stream that will provide the completion result in the future. For the latter, the token count will be evaluated lazily and asynchronously. StaticTokenCounter count tokens synchronously, while AsynchronousTokenCounter supports the streaming use-cases.

func NewTokenCount

func NewTokenCount() *TokenCount

NewTokenCount initializes a new TokenCount struct.

func (*TokenCount) AddCompletionCounter

func (tc *TokenCount) AddCompletionCounter(completion TokenCounter)

AddCompletionCounter adds a TokenCounter to the Completion list.

func (*TokenCount) AddPromptCounter

func (tc *TokenCount) AddPromptCounter(prompt TokenCounter)

AddPromptCounter adds a TokenCounter to the Prompt list.

func (*TokenCount) CountAll

func (tc *TokenCount) CountAll() (int, int)

CountAll iterates over all counters and returns how many prompt and completion tokens were used. As completion token counting can require waiting for a response to be streamed, the caller should pass a context and use it to implement some kind of deadline to avoid hanging infinitely if something goes wrong (e.g. use `context.WithTimeout()`).

type TokenCounter

type TokenCounter interface {
	TokenCount() int
}

TokenCounter is an interface for all token counters, regardless of the kind of token they count (prompt/completion) or the tokenizer used. TokenCount must be idempotent.

type TokenCounters

type TokenCounters []TokenCounter

TokenCounters is a list of TokenCounter and offers function to iterate over all counters and compute the total.

func (TokenCounters) CountAll

func (tc TokenCounters) CountAll() int

CountAll iterates over a list of TokenCounter and returns the sum of the results of all counters. As the counting process might be blocking/take some time, the caller should set a Deadline on the context.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL