OpenAI - Added Support for Token Callback / Bug Fixes (newrelic#875)

* Bug Fixes for OpenAI
nr-swilloughby · Mar 25, 2024 · 67c0a92 · 67c0a92
1 parent 3480bb1
commit 67c0a92
Show file tree

Hide file tree

Showing 6 changed files with 220 additions and 111 deletions.
diff --git a/v3/integrations/nropenai/examples/chatcompletion/chatcompletion_example.go b/v3/integrations/nropenai/examples/chatcompletion/chatcompletion_example.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/newrelic/go-agent/v3/integrations/nropenai"
 	"github.com/newrelic/go-agent/v3/newrelic"
+	"github.com/pkoukk/tiktoken-go"
 	openai "github.com/sashabaranov/go-openai"
 )
 
@@ -25,6 +26,34 @@ func main() {
 	}
 	app.WaitForConnection(10 * time.Second)
 
+	// SetLLMTokenCountCallback allows for custom token counting, if left unset and if newrelic.ConfigAIMonitoringRecordContentEnabled()
+	// is disabled, no token counts will be reported
+	app.SetLLMTokenCountCallback(func(modelName string, content string) int {
+		var tokensPerMessage, tokensPerName int
+		switch modelName {
+		case "gpt-3.5-turbo-0613",
+			"gpt-3.5-turbo-16k-0613",
+			"gpt-4-0314",
+			"gpt-4-32k-0314",
+			"gpt-4-0613",
+			"gpt-4-32k-0613":
+			tokensPerMessage = 3
+			tokensPerName = 1
+		case "gpt-3.5-turbo-0301":
+			tokensPerMessage = 4
+			tokensPerName = -1
+		}
+
+		tkm, err := tiktoken.EncodingForModel(modelName)
+		if err != nil {
+			fmt.Println("error getting tokens", err)
+			return 0
+		}
+		token := tkm.Encode(content, nil, nil)
+		totalTokens := len(token) + tokensPerMessage + tokensPerName
+		return totalTokens
+	})
+
 	// OpenAI Config - Additionally, NRDefaultAzureConfig(apiKey, baseURL string) can be used for Azure
 	cfg := nropenai.NRDefaultConfig(os.Getenv("OPEN_AI_API_KEY"))
 
@@ -40,13 +69,13 @@ func main() {
 
 	// GPT Request
 	req := openai.ChatCompletionRequest{
-		Model:       openai.GPT3Dot5Turbo,
+		Model:       openai.GPT4,
 		Temperature: 0.7,
 		MaxTokens:   150,
 		Messages: []openai.ChatCompletionMessage{
 			{
 				Role:    openai.ChatMessageRoleUser,
-				Content: "What is 8*5",
+				Content: "What is Observability in Software Engineering?",
 			},
 		},
 	}
@@ -56,8 +85,9 @@ func main() {
 	if err != nil {
 		panic(err)
 	}
-
-	fmt.Println(resp.ChatCompletionResponse.Choices[0].Message.Content)
+	if len(resp.ChatCompletionResponse.Choices) == 0 {
+		fmt.Println("No choices returned")
+	}
 
 	// Shutdown Application
 	app.Shutdown(5 * time.Second)

diff --git a/v3/integrations/nropenai/examples/chatcompletionfeedback/chatcompletionfeedback.go b/v3/integrations/nropenai/examples/chatcompletionfeedback/chatcompletionfeedback.go
@@ -49,7 +49,7 @@ func main() {
 		Messages: []openai.ChatCompletionMessage{
 			{
 				Role:    openai.ChatMessageRoleUser,
-				Content: "What is 8*5",
+				Content: "What is observability in software engineering?",
 			},
 		},
 	}

diff --git a/v3/integrations/nropenai/examples/embeddings/embeddings_example.go b/v3/integrations/nropenai/examples/embeddings/embeddings_example.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/newrelic/go-agent/v3/integrations/nropenai"
 	"github.com/newrelic/go-agent/v3/newrelic"
+	"github.com/pkoukk/tiktoken-go"
 	openai "github.com/sashabaranov/go-openai"
 )
 
@@ -23,7 +24,31 @@ func main() {
 		panic(err)
 	}
 	app.WaitForConnection(10 * time.Second)
+	app.SetLLMTokenCountCallback(func(modelName string, content string) int {
+		var tokensPerMessage, tokensPerName int
+		switch modelName {
+		case "gpt-3.5-turbo-0613",
+			"gpt-3.5-turbo-16k-0613",
+			"gpt-4-0314",
+			"gpt-4-32k-0314",
+			"gpt-4-0613",
+			"gpt-4-32k-0613":
+			tokensPerMessage = 3
+			tokensPerName = 1
+		case "gpt-3.5-turbo-0301":
+			tokensPerMessage = 4
+			tokensPerName = -1
+		}
 
+		tkm, err := tiktoken.EncodingForModel(modelName)
+		if err != nil {
+			fmt.Println("error getting tokens", err)
+			return 0
+		}
+		token := tkm.Encode(content, nil, nil)
+		totalTokens := len(token) + tokensPerMessage + tokensPerName
+		return totalTokens
+	})
 	// OpenAI Config - Additionally, NRDefaultAzureConfig(apiKey, baseURL string) can be used for Azure
 	cfg := nropenai.NRDefaultConfig(os.Getenv("OPEN_AI_API_KEY"))
 

diff --git a/v3/integrations/nropenai/go.mod b/v3/integrations/nropenai/go.mod
@@ -5,10 +5,12 @@ go 1.21.0
 require (
 	github.com/google/uuid v1.6.0
 	github.com/newrelic/go-agent/v3 v3.30.0
+	github.com/pkoukk/tiktoken-go v0.1.6
 	github.com/sashabaranov/go-openai v1.20.2
 )
 
 require (
+	github.com/dlclark/regexp2 v1.10.0 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	golang.org/x/net v0.9.0 // indirect
 	golang.org/x/sys v0.7.0 // indirect
@@ -18,4 +20,4 @@ require (
 	google.golang.org/protobuf v1.30.0 // indirect
 )
 
-replace github.com/newrelic/go-agent/v3 => ../..
+replace github.com/newrelic/go-agent/v3 => ../..