Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@mirackara
mirackara committed Mar 22, 2024
1 parent 63690ec commit cb5a85d
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 77 deletions.
2 changes: 2 additions & 0 deletions v3/integrations/nropenai/examples/chatcompletion/chatcompletion_example.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ 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 {
Expand Down
138 changes: 61 additions & 77 deletions v3/integrations/nropenai/nropenai_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -237,28 +237,24 @@ func TestNRCreateChatCompletion(t *testing.T) {
"timestamp": internal.MatchAnything,
},
UserAttributes: map[string]interface{}{
"ingest_source": "Go",
"vendor": "OpenAI",
"model": "gpt-3.5-turbo",
"id": internal.MatchAnything,
"transaction_id": internal.MatchAnything,
"trace_id": internal.MatchAnything,
"span_id": internal.MatchAnything,
"appName": "my app",
"duration": 0,
"response.choices.finish_reason": internal.MatchAnything,
"request.temperature": 0,
"api_key_last_four_digits": "sk-mnop",
"request_id": "chatcmpl-123",
"request.model": "gpt-3.5-turbo",
"request.max_tokens": 150,
"response.number_of_messages": 1,
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"response.usage.completion_tokens": 12,
"response.model": "gpt-3.5-turbo",
"response.usage.total_tokens": 21,
"response.usage.prompt_tokens": 9,
"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,
"response.choices.finish_reason": internal.MatchAnything,
"request.temperature": 0,
"api_key_last_four_digits": "sk-mnop",
"request_id": "chatcmpl-123",
"request.model": "gpt-3.5-turbo",
"request.max_tokens": 150,
"response.number_of_messages": 1,
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"response.model": "gpt-3.5-turbo",
"response.headers.ratelimitRemainingTokens": "100",
"response.headers.ratelimitRemainingRequests": "10000",
"response.headers.ratelimitResetTokens": "100",
Expand All @@ -274,7 +270,6 @@ func TestNRCreateChatCompletion(t *testing.T) {
},
UserAttributes: map[string]interface{}{
"trace_id": internal.MatchAnything,
"transaction_id": internal.MatchAnything,
"span_id": internal.MatchAnything,
"id": "chatcmpl-123",
"sequence": 0,
Expand Down Expand Up @@ -350,28 +345,24 @@ func TestNRCreateChatCompletionError(t *testing.T) {
"timestamp": internal.MatchAnything,
},
UserAttributes: map[string]interface{}{
"error": true,
"ingest_source": "Go",
"vendor": "OpenAI",
"model": "gpt-3.5-turbo",
"id": internal.MatchAnything,
"transaction_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_id": "",
"request.model": "gpt-3.5-turbo",
"request.max_tokens": 150,
"response.number_of_messages": 0,
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"response.usage.completion_tokens": 0,
"response.model": "",
"response.usage.total_tokens": 0,
"response.usage.prompt_tokens": 0,
"error": true,
"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_id": "",
"request.model": "gpt-3.5-turbo",
"request.max_tokens": 150,
"response.number_of_messages": 0,
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"response.model": "",
"response.headers.ratelimitRemainingTokens": "100",
"response.headers.ratelimitRemainingRequests": "10000",
"response.headers.ratelimitResetTokens": "100",
Expand Down Expand Up @@ -428,21 +419,18 @@ func TestNRCreateEmbedding(t *testing.T) {
"timestamp": internal.MatchAnything,
},
UserAttributes: map[string]interface{}{
"ingest_source": "Go",
"vendor": "OpenAI",
"id": internal.MatchAnything,
"transaction_id": internal.MatchAnything,
"trace_id": internal.MatchAnything,
"span_id": internal.MatchAnything,
"duration": 0,
"api_key_last_four_digits": "sk-mnop",
"request.model": "text-embedding-ada-002",
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"response.model": "text-embedding-ada-002",
"response.usage.total_tokens": 21,
"response.usage.prompt_tokens": 9,
"input": "The food was delicious and the waiter",
"ingest_source": "Go",
"vendor": "OpenAI",
"id": internal.MatchAnything,
"trace_id": internal.MatchAnything,
"span_id": internal.MatchAnything,
"duration": 0,
"api_key_last_four_digits": "sk-mnop",
"request.model": "text-embedding-ada-002",
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"response.model": "text-embedding-ada-002",
"input": "The food was delicious and the waiter",
"response.headers.ratelimitRemainingTokens": "100",
"response.headers.ratelimitRemainingRequests": "10000",
"response.headers.ratelimitResetTokens": "100",
Expand Down Expand Up @@ -509,22 +497,19 @@ func TestNRCreateEmbeddingError(t *testing.T) {
"timestamp": internal.MatchAnything,
},
UserAttributes: map[string]interface{}{
"ingest_source": "Go",
"vendor": "OpenAI",
"id": internal.MatchAnything,
"transaction_id": internal.MatchAnything,
"trace_id": internal.MatchAnything,
"span_id": internal.MatchAnything,
"duration": 0,
"api_key_last_four_digits": "sk-mnop",
"request.model": "text-embedding-ada-002",
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"error": true,
"response.model": "",
"response.usage.total_tokens": 0,
"response.usage.prompt_tokens": 0,
"input": "testError",
"ingest_source": "Go",
"vendor": "OpenAI",
"id": internal.MatchAnything,
"trace_id": internal.MatchAnything,
"span_id": internal.MatchAnything,
"duration": 0,
"api_key_last_four_digits": "sk-mnop",
"request.model": "text-embedding-ada-002",
"response.headers.llmVersion": "2020-10-01",
"response.organization": "user-123",
"error": true,
"response.model": "",
"input": "testError",
"response.headers.ratelimitRemainingTokens": "100",
"response.headers.ratelimitRemainingRequests": "10000",
"response.headers.ratelimitResetTokens": "100",
Expand Down Expand Up @@ -589,7 +574,6 @@ func TestNRCreateStream(t *testing.T) {
"vendor": "OpenAI",
"model": "gpt-3.5-turbo",
"id": internal.MatchAnything,
"transaction_id": internal.MatchAnything,
"trace_id": internal.MatchAnything,
"span_id": internal.MatchAnything,
"appName": "my app",
Expand Down

0 comments on commit cb5a85d

Please sign in to comment.