package main

import (
	"fmt"
	"github.com/chrisdinn/tokens"
)

func main() {
	// Create a new token counter
	tc := tokens.NewCounter()

	// Count tokens in a string.
	someString = "You string for/from model goes here."
	fmt.Printf("%q: %d\n", someString, tc.CountTokens(someString))


	// Count tokens in a ChatCompletionRequest.
	req := openai.ChatCompletionRequest{
		Messages: []openai.ChatCompletionMessage{{
			Role:    openai.ChatMessageRoleSystem,
			Content: "This is a system message.",
		}, {
			Role:    openai.ChatMessageRoleUser,
			Content: "This is a user message.",
		}},
	},
	fmt.Printf("Req tokens: %d\n", req, tc.CountRequestTokens(req))

	// Count tokens in a ChatCompletionMessage.
	msg := openai.ChatCompletionMessage{
		Role: openai.ChatMessageRoleAssistant,
		ToolCalls: []openai.ToolCall{{
			ID:   "testcall_20240330",
			Type: openai.ToolTypeFunction,
			Function: openai.FunctionCall{
				Name:      "get_current_weather",
				Arguments: "{\"location\": \"Park City, UT\"}",
			},
		}},
	}
	fmt.Printf("Msg tokens: %d\n", msg, tc.CountMessageTokens(msg))

	// Count tool tokens on their own.
	tools := []openai.Tool{{
		Type: openai.ToolTypeFunction,
		Function: &openai.FunctionDefinition{
			Name:        "get_current_weather",
			Description: "Get the current weather in a given location.",
			Parameters: jsonschema.Definition{
				Type: jsonschema.Object,
				Properties: map[string]jsonschema.Definition{
					"location": {
						Type:        jsonschema.String,
						Description: "The city and state, e.g. San Francisco, CA",
					},
					"unit": {
						Type: jsonschema.String,
						Enum: []string{"celcius", "fahrenheit"},
					},
					"date": {
						Type:        jsonschema.String,
						Description: "The date for which to get the weather.",
					},
				},
				Required: []string{"location"},
			},
		},
	}},
	fmt.Printf("Tools tokens: %d\n", tools, tc.CountToolTokens(tools))

}