feat: add token utils to count the total tokens utilized

e2e/factories/factories.go

@@ -3,6 +3,7 @@ package factories
 import (
 	"context"
 	"encoding/json"
+
 	openaiconnector "github.com/basemind-ai/monorepo/gen/go/openai/v1"
 	"github.com/basemind-ai/monorepo/shared/go/datatypes"
 	"github.com/basemind-ai/monorepo/shared/go/db"

go.mod

@@ -21,6 +21,7 @@ require (
 	github.com/rs/zerolog v1.31.0
 	github.com/sethvargo/go-envconfig v0.9.0
 	github.com/stretchr/testify v1.8.4
+	github.com/tiktoken-go/tokenizer v0.1.0
 	golang.org/x/sync v0.3.0
 	google.golang.org/grpc v1.58.2
 	google.golang.org/protobuf v1.31.0
@@ -42,6 +43,7 @@ require (
 	github.com/containerd/continuity v0.4.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
+	github.com/dlclark/regexp2 v1.9.0 // indirect
 	github.com/docker/cli v24.0.6+incompatible // indirect
 	github.com/docker/docker v24.0.6+incompatible // indirect
 	github.com/docker/go-connections v0.4.0 // indirect

go.sum

@@ -50,6 +50,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
+github.com/dlclark/regexp2 v1.9.0 h1:pTK/l/3qYIKaRXuHnEnIf7Y5NxfRPfpb7dis6/gdlVI=
+github.com/dlclark/regexp2 v1.9.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/docker/cli v24.0.6+incompatible h1:fF+XCQCgJjjQNIMjzaSmiKJSCcfcXb3TWTcc7GAneOY=
 github.com/docker/cli v24.0.6+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=
 github.com/docker/docker v24.0.6+incompatible h1:hceabKCtUgDqPu+qm0NgsaXf28Ljf4/pWFL7xjWWDgE=
@@ -249,6 +251,8 @@ github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o
 github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/tiktoken-go/tokenizer v0.1.0 h1:c1fXriHSR/NmhMDTwUDLGiNhHwTV+ElABGvqhCWLRvY=
+github.com/tiktoken-go/tokenizer v0.1.0/go.mod h1:7SZW3pZUKWLJRilTvWCa86TOVIiiJhYj3FQ5V3alWcg=
 github.com/vmihailenco/go-tinylfu v0.2.2 h1:H1eiG6HM36iniK6+21n9LLpzx1G9R3DJa2UjUjbynsI=
 github.com/vmihailenco/go-tinylfu v0.2.2/go.mod h1:CutYi2Q9puTxfcolkliPq4npPuofg9N9t8JVrjzwa3Q=
 github.com/vmihailenco/msgpack/v5 v5.3.4/go.mod h1:7xyJ9e+0+9SaZT0Wt1RGleJXzli6Q/V5KbhBonMG9jc=

services/api-gateway/connectors/openai/openai.go

@@ -3,9 +3,12 @@ package openai
 import (
 	"context"
 	"errors"
-	"github.com/basemind-ai/monorepo/shared/go/datatypes"
+	"fmt"
 	"io"
 
+	"github.com/basemind-ai/monorepo/shared/go/datatypes"
+	"github.com/basemind-ai/monorepo/shared/go/tokenutils"
+
 	openaiconnector "github.com/basemind-ai/monorepo/gen/go/openai/v1"
 	"github.com/rs/zerolog/log"
 	"google.golang.org/grpc"
@@ -48,7 +51,20 @@ func (c *Client) RequestPrompt(
 	if requestErr != nil {
 		return "", requestErr
 	}
-	// TODO handle token related logic here by using the response token properties.
+
+	// Count the total number of tokens utilized for openai prompt
+	reqPromptString := GetRequestPromptString(promptRequest.Messages)
+	promptReqTokenCount, tokenizationErr := tokenutils.GetPromptTokenCount(reqPromptString, applicationPromptConfig.PromptConfigData.ModelType)
+	if tokenizationErr != nil {
+		log.Err(tokenizationErr).Msg("failed to get prompt token count")
+	}
+
+	promptResTokenCount, tokenizationErr := tokenutils.GetPromptTokenCount(response.Content, applicationPromptConfig.PromptConfigData.ModelType)
+	if tokenizationErr != nil {
+		log.Err(tokenizationErr).Msg("failed to get prompt token count")
+	}
+
+	log.Debug().Msg(fmt.Sprintf("Total tokens utilized: Request-%d, Response-%d", promptReqTokenCount, promptResTokenCount))
 	return response.Content, nil
 }
 
@@ -78,17 +94,31 @@ func (c *Client) RequestStream(
 		return
 	}
 
+	reqPromptString := GetRequestPromptString(promptRequest.Messages)
+	promptReqTokenCount, tokenizationErr := tokenutils.GetPromptTokenCount(reqPromptString, applicationPromptConfig.PromptConfigData.ModelType)
+	if tokenizationErr != nil {
+		log.Err(tokenizationErr).Msg("failed to get prompt token count")
+	}
+	log.Debug().Msg(fmt.Sprintf("Total tokens utilized for request prompt - %d", promptReqTokenCount))
+
+	var promptResTokenCount int
+
 	for {
 		msg, receiveErr := stream.Recv()
 		if receiveErr != nil {
 			if !errors.Is(receiveErr, io.EOF) {
 				errChannel <- receiveErr
 			}
 			close(contentChannel)
+			log.Debug().Msg(fmt.Sprintf("Tokens utilized for streaming response-%d", promptResTokenCount))
 			return
 		}
 
-		// TODO handle token related logic here
+		streamResTokenCount, tokenizationErr := tokenutils.GetPromptTokenCount(msg.Content, applicationPromptConfig.PromptConfigData.ModelType)
+		if tokenizationErr != nil {
+			log.Err(tokenizationErr).Msg("failed to get prompt token count")
+		}
+		promptResTokenCount += streamResTokenCount
 		contentChannel <- msg.Content
 	}
 }

services/api-gateway/connectors/openai/utils.go

@@ -3,9 +3,10 @@ package openai
 import (
 	"encoding/json"
 	"fmt"
+	"strings"
+
 	"github.com/basemind-ai/monorepo/shared/go/datatypes"
 	"github.com/basemind-ai/monorepo/shared/go/db"
-	"strings"
 
 	openaiconnector "github.com/basemind-ai/monorepo/gen/go/openai/v1"
 )
@@ -86,3 +87,12 @@ func CreatePromptRequest(
 
 	return promptRequest, nil
 }
+
+func GetRequestPromptString(messages []*openaiconnector.OpenAIMessage) string {
+	var promptMessages string
+	for _, message := range messages {
+		promptMessages += *message.Content
+		promptMessages += "\n"
+	}
+	return strings.TrimRight(promptMessages, "\n")
+}

services/api-gateway/connectors/openai/utils_test.go

@@ -2,11 +2,12 @@ package openai_test
 
 import (
 	"fmt"
+	"testing"
+
 	"github.com/basemind-ai/monorepo/e2e/factories"
 	"github.com/basemind-ai/monorepo/services/api-gateway/connectors/openai"
 	"github.com/basemind-ai/monorepo/shared/go/db"
 	"github.com/stretchr/testify/assert"
-	"testing"
 
 	openaiconnector "github.com/basemind-ai/monorepo/gen/go/openai/v1"
 )
@@ -169,4 +170,42 @@ func TestUtils(t *testing.T) {
 			assert.Error(t, err)
 		})
 	})
