From a3dbcec3fccbeb9a77b1846f912ddbed6ebda5fd Mon Sep 17 00:00:00 2001 From: mirackara Date: Mon, 25 Mar 2024 14:58:14 -0500 Subject: [PATCH 1/5] openai changes --- v3/integrations/nropenai/nropenai.go | 105 +++++++++++++++++++--- v3/integrations/nropenai/nropenai_test.go | 93 +++++++++++-------- 2 files changed, 151 insertions(+), 47 deletions(-) diff --git a/v3/integrations/nropenai/nropenai.go b/v3/integrations/nropenai/nropenai.go index 8549780c1..0a1cf1434 100644 --- a/v3/integrations/nropenai/nropenai.go +++ b/v3/integrations/nropenai/nropenai.go @@ -149,23 +149,49 @@ type ChatCompletionResponseWrapper struct { // Wrapper for ChatCompletionStream that is returned from NRCreateChatCompletionStream type ChatCompletionStreamWrapper struct { - stream *openai.ChatCompletionStream - txn *newrelic.Transaction + app *newrelic.Application + stream *openai.ChatCompletionStream + streamResp openai.ChatCompletionResponse + responseStr string + uuid string + txn *newrelic.Transaction + cw *ClientWrapper + role string + model string + StreamingData map[string]interface{} + isRoleAdded bool + TraceID string } // Wrapper for Recv() method that calls the underlying stream's Recv() method func (w *ChatCompletionStreamWrapper) Recv() (openai.ChatCompletionStreamResponse, error) { response, err := w.stream.Recv() - if err != nil { return response, err } + if !w.isRoleAdded && (response.Choices[0].Delta.Role == "assistant" || response.Choices[0].Delta.Role == "user" || response.Choices[0].Delta.Role == "system") { + w.isRoleAdded = true + w.role = response.Choices[0].Delta.Role + + } + if response.Choices[0].FinishReason != "stop" { + w.responseStr += response.Choices[0].Delta.Content + w.streamResp.ID = response.ID + w.streamResp.Model = response.Model + w.model = response.Model + } return response, nil } func (w *ChatCompletionStreamWrapper) Close() { + w.StreamingData["response.model"] = w.model + w.app.RecordCustomEvent("LlmChatCompletionSummary", w.StreamingData) + + NRCreateChatCompletionMessageStream(w.app, uuid.MustParse(w.uuid), w, w.cw) + + w.txn.End() w.stream.Close() } @@ -188,7 +214,8 @@ func NRCreateChatCompletionSummary(txn *newrelic.Transaction, app *newrelic.Appl } } // Start span - integrationsupport.AddAgentAttribute(txn, "llm", "", true) + txn.AddAttribute("llm", true) + chatCompletionSpan := txn.StartSegment("Llm/completion/OpenAI/CreateChatCompletion") // Track Total time taken for the chat completion or embedding call to complete in milliseconds start := time.Now() @@ -272,6 +299,52 @@ func NRCreateChatCompletionSummary(txn *newrelic.Transaction, app *newrelic.Appl TraceID: traceID, } } + +func NRCreateChatCompletionMessageStream(app *newrelic.Application, uuid uuid.UUID, sw *ChatCompletionStreamWrapper, cw *ClientWrapper) { + + spanID := sw.txn.GetTraceMetadata().SpanID + traceID := sw.txn.GetTraceMetadata().TraceID + + appCfg, configErr := app.Config() + if !configErr { + appCfg.AppName = "Unknown" + } + integrationsupport.AddAgentAttribute(sw.txn, "llm", "", true) + chatCompletionMessageSpan := sw.txn.StartSegment("Llm/completion/OpenAI/CreateChatCompletionMessageStream") + + ChatCompletionMessageData := map[string]interface{}{} + // if the response doesn't have an ID, use the UUID from the summary + + ChatCompletionMessageData["id"] = sw.streamResp.ID + + // Response Data + ChatCompletionMessageData["request.model"] = sw.model + + if appCfg.AIMonitoring.RecordContent.Enabled { + ChatCompletionMessageData["content"] = sw.responseStr + } + + ChatCompletionMessageData["role"] = sw.role + + // New Relic Attributes + ChatCompletionMessageData["sequence"] = 1 + ChatCompletionMessageData["vendor"] = "OpenAI" + ChatCompletionMessageData["ingest_source"] = "Go" + ChatCompletionMessageData["span_id"] = spanID + ChatCompletionMessageData["trace_id"] = traceID + contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(sw.model, sw.responseStr) + if contentCounted { + ChatCompletionMessageData["token_count"] = contentTokens + } + + // If custom attributes are set, add them to the data + ChatCompletionMessageData = AppendCustomAttributesToEvent(cw, ChatCompletionMessageData) + chatCompletionMessageSpan.End() + // Record Custom Event for each message + app.RecordCustomEvent("LlmChatCompletionMessage", ChatCompletionMessageData) + +} + func NRCreateChatCompletionMessageInput(txn *newrelic.Transaction, app *newrelic.Application, req openai.ChatCompletionRequest, uuid uuid.UUID, cw *ClientWrapper) { spanID := txn.GetTraceMetadata().SpanID traceID := txn.GetTraceMetadata().TraceID @@ -297,8 +370,8 @@ func NRCreateChatCompletionMessageInput(txn *newrelic.Transaction, app *newrelic // New Relic Attributes ChatCompletionMessageData["sequence"] = 0 - ChatCompletionMessageData["vendor"] = "openai" - ChatCompletionMessageData["ingest_source"] = "go" + ChatCompletionMessageData["vendor"] = "OpenAI" + ChatCompletionMessageData["ingest_source"] = "Go" ChatCompletionMessageData["span_id"] = spanID ChatCompletionMessageData["trace_id"] = traceID contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(req.Model, req.Messages[0].Content) @@ -348,8 +421,8 @@ func NRCreateChatCompletionMessage(txn *newrelic.Transaction, app *newrelic.Appl // New Relic Attributes ChatCompletionMessageData["sequence"] = i + 1 - ChatCompletionMessageData["vendor"] = "openai" - ChatCompletionMessageData["ingest_source"] = "go" + ChatCompletionMessageData["vendor"] = "OpenAI" + ChatCompletionMessageData["ingest_source"] = "Go" ChatCompletionMessageData["span_id"] = spanID ChatCompletionMessageData["trace_id"] = traceID tokenCount, tokensCounted := TokenCountingHelper(app, choice.Message, resp.Model) @@ -378,6 +451,16 @@ func TokenCountingHelper(app *newrelic.Application, message openai.ChatCompletio return numTokens, (contentCounted && roleCounted && messageCounted) } +func TokenCountingHelperStream(app *newrelic.Application, model string, content string, role string, messageName string) (numTokens int, tokensCounted bool) { + + contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(model, content) + roleTokens, roleCounted := app.InvokeLLMTokenCountCallback(model, role) + messageTokens, messageCounted := app.InvokeLLMTokenCountCallback(model, messageName) + numTokens += contentTokens + roleTokens + messageTokens + + return numTokens, (contentCounted && roleCounted && messageCounted) +} + // NRCreateChatCompletion is a wrapper for the OpenAI CreateChatCompletion method. // If AI Monitoring is disabled, the wrapped function will still call the OpenAI CreateChatCompletion method and return the response with no New Relic instrumentation func NRCreateChatCompletion(cw *ClientWrapper, req openai.ChatCompletionRequest, app *newrelic.Application) (ChatCompletionResponseWrapper, error) { @@ -552,8 +635,8 @@ func NRCreateChatCompletionStream(cw *ClientWrapper, ctx context.Context, req op StreamingData["vendor"] = "OpenAI" StreamingData["ingest_source"] = "Go" StreamingData["appName"] = config.AppName - app.RecordCustomEvent("LlmChatCompletionSummary", StreamingData) - txn.End() - return &ChatCompletionStreamWrapper{stream: stream, txn: txn}, nil + + NRCreateChatCompletionMessageInput(txn, app, req, uuid, cw) + return &ChatCompletionStreamWrapper{app: app, stream: stream, txn: txn, uuid: uuid.String(), cw: cw, StreamingData: StreamingData, TraceID: traceID}, nil } diff --git a/v3/integrations/nropenai/nropenai_test.go b/v3/integrations/nropenai/nropenai_test.go index 68498eaff..76c762b6a 100644 --- a/v3/integrations/nropenai/nropenai_test.go +++ b/v3/integrations/nropenai/nropenai_test.go @@ -251,7 +251,7 @@ func TestNRCreateChatCompletion(t *testing.T) { "request_id": "chatcmpl-123", "request.model": "gpt-3.5-turbo", "request.max_tokens": 150, - "response.number_of_messages": 1, + "response.number_of_messages": 2, "response.headers.llmVersion": "2020-10-01", "response.organization": "user-123", "response.model": "gpt-3.5-turbo", @@ -271,13 +271,31 @@ func TestNRCreateChatCompletion(t *testing.T) { UserAttributes: map[string]interface{}{ "trace_id": internal.MatchAnything, "span_id": internal.MatchAnything, - "id": "chatcmpl-123", + "id": internal.MatchAnything, "sequence": 0, + "role": "user", + "content": "What is 8*5", + "vendor": "OpenAI", + "ingest_source": "Go", + "response.model": "gpt-3.5-turbo", + }, + AgentAttributes: map[string]interface{}{}, + }, + { + Intrinsics: map[string]interface{}{ + "type": "LlmChatCompletionMessage", + "timestamp": internal.MatchAnything, + }, + UserAttributes: map[string]interface{}{ + "trace_id": internal.MatchAnything, + "span_id": internal.MatchAnything, + "id": "chatcmpl-123", + "sequence": 1, "role": "assistant", "content": "\n\nHello there, how may I assist you today?", "request_id": "chatcmpl-123", - "vendor": "openai", - "ingest_source": "go", + "vendor": "OpenAI", + "ingest_source": "Go", "response.model": "gpt-3.5-turbo", }, AgentAttributes: map[string]interface{}{}, @@ -359,7 +377,7 @@ func TestNRCreateChatCompletionError(t *testing.T) { "request_id": "", "request.model": "gpt-3.5-turbo", "request.max_tokens": 150, - "response.number_of_messages": 0, + "response.number_of_messages": 1, "response.headers.llmVersion": "2020-10-01", "response.organization": "user-123", "response.model": "", @@ -371,6 +389,23 @@ func TestNRCreateChatCompletionError(t *testing.T) { "response.headers.ratelimitLimitRequests": "10000", }, }, + { + Intrinsics: map[string]interface{}{ + "type": "LlmChatCompletionMessage", + "timestamp": internal.MatchAnything, + }, + UserAttributes: map[string]interface{}{ + "ingest_source": "Go", + "vendor": "OpenAI", + "id": internal.MatchAnything, + "trace_id": internal.MatchAnything, + "span_id": internal.MatchAnything, + "content": "testError", + "role": "user", + "response.model": "gpt-3.5-turbo", + "sequence": 0, + }, + }, }) app.ExpectErrorEvents(t, []internal.WantEvent{ { @@ -385,11 +420,11 @@ func TestNRCreateChatCompletionError(t *testing.T) { "error.message": "test error", }, UserAttributes: map[string]interface{}{ - "error.code": "404", - "http.status": "404", "completion_id": internal.MatchAnything, + "llm": true, }, - }}) + }, + }) } func TestNRCreateEmbedding(t *testing.T) { mockClient := &MockOpenAIClient{} @@ -425,6 +460,7 @@ func TestNRCreateEmbedding(t *testing.T) { "trace_id": internal.MatchAnything, "span_id": internal.MatchAnything, "duration": 0, + "request_id": "chatcmpl-123", "api_key_last_four_digits": "sk-mnop", "request.model": "text-embedding-ada-002", "response.headers.llmVersion": "2020-10-01", @@ -504,6 +540,7 @@ func TestNRCreateEmbeddingError(t *testing.T) { "span_id": internal.MatchAnything, "duration": 0, "api_key_last_four_digits": "sk-mnop", + "request_id": "chatcmpl-123", "request.model": "text-embedding-ada-002", "response.headers.llmVersion": "2020-10-01", "response.organization": "user-123", @@ -533,8 +570,6 @@ func TestNRCreateEmbeddingError(t *testing.T) { "error.message": "test error", }, UserAttributes: map[string]interface{}{ - "error.code": "404", - "http.status": "404", "embedding_id": internal.MatchAnything, }, }}) @@ -566,36 +601,21 @@ func TestNRCreateStream(t *testing.T) { app.ExpectCustomEvents(t, []internal.WantEvent{ { Intrinsics: map[string]interface{}{ - "type": "LlmChatCompletionSummary", + "type": "LlmChatCompletionMessage", "timestamp": internal.MatchAnything, }, UserAttributes: map[string]interface{}{ - "ingest_source": "Go", - "vendor": "OpenAI", - "model": "gpt-3.5-turbo", - "id": internal.MatchAnything, - "trace_id": internal.MatchAnything, - "span_id": internal.MatchAnything, - "appName": "my app", - "duration": 0, - "request.temperature": 0, - "api_key_last_four_digits": "sk-mnop", - "request.max_tokens": 1500, - "request.model": "gpt-3.5-turbo", - }, - }, - }) - app.ExpectTxnEvents(t, []internal.WantEvent{ - { - Intrinsics: map[string]interface{}{ - "type": "Transaction", - "name": "OtherTransaction/Go/OpenAIChatCompletionStream", - "timestamp": internal.MatchAnything, - "traceId": internal.MatchAnything, - "priority": internal.MatchAnything, - "sampled": internal.MatchAnything, - "guid": internal.MatchAnything, + "trace_id": internal.MatchAnything, + "span_id": internal.MatchAnything, + "id": internal.MatchAnything, + "sequence": 0, + "role": "user", + "content": "Say this is a test", + "vendor": "OpenAI", + "ingest_source": "Go", + "response.model": "gpt-3.5-turbo", }, + AgentAttributes: map[string]interface{}{}, }, }) } @@ -625,6 +645,7 @@ func TestNRCreateStreamAIMonitoringNotEnabled(t *testing.T) { } app.ExpectCustomEvents(t, []internal.WantEvent{}) app.ExpectTxnEvents(t, []internal.WantEvent{}) + } func TestNRCreateStreamError(t *testing.T) { From 13475968633e91419a5bbf0aee1ace55c14406a5 Mon Sep 17 00:00:00 2001 From: mirackara Date: Mon, 25 Mar 2024 15:01:34 -0500 Subject: [PATCH 2/5] Close stream after response is finished --- .../chatcompletionstreaming/chatcompletionstreaming.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go index 4745aae3b..e4235930f 100644 --- a/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go +++ b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go @@ -61,8 +61,6 @@ func main() { if err != nil { panic(err) } - defer stream.Close() - fmt.Printf("Stream response: ") for { var response openai.ChatCompletionStreamResponse @@ -78,6 +76,7 @@ func main() { fmt.Printf(response.Choices[0].Delta.Content) } + stream.Close() // Shutdown Application app.Shutdown(5 * time.Second) } From 7c7f73834fbf6023f0805780e0911d4499456aaa Mon Sep 17 00:00:00 2001 From: mirackara Date: Mon, 25 Mar 2024 15:04:54 -0500 Subject: [PATCH 3/5] added feedback to streaming example --- .../chatcompletionstreaming.go | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go index e4235930f..50f73ec13 100644 --- a/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go +++ b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go @@ -13,12 +13,26 @@ import ( openai "github.com/sashabaranov/go-openai" ) +// Simulates feedback being sent to New Relic. Feedback on a chat completion requires +// having access to the ChatCompletionResponseWrapper which is returned by the NRCreateChatCompletion function. +func SendFeedback(app *newrelic.Application, resp nropenai.ChatCompletionStreamWrapper) { + trace_id := resp.TraceID + rating := "5" + category := "informative" + message := "The response was concise yet thorough." + customMetadata := map[string]interface{}{ + "foo": "bar", + "pi": 3.14, + } + + app.RecordLLMFeedbackEvent(trace_id, rating, category, message, customMetadata) +} + func main() { // Start New Relic Application app, err := newrelic.NewApplication( newrelic.ConfigAppName("Basic OpenAI App"), newrelic.ConfigLicense(os.Getenv("NEW_RELIC_LICENSE_KEY")), - newrelic.ConfigDebugLogger(os.Stdout), // Enable AI Monitoring // NOTE - If High Security Mode is enabled, AI Monitoring will always be disabled newrelic.ConfigAIMonitoringEnabled(true), @@ -49,7 +63,7 @@ func main() { Messages: []openai.ChatCompletionMessage{ { Role: openai.ChatMessageRoleUser, - Content: "Say this is a test", + Content: "What is observability in software engineering?", }, }, Stream: true, @@ -59,6 +73,7 @@ func main() { stream, err := nropenai.NRCreateChatCompletionStream(client, ctx, req, app) if err != nil { + panic(err) } fmt.Printf("Stream response: ") @@ -77,6 +92,7 @@ func main() { fmt.Printf(response.Choices[0].Delta.Content) } stream.Close() + SendFeedback(app, *stream) // Shutdown Application app.Shutdown(5 * time.Second) } From 95cd5b8ea36fb14af7f21b5757acf9302e09f899 Mon Sep 17 00:00:00 2001 From: mirackara Date: Mon, 25 Mar 2024 15:12:53 -0500 Subject: [PATCH 4/5] token callback extra check --- v3/integrations/nropenai/nropenai.go | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/v3/integrations/nropenai/nropenai.go b/v3/integrations/nropenai/nropenai.go index 0a1cf1434..36b262586 100644 --- a/v3/integrations/nropenai/nropenai.go +++ b/v3/integrations/nropenai/nropenai.go @@ -333,8 +333,10 @@ func NRCreateChatCompletionMessageStream(app *newrelic.Application, uuid uuid.UU ChatCompletionMessageData["span_id"] = spanID ChatCompletionMessageData["trace_id"] = traceID contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(sw.model, sw.responseStr) - if contentCounted { - ChatCompletionMessageData["token_count"] = contentTokens + roleTokens, roleCounted := app.InvokeLLMTokenCountCallback(sw.model, sw.role) + + if (contentCounted && roleCounted) && app.HasLLMTokenCountCallback() { + ChatCompletionMessageData["token_count"] = contentTokens + roleTokens } // If custom attributes are set, add them to the data @@ -376,7 +378,7 @@ func NRCreateChatCompletionMessageInput(txn *newrelic.Transaction, app *newrelic ChatCompletionMessageData["trace_id"] = traceID contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(req.Model, req.Messages[0].Content) - if contentCounted { + if contentCounted && app.HasLLMTokenCountCallback() { ChatCompletionMessageData["token_count"] = contentTokens } @@ -548,7 +550,8 @@ func NRCreateEmbedding(cw *ClientWrapper, req openai.EmbeddingRequest, app *newr // cast input as string input := GetInput(req.Input).(string) tokenCount, tokensCounted := app.InvokeLLMTokenCountCallback(string(resp.Model), input) - if tokensCounted { + + if tokensCounted && app.HasLLMTokenCountCallback() { EmbeddingsData["token_count"] = tokenCount } From 8bd6b0bb365b1f3d65e3ace9ace707864ef1b16a Mon Sep 17 00:00:00 2001 From: mirackara Date: Mon, 25 Mar 2024 15:47:13 -0500 Subject: [PATCH 5/5] Corrected token count issue and added token callback for streaming --- .../chatcompletionstreaming.go | 29 ++++++++++++++++- v3/integrations/nropenai/nropenai.go | 31 +++++++++---------- 2 files changed, 42 insertions(+), 18 deletions(-) diff --git a/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go index 50f73ec13..f43ddba70 100644 --- a/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go +++ b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go @@ -10,6 +10,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" ) @@ -41,7 +42,33 @@ func main() { panic(err) } 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")) @@ -57,7 +84,7 @@ func main() { // GPT Request req := openai.ChatCompletionRequest{ - Model: openai.GPT3Dot5Turbo, + Model: openai.GPT4, Temperature: 0.7, MaxTokens: 1500, Messages: []openai.ChatCompletionMessage{ diff --git a/v3/integrations/nropenai/nropenai.go b/v3/integrations/nropenai/nropenai.go index 36b262586..8c2b6c248 100644 --- a/v3/integrations/nropenai/nropenai.go +++ b/v3/integrations/nropenai/nropenai.go @@ -332,11 +332,15 @@ func NRCreateChatCompletionMessageStream(app *newrelic.Application, uuid uuid.UU ChatCompletionMessageData["ingest_source"] = "Go" ChatCompletionMessageData["span_id"] = spanID ChatCompletionMessageData["trace_id"] = traceID - contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(sw.model, sw.responseStr) - roleTokens, roleCounted := app.InvokeLLMTokenCountCallback(sw.model, sw.role) - - if (contentCounted && roleCounted) && app.HasLLMTokenCountCallback() { - ChatCompletionMessageData["token_count"] = contentTokens + roleTokens + tmpMessage := openai.ChatCompletionMessage{ + Content: sw.responseStr, + Role: sw.role, + // Name is not provided in the stream response, so we don't include it in token counting + Name: "", + } + tokenCount, tokensCounted := TokenCountingHelper(app, tmpMessage, sw.model) + if tokensCounted { + ChatCompletionMessageData["token_count"] = tokenCount } // If custom attributes are set, add them to the data @@ -444,23 +448,16 @@ func NRCreateChatCompletionMessage(txn *newrelic.Transaction, app *newrelic.Appl } func TokenCountingHelper(app *newrelic.Application, message openai.ChatCompletionMessage, model string) (numTokens int, tokensCounted bool) { - contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(model, message.Content) roleTokens, roleCounted := app.InvokeLLMTokenCountCallback(model, message.Role) - messageTokens, messageCounted := app.InvokeLLMTokenCountCallback(model, message.Name) - numTokens += contentTokens + roleTokens + messageTokens + var messageTokens int + if message.Name != "" { + messageTokens, _ = app.InvokeLLMTokenCountCallback(model, message.Name) - return numTokens, (contentCounted && roleCounted && messageCounted) -} - -func TokenCountingHelperStream(app *newrelic.Application, model string, content string, role string, messageName string) (numTokens int, tokensCounted bool) { - - contentTokens, contentCounted := app.InvokeLLMTokenCountCallback(model, content) - roleTokens, roleCounted := app.InvokeLLMTokenCountCallback(model, role) - messageTokens, messageCounted := app.InvokeLLMTokenCountCallback(model, messageName) + } numTokens += contentTokens + roleTokens + messageTokens - return numTokens, (contentCounted && roleCounted && messageCounted) + return numTokens, (contentCounted && roleCounted) } // NRCreateChatCompletion is a wrapper for the OpenAI CreateChatCompletion method.