+
+	t.Run("GetRequestPromptString", func(t *testing.T) {
+		t.Run("returns the request prompt as string", func(t *testing.T) {
+			floatValue := float32(1)
+			uintValue := uint32(1)
+
+			expectedModelParameters := &openaiconnector.OpenAIModelParameters{
+				Temperature:      &floatValue,
+				TopP:             &floatValue,
+				MaxTokens:        &uintValue,
+				PresencePenalty:  &floatValue,
+				FrequencyPenalty: &floatValue,
+			}
+
+			systemMessage := "You are a helpful chat bot."
+			applicationId := "12345"
+			userInput := "Please write an essay on Dogs."
+			content := fmt.Sprintf("This is what the user asked for: %s", userInput)
+
+			promptRequest := &openaiconnector.OpenAIPromptRequest{
+				Model:         openaiconnector.OpenAIModel_OPEN_AI_MODEL_GPT3_5_TURBO_4K,
+				ApplicationId: &applicationId,
+				Parameters:    expectedModelParameters,
+				Messages: []*openaiconnector.OpenAIMessage{
+					{
+						Content: &systemMessage,
+						Role:    openaiconnector.OpenAIMessageRole_OPEN_AI_MESSAGE_ROLE_SYSTEM,
+					},
+					{
+						Content: &content,
+						Role:    openaiconnector.OpenAIMessageRole_OPEN_AI_MESSAGE_ROLE_USER,
+					},
+				},
+			}
+			reqPromptString := openai.GetRequestPromptString(promptRequest.Messages)
+			assert.Equal(t, "You are a helpful chat bot.\nThis is what the user asked for: Please write an essay on Dogs.", reqPromptString)
+		})
+	})
 }

shared/go/tokenutils/tokenutils.go

@@ -0,0 +1,24 @@
+package tokenutils
+
+import (
+	"github.com/basemind-ai/monorepo/shared/go/db"
+	"github.com/tiktoken-go/tokenizer"
+)
+
+var modelEncodingMap map[db.ModelType]tokenizer.Encoding = map[db.ModelType]tokenizer.Encoding{
+	db.ModelTypeGpt35Turbo:    tokenizer.Cl100kBase,
+	db.ModelTypeGpt35Turbo16k: tokenizer.Cl100kBase,
+	db.ModelTypeGpt4:          tokenizer.Cl100kBase,
+	db.ModelTypeGpt432k:       tokenizer.Cl100kBase,
+}
+
+func GetPromptTokenCount(prompt string, modelType db.ModelType) (int, error) {
+	encoding := modelEncodingMap[modelType]
+	enc, err := tokenizer.Get(encoding)
+	if err != nil {
+		return -1, err
+	}
+
+	ids, _, _ := enc.Encode(prompt)
+	return len(ids), nil
+}

shared/go/tokenutils/tokenutils_test.go

@@ -0,0 +1,42 @@
+package tokenutils_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/basemind-ai/monorepo/shared/go/db"
+	"github.com/basemind-ai/monorepo/shared/go/tokenutils"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetPromptTokenCount(t *testing.T) {
+	testCases := []struct {
+		input    string
+		expected int
+	}{
+		{
+			input:    "Hello world!",
+			expected: 3,
+		},
+		{
+			input:    "",
+			expected: 0,
+		},
+		{
+			input:    "Goodbye world!",
+			expected: 4,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(fmt.Sprintf("Test: %d", testCase.expected), func(t *testing.T) {
+			actual, _ := tokenutils.GetPromptTokenCount(testCase.input, db.ModelTypeGpt35Turbo)
+			assert.Equal(t, testCase.expected, actual)
+		})
+	}
+
+	// test for the invalid encodings
+	tokenCnt, err := tokenutils.GetPromptTokenCount("Hello world!", db.ModelType("invalid"))
+	assert.NotNil(t, err)
+	assert.Equal(t, -1, tokenCnt)
+}