diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 9fd0c4ad3..efaae3ba6 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -66,6 +66,7 @@ jobs: - dirs: v3/integrations/nrmongo - dirs: v3/integrations/nrgraphqlgo,v3/integrations/nrgraphqlgo/example - dirs: v3/integrations/nrmssql + - dirs: v3/integrations/nropenai steps: - name: Checkout Code uses: actions/checkout@v2 diff --git a/v3/integrations/nropenai/examples/chatcompletion/chatcompletion_example.go b/v3/integrations/nropenai/examples/chatcompletion/chatcompletion_example.go new file mode 100644 index 000000000..6aeef8da1 --- /dev/null +++ b/v3/integrations/nropenai/examples/chatcompletion/chatcompletion_example.go @@ -0,0 +1,64 @@ +package main + +import ( + "fmt" + "os" + "time" + + "github.com/newrelic/go-agent/v3/integrations/nropenai" + "github.com/newrelic/go-agent/v3/newrelic" + openai "github.com/sashabaranov/go-openai" +) + +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), + ) + if nil != err { + panic(err) + } + app.WaitForConnection(10 * time.Second) + + // OpenAI Config - Additionally, NRDefaultAzureConfig(apiKey, baseURL string) can be used for Azure + cfg := nropenai.NRDefaultConfig(os.Getenv("OPEN_AI_API_KEY")) + + // Create OpenAI Client - Additionally, NRNewClient(authToken string) can be used + client := nropenai.NRNewClientWithConfig(cfg) + + // Add any custom attributes + // NOTE: Attributes must start with "llm.", otherwise they will be ignored + client.AddCustomAttributes(map[string]interface{}{ + "llm.foo": "bar", + "llm.pi": 3.14, + }) + + // GPT Request + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0.7, + MaxTokens: 150, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "What is 8*5", + }, + }, + } + // NRCreateChatCompletion returns a wrapped version of openai.ChatCompletionResponse + resp, err := nropenai.NRCreateChatCompletion(client, req, app) + + if err != nil { + panic(err) + } + + fmt.Println(resp.ChatCompletionResponse.Choices[0].Message.Content) + + // 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 new file mode 100644 index 000000000..24eba417d --- /dev/null +++ b/v3/integrations/nropenai/examples/chatcompletionfeedback/chatcompletionfeedback.go @@ -0,0 +1,68 @@ +package main + +import ( + "fmt" + "os" + "time" + + "github.com/newrelic/go-agent/v3/integrations/nropenai" + "github.com/newrelic/go-agent/v3/newrelic" + 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.ChatCompletionResponseWrapper) { + 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), + newrelic.ConfigAIMonitoringEnabled(true), + ) + if nil != err { + panic(err) + } + app.WaitForConnection(10 * time.Second) + + // OpenAI Config - Additionally, NRDefaultAzureConfig(apiKey, baseURL string) can be used for Azure + cfg := nropenai.NRDefaultConfig(os.Getenv("OPEN_AI_API_KEY")) + client := nropenai.NRNewClientWithConfig(cfg) + // GPT Request + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0.7, + MaxTokens: 150, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "What is 8*5", + }, + }, + } + // NRCreateChatCompletion returns a wrapped version of openai.ChatCompletionResponse + resp, err := nropenai.NRCreateChatCompletion(client, req, app) + + if err != nil { + panic(err) + } + // Print the contents of the message + fmt.Println("Message Response: ", resp.ChatCompletionResponse.Choices[0].Message.Content) + SendFeedback(app, resp) + + // Shutdown Application + app.Shutdown(5 * time.Second) +} diff --git a/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go new file mode 100644 index 000000000..4745aae3b --- /dev/null +++ b/v3/integrations/nropenai/examples/chatcompletionstreaming/chatcompletionstreaming.go @@ -0,0 +1,83 @@ +package main + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "time" + + "github.com/newrelic/go-agent/v3/integrations/nropenai" + "github.com/newrelic/go-agent/v3/newrelic" + openai "github.com/sashabaranov/go-openai" +) + +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), + ) + if nil != err { + panic(err) + } + app.WaitForConnection(10 * time.Second) + + // OpenAI Config - Additionally, NRDefaultAzureConfig(apiKey, baseURL string) can be used for Azure + cfg := nropenai.NRDefaultConfig(os.Getenv("OPEN_AI_API_KEY")) + + // Create OpenAI Client - Additionally, NRNewClient(authToken string) can be used + client := nropenai.NRNewClientWithConfig(cfg) + + // Add any custom attributes + // NOTE: Attributes must start with "llm.", otherwise they will be ignored + client.AddCustomAttributes(map[string]interface{}{ + "llm.foo": "bar", + "llm.pi": 3.14, + }) + + // GPT Request + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0.7, + MaxTokens: 1500, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "Say this is a test", + }, + }, + Stream: true, + } + ctx := context.Background() + + stream, err := nropenai.NRCreateChatCompletionStream(client, ctx, req, app) + + if err != nil { + panic(err) + } + defer stream.Close() + + fmt.Printf("Stream response: ") + for { + var response openai.ChatCompletionStreamResponse + response, err = stream.Recv() + if errors.Is(err, io.EOF) { + fmt.Println("\nStream finished") + break + } + if err != nil { + fmt.Printf("\nStream error: %v\n", err) + return + } + + fmt.Printf(response.Choices[0].Delta.Content) + } + // Shutdown Application + app.Shutdown(5 * time.Second) +} diff --git a/v3/integrations/nropenai/examples/embeddings/embeddings_example.go b/v3/integrations/nropenai/examples/embeddings/embeddings_example.go new file mode 100644 index 000000000..ff50d7428 --- /dev/null +++ b/v3/integrations/nropenai/examples/embeddings/embeddings_example.go @@ -0,0 +1,59 @@ +package main + +import ( + "fmt" + "os" + "time" + + "github.com/newrelic/go-agent/v3/integrations/nropenai" + "github.com/newrelic/go-agent/v3/newrelic" + openai "github.com/sashabaranov/go-openai" +) + +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 + newrelic.ConfigAIMonitoringEnabled(true), + ) + if nil != err { + panic(err) + } + app.WaitForConnection(10 * time.Second) + + // OpenAI Config - Additionally, NRDefaultAzureConfig(apiKey, baseURL string) can be used for Azure + cfg := nropenai.NRDefaultConfig(os.Getenv("OPEN_AI_API_KEY")) + + // Create OpenAI Client - Additionally, NRNewClient(authToken string) can be used + client := nropenai.NRNewClientWithConfig(cfg) + + // Add any custom attributes + // NOTE: Attributes must start with "llm.", otherwise they will be ignored + client.CustomAttributes = map[string]interface{}{ + "llm.foo": "bar", + "llm.pi": 3.14, + } + + fmt.Println("Creating Embedding Request...") + // Create Embeddings + embeddingReq := openai.EmbeddingRequest{ + Input: []string{ + "The food was delicious and the waiter", + "Other examples of embedding request", + }, + Model: openai.AdaEmbeddingV2, + EncodingFormat: openai.EmbeddingEncodingFormatFloat, + } + resp, err := nropenai.NRCreateEmbedding(client, embeddingReq, app) + if err != nil { + panic(err) + } + + fmt.Println("Embedding Created!") + fmt.Println(resp.Usage.PromptTokens) + // Shutdown Application + app.Shutdown(5 * time.Second) +} diff --git a/v3/integrations/nropenai/go.mod b/v3/integrations/nropenai/go.mod new file mode 100644 index 000000000..338ad0aef --- /dev/null +++ b/v3/integrations/nropenai/go.mod @@ -0,0 +1,21 @@ +module github.com/newrelic/go-agent/v3/integrations/nropenai + +go 1.21.0 + +require ( + github.com/google/uuid v1.6.0 + github.com/newrelic/go-agent/v3 v3.30.0 + github.com/sashabaranov/go-openai v1.20.2 +) + +require ( + github.com/golang/protobuf v1.5.3 // indirect + golang.org/x/net v0.9.0 // indirect + golang.org/x/sys v0.7.0 // indirect + golang.org/x/text v0.9.0 // indirect + google.golang.org/genproto v0.0.0-20230410155749-daa745c078e1 // indirect + google.golang.org/grpc v1.56.3 // indirect + google.golang.org/protobuf v1.30.0 // indirect +) + +replace github.com/newrelic/go-agent/v3 => ../.. \ No newline at end of file diff --git a/v3/integrations/nropenai/nropenai.go b/v3/integrations/nropenai/nropenai.go new file mode 100644 index 000000000..33054cec9 --- /dev/null +++ b/v3/integrations/nropenai/nropenai.go @@ -0,0 +1,491 @@ +// Copyright 2020 New Relic Corporation. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package nropenai + +import ( + "context" + "errors" + "reflect" + "runtime/debug" + "strings" + "time" + + "github.com/google/uuid" + "github.com/newrelic/go-agent/v3/internal" + "github.com/newrelic/go-agent/v3/newrelic" + "github.com/sashabaranov/go-openai" +) + +func init() { + // Get current go-openai version + info, ok := debug.ReadBuildInfo() + if info != nil && ok { + for _, module := range info.Deps { + if module != nil && strings.Contains(module.Path, "go-openai") { + internal.TrackUsage("Go", "ML", "OpenAI", module.Version) + return + } + } + } + + internal.TrackUsage("Go", "ML", "OpenAI", "unknown") + +} + +var ( + errAIMonitoringDisabled = errors.New("AI Monitoring is set to disabled or High Security Mode is enabled. Please enable AI Monitoring and ensure High Security Mode is disabled") +) + +type OpenAIClient interface { + CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest) (response openai.ChatCompletionResponse, err error) + CreateChatCompletionStream(ctx context.Context, request openai.ChatCompletionRequest) (stream *openai.ChatCompletionStream, err error) + CreateEmbeddings(ctx context.Context, conv openai.EmbeddingRequestConverter) (res openai.EmbeddingResponse, err error) +} + +// Wrapper for OpenAI Configuration +type ConfigWrapper struct { + Config *openai.ClientConfig + LicenseKeyLastFour string +} + +// Wrapper for OpenAI Client with Custom Attributes that can be set for all LLM Events +type ClientWrapper struct { + Client OpenAIClient + LicenseKeyLastFour string + // Set of Custom Attributes that get tied to all LLM Events + CustomAttributes map[string]interface{} +} + +func FormatAPIKey(apiKey string) string { + return "sk-" + apiKey[len(apiKey)-4:] +} + +// Default Config +func NRDefaultConfig(authToken string) *ConfigWrapper { + cfg := openai.DefaultConfig(authToken) + return &ConfigWrapper{ + Config: &cfg, + LicenseKeyLastFour: FormatAPIKey(authToken), + } +} + +// Azure Config +func NRDefaultAzureConfig(apiKey, baseURL string) *ConfigWrapper { + cfg := openai.DefaultAzureConfig(apiKey, baseURL) + return &ConfigWrapper{ + Config: &cfg, + LicenseKeyLastFour: FormatAPIKey(apiKey), + } +} + +// Call to Create Client Wrapper +func NRNewClient(authToken string) *ClientWrapper { + client := openai.NewClient(authToken) + return &ClientWrapper{ + Client: client, + LicenseKeyLastFour: FormatAPIKey(authToken), + } +} + +// NewClientWithConfig creates new OpenAI API client for specified config. +func NRNewClientWithConfig(config *ConfigWrapper) *ClientWrapper { + client := openai.NewClientWithConfig(*config.Config) + return &ClientWrapper{ + Client: client, + LicenseKeyLastFour: config.LicenseKeyLastFour, + } +} + +// Adds Custom Attributes to the ClientWrapper +func (cw *ClientWrapper) AddCustomAttributes(attributes map[string]interface{}) { + if cw.CustomAttributes == nil { + cw.CustomAttributes = make(map[string]interface{}) + } + + for key, value := range attributes { + if strings.HasPrefix(key, "llm.") { + cw.CustomAttributes[key] = value + } + } +} + +func AppendCustomAttributesToEvent(cw *ClientWrapper, data map[string]interface{}) map[string]interface{} { + for k, v := range cw.CustomAttributes { + data[k] = v + } + return data +} + +// If multiple messages are sent, only the first message is used for the "content" field +func GetInput(any interface{}) any { + v := reflect.ValueOf(any) + if v.Kind() == reflect.Array || v.Kind() == reflect.Slice { + if v.Len() > 0 { + // Return the first element + return v.Index(0).Interface() + } + // Input passed in is empty + return "" + } + return any + +} + +// Wrapper for ChatCompletionResponse that is returned from NRCreateChatCompletion. It also includes the TraceID of the transaction for linking a chat response with it's feedback +type ChatCompletionResponseWrapper struct { + ChatCompletionResponse openai.ChatCompletionResponse + TraceID string +} + +// Wrapper for ChatCompletionStream that is returned from NRCreateChatCompletionStream +type ChatCompletionStreamWrapper struct { + stream *openai.ChatCompletionStream + txn *newrelic.Transaction +} + +// 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 + } + + return response, nil + +} + +func (w *ChatCompletionStreamWrapper) Close() { + w.stream.Close() +} + +// NRCreateChatCompletionSummary captures the request and response data for a chat completion request and records a custom event in New Relic. It also captures the completion messages +// With a call to NRCreateChatCompletionMessage +func NRCreateChatCompletionSummary(txn *newrelic.Transaction, app *newrelic.Application, cw *ClientWrapper, req openai.ChatCompletionRequest) ChatCompletionResponseWrapper { + // Get App Config for setting App Name Attribute + appConfig, configErr := app.Config() + if !configErr { + appConfig.AppName = "Unknown" + } + uuid := uuid.New() + spanID := txn.GetTraceMetadata().SpanID + traceID := txn.GetTraceMetadata().TraceID + transactionID := traceID[:16] + + ChatCompletionSummaryData := map[string]interface{}{} + + // Start span + 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() + resp, err := cw.Client.CreateChatCompletion( + context.Background(), + req, + ) + duration := time.Since(start).Milliseconds() + chatCompletionSpan.End() + if err != nil { + ChatCompletionSummaryData["error"] = true + // notice error with custom attributes + txn.NoticeError(newrelic.Error{ + Message: err.Error(), + Class: "OpenAIError", + Attributes: map[string]interface{}{ + "http.status": resp.Header().Get("Status"), + "error.code": resp.Header().Get("Error-Code"), + "completion_id": uuid.String(), + }, + }) + } + + // ratelimitLimitTokensUsageBased, ratelimitResetTokensUsageBased, and ratelimitRemainingTokensUsageBased are not in the response + // Request Headers + ChatCompletionSummaryData["request.temperature"] = req.Temperature + ChatCompletionSummaryData["request.max_tokens"] = req.MaxTokens + ChatCompletionSummaryData["request.model"] = req.Model + ChatCompletionSummaryData["model"] = req.Model + ChatCompletionSummaryData["duration"] = duration + + // Response Data + ChatCompletionSummaryData["response.number_of_messages"] = len(resp.Choices) + ChatCompletionSummaryData["response.model"] = resp.Model + ChatCompletionSummaryData["request_id"] = resp.ID + ChatCompletionSummaryData["response.organization"] = resp.Header().Get("Openai-Organization") + ChatCompletionSummaryData["response.number_of_messages"] = len(resp.Choices) + ChatCompletionSummaryData["response.usage.total_tokens"] = resp.Usage.TotalTokens + ChatCompletionSummaryData["response.usage.prompt_tokens"] = resp.Usage.PromptTokens + ChatCompletionSummaryData["response.usage.completion_tokens"] = resp.Usage.CompletionTokens + if len(resp.Choices) > 0 { + finishReason, err := resp.Choices[0].FinishReason.MarshalJSON() + if err != nil { + ChatCompletionSummaryData["error"] = true + txn.NoticeError(newrelic.Error{ + Message: err.Error(), + Class: "OpenAIError", + }) + } else { + ChatCompletionSummaryData["response.choices.finish_reason"] = string(finishReason) + } + } + + // Response Headers + ChatCompletionSummaryData["response.headers.llmVersion"] = resp.Header().Get("Openai-Version") + ChatCompletionSummaryData["response.headers.ratelimitLimitRequests"] = resp.Header().Get("X-Ratelimit-Limit-Requests") + ChatCompletionSummaryData["response.headers.ratelimitLimitTokens"] = resp.Header().Get("X-Ratelimit-Limit-Tokens") + ChatCompletionSummaryData["response.headers.ratelimitResetTokens"] = resp.Header().Get("X-Ratelimit-Reset-Tokens") + ChatCompletionSummaryData["response.headers.ratelimitResetRequests"] = resp.Header().Get("X-Ratelimit-Reset-Requests") + ChatCompletionSummaryData["response.headers.ratelimitRemainingTokens"] = resp.Header().Get("X-Ratelimit-Remaining-Tokens") + ChatCompletionSummaryData["response.headers.ratelimitRemainingRequests"] = resp.Header().Get("X-Ratelimit-Remaining-Requests") + + // New Relic Attributes + ChatCompletionSummaryData["id"] = uuid.String() + ChatCompletionSummaryData["span_id"] = spanID + ChatCompletionSummaryData["transaction_id"] = transactionID + ChatCompletionSummaryData["trace_id"] = traceID + ChatCompletionSummaryData["api_key_last_four_digits"] = cw.LicenseKeyLastFour + ChatCompletionSummaryData["vendor"] = "OpenAI" + ChatCompletionSummaryData["ingest_source"] = "Go" + ChatCompletionSummaryData["appName"] = appConfig.AppName + + // Record any custom attributes if they exist + ChatCompletionSummaryData = AppendCustomAttributesToEvent(cw, ChatCompletionSummaryData) + + // Record Custom Event + app.RecordCustomEvent("LlmChatCompletionSummary", ChatCompletionSummaryData) + + // Capture completion messages + NRCreateChatCompletionMessage(txn, app, resp, uuid, cw) + txn.End() + + return ChatCompletionResponseWrapper{ + ChatCompletionResponse: resp, + TraceID: traceID, + } +} + +// NRCreateChatCompletionMessage captures the completion messages and records a custom event in New Relic for each message +func NRCreateChatCompletionMessage(txn *newrelic.Transaction, app *newrelic.Application, resp openai.ChatCompletionResponse, uuid uuid.UUID, cw *ClientWrapper) { + spanID := txn.GetTraceMetadata().SpanID + traceID := txn.GetTraceMetadata().TraceID + transactionID := traceID[:16] + appCfg, configErr := app.Config() + if !configErr { + appCfg.AppName = "Unknown" + } + chatCompletionMessageSpan := txn.StartSegment("Llm/completion/OpenAI/CreateChatCompletionMessage") + for i, choice := range resp.Choices { + ChatCompletionMessageData := map[string]interface{}{} + // if the response doesn't have an ID, use the UUID from the summary + if resp.ID == "" { + ChatCompletionMessageData["id"] = uuid.String() + } else { + ChatCompletionMessageData["id"] = resp.ID + } + + // Response Data + ChatCompletionMessageData["response.model"] = resp.Model + + if appCfg.AIMonitoring.RecordContent.Enabled { + ChatCompletionMessageData["content"] = choice.Message.Content + } + + ChatCompletionMessageData["role"] = choice.Message.Role + + // Request Headers + ChatCompletionMessageData["request_id"] = resp.Header().Get("X-Request-Id") + + // New Relic Attributes + ChatCompletionMessageData["sequence"] = i + ChatCompletionMessageData["vendor"] = "openai" + ChatCompletionMessageData["ingest_source"] = "go" + ChatCompletionMessageData["span_id"] = spanID + ChatCompletionMessageData["trace_id"] = traceID + ChatCompletionMessageData["transaction_id"] = transactionID + // TO:DO completion_id set in CompletionSummary which is a UUID generated by the agent to identify the event + // TO:DO - llm.conversation_id + + // If custom attributes are set, add them to the data + ChatCompletionMessageData = AppendCustomAttributesToEvent(cw, ChatCompletionMessageData) + + // Record Custom Event for each message + app.RecordCustomEvent("LlmChatCompletionMessage", ChatCompletionMessageData) + + } + + chatCompletionMessageSpan.End() +} + +// 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) { + config, cfgErr := app.Config() + if !cfgErr { + config.AppName = "Unknown" + } + resp := ChatCompletionResponseWrapper{} + // If AI Monitoring is disabled, do not start a transaction but still perform the request + if !config.AIMonitoring.Enabled { + chatresp, err := cw.Client.CreateChatCompletion(context.Background(), req) + resp.ChatCompletionResponse = chatresp + if err != nil { + + return resp, err + } + return resp, errAIMonitoringDisabled + } + // Start NR Transaction + txn := app.StartTransaction("OpenAIChatCompletion") + resp = NRCreateChatCompletionSummary(txn, app, cw, req) + + return resp, nil +} + +// NRCreateEmbedding is a wrapper for the OpenAI CreateEmbedding method. +// If AI Monitoring is disabled, the wrapped function will still call the OpenAI CreateEmbedding method and return the response with no New Relic instrumentation +func NRCreateEmbedding(cw *ClientWrapper, req openai.EmbeddingRequest, app *newrelic.Application) (openai.EmbeddingResponse, error) { + config, cfgErr := app.Config() + if !cfgErr { + config.AppName = "Unknown" + } + + resp := openai.EmbeddingResponse{} + + // If AI Monitoring is disabled, do not start a transaction but still perform the request + if !config.AIMonitoring.Enabled { + resp, err := cw.Client.CreateEmbeddings(context.Background(), req) + if err != nil { + + return resp, err + } + return resp, errAIMonitoringDisabled + } + + // Start NR Transaction + txn := app.StartTransaction("OpenAIEmbedding") + + spanID := txn.GetTraceMetadata().SpanID + traceID := txn.GetTraceMetadata().TraceID + transactionID := traceID[:16] + EmbeddingsData := map[string]interface{}{} + uuid := uuid.New() + + embeddingSpan := txn.StartSegment("Llm/completion/OpenAI/CreateEmbedding") + start := time.Now() + resp, err := cw.Client.CreateEmbeddings(context.Background(), req) + duration := time.Since(start).Milliseconds() + embeddingSpan.End() + + if err != nil { + EmbeddingsData["error"] = true + txn.NoticeError(newrelic.Error{ + Message: err.Error(), + Class: "OpenAIError", + Attributes: map[string]interface{}{ + "http.status": resp.Header().Get("Status"), + "error.code": resp.Header().Get("Error-Code"), + "embedding_id": uuid.String(), + }, + }) + } + + // Request Data + if config.AIMonitoring.RecordContent.Enabled { + EmbeddingsData["input"] = GetInput(req.Input) + } + EmbeddingsData["api_key_last_four_digits"] = cw.LicenseKeyLastFour + EmbeddingsData["request.model"] = string(req.Model) + EmbeddingsData["duration"] = duration + + // Response Data + EmbeddingsData["response.model"] = string(resp.Model) + EmbeddingsData["response.usage.total_tokens"] = resp.Usage.TotalTokens + EmbeddingsData["response.usage.prompt_tokens"] = resp.Usage.PromptTokens + + // Response Headers + EmbeddingsData["response.organization"] = resp.Header().Get("Openai-Organization") + EmbeddingsData["response.headers.llmVersion"] = resp.Header().Get("Openai-Version") + EmbeddingsData["response.headers.ratelimitLimitRequests"] = resp.Header().Get("X-Ratelimit-Limit-Requests") + EmbeddingsData["response.headers.ratelimitLimitTokens"] = resp.Header().Get("X-Ratelimit-Limit-Tokens") + EmbeddingsData["response.headers.ratelimitResetTokens"] = resp.Header().Get("X-Ratelimit-Reset-Tokens") + EmbeddingsData["response.headers.ratelimitResetRequests"] = resp.Header().Get("X-Ratelimit-Reset-Requests") + EmbeddingsData["response.headers.ratelimitRemainingTokens"] = resp.Header().Get("X-Ratelimit-Remaining-Tokens") + EmbeddingsData["response.headers.ratelimitRemainingRequests"] = resp.Header().Get("X-Ratelimit-Remaining-Requests") + + EmbeddingsData = AppendCustomAttributesToEvent(cw, EmbeddingsData) + + // New Relic Attributes + EmbeddingsData["id"] = uuid.String() + EmbeddingsData["vendor"] = "OpenAI" + EmbeddingsData["ingest_source"] = "Go" + EmbeddingsData["span_id"] = spanID + EmbeddingsData["transaction_id"] = transactionID + EmbeddingsData["trace_id"] = traceID + + app.RecordCustomEvent("LlmEmbedding", EmbeddingsData) + txn.End() + return resp, nil +} + +func NRCreateChatCompletionStream(cw *ClientWrapper, ctx context.Context, req openai.ChatCompletionRequest, app *newrelic.Application) (*ChatCompletionStreamWrapper, error) { + config, cfgErr := app.Config() + if !cfgErr { + config.AppName = "Unknown" + } + + // If AI Monitoring OR AIMonitoring.Streaming is disabled, do not start a transaction but still perform the request + if !config.AIMonitoring.Enabled || !config.AIMonitoring.Streaming.Enabled { + stream, err := cw.Client.CreateChatCompletionStream(ctx, req) + if err != nil { + + return &ChatCompletionStreamWrapper{stream: stream}, err + } + return &ChatCompletionStreamWrapper{stream: stream}, errAIMonitoringDisabled + } + + txn := app.StartTransaction("OpenAIChatCompletionStream") + spanID := txn.GetTraceMetadata().SpanID + traceID := txn.GetTraceMetadata().TraceID + transactionID := traceID[:16] + StreamingData := map[string]interface{}{} + uuid := uuid.New() + + streamSpan := txn.StartSegment("Llm/completion/OpenAI/stream") + start := time.Now() + stream, err := cw.Client.CreateChatCompletionStream(ctx, req) + duration := time.Since(start).Milliseconds() + streamSpan.End() + + if err != nil { + StreamingData["error"] = true + txn.NoticeError(newrelic.Error{ + Message: err.Error(), + Class: "OpenAIError", + }) + txn.End() + return nil, err + } + + // Request Data + StreamingData["api_key_last_four_digits"] = cw.LicenseKeyLastFour + StreamingData["request.model"] = string(req.Model) + StreamingData["request.temperature"] = req.Temperature + StreamingData["request.max_tokens"] = req.MaxTokens + StreamingData["model"] = req.Model + + StreamingData["duration"] = duration + + // New Relic Attributes + StreamingData["id"] = uuid.String() + StreamingData["span_id"] = spanID + StreamingData["transaction_id"] = transactionID + StreamingData["trace_id"] = traceID + StreamingData["api_key_last_four_digits"] = cw.LicenseKeyLastFour + StreamingData["vendor"] = "OpenAI" + StreamingData["ingest_source"] = "Go" + StreamingData["appName"] = config.AppName + app.RecordCustomEvent("LlmChatCompletionSummary", StreamingData) + txn.End() + return &ChatCompletionStreamWrapper{stream: stream, txn: txn}, nil + +} diff --git a/v3/integrations/nropenai/nropenai_test.go b/v3/integrations/nropenai/nropenai_test.go new file mode 100644 index 000000000..6648de26f --- /dev/null +++ b/v3/integrations/nropenai/nropenai_test.go @@ -0,0 +1,684 @@ +package nropenai + +import ( + "context" + "errors" + "net/http" + "testing" + + "github.com/newrelic/go-agent/v3/internal" + "github.com/newrelic/go-agent/v3/internal/integrationsupport" + "github.com/newrelic/go-agent/v3/newrelic" + "github.com/sashabaranov/go-openai" +) + +type MockOpenAIClient struct { + MockCreateChatCompletionResp openai.ChatCompletionResponse + MockCreateEmbeddingsResp openai.EmbeddingResponse + MockCreateChatCompletionStream *openai.ChatCompletionStream + MockCreateChatCompletionErr error +} + +// Mock CreateChatCompletion function that returns a mock response +func (m *MockOpenAIClient) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (openai.ChatCompletionResponse, error) { + + MockResponse := openai.ChatCompletionResponse{ + ID: "chatcmpl-123", + Object: "chat.completion", + Created: 1677652288, + Model: openai.GPT3Dot5Turbo, + SystemFingerprint: "fp_44709d6fcb", + Usage: openai.Usage{ + PromptTokens: 9, + CompletionTokens: 12, + TotalTokens: 21, + }, + Choices: []openai.ChatCompletionChoice{ + { + Index: 0, + Message: openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleAssistant, + Content: "\n\nHello there, how may I assist you today?", + }, + }, + }, + } + hdrs := http.Header{} + hdrs.Add("X-Request-Id", "chatcmpl-123") + hdrs.Add("ratelimit-limit-tokens", "100") + hdrs.Add("Openai-Version", "2020-10-01") + hdrs.Add("X-Ratelimit-Limit-Requests", "10000") + hdrs.Add("X-Ratelimit-Limit-Tokens", "100") + hdrs.Add("X-Ratelimit-Reset-Tokens", "100") + hdrs.Add("X-Ratelimit-Reset-Requests", "10000") + hdrs.Add("X-Ratelimit-Remaining-Tokens", "100") + hdrs.Add("X-Ratelimit-Remaining-Requests", "10000") + hdrs.Add("Openai-Organization", "user-123") + + if req.Messages[0].Content == "testError" { + mockRespErr := openai.ChatCompletionResponse{} + hdrs.Add("Status", "404") + hdrs.Add("Error-Code", "404") + mockRespErr.SetHeader(hdrs) + return mockRespErr, errors.New("test error") + } + MockResponse.SetHeader(hdrs) + + return MockResponse, m.MockCreateChatCompletionErr +} + +func (m *MockOpenAIClient) CreateEmbeddings(ctx context.Context, conv openai.EmbeddingRequestConverter) (res openai.EmbeddingResponse, err error) { + MockResponse := openai.EmbeddingResponse{ + Model: openai.AdaEmbeddingV2, + Usage: openai.Usage{ + PromptTokens: 9, + CompletionTokens: 12, + TotalTokens: 21, + }, + Data: []openai.Embedding{ + { + Embedding: []float32{0.1, 0.2, 0.3}, + }, + }, + } + hdrs := http.Header{} + hdrs.Add("X-Request-Id", "chatcmpl-123") + hdrs.Add("ratelimit-limit-tokens", "100") + hdrs.Add("Openai-Version", "2020-10-01") + hdrs.Add("X-Ratelimit-Limit-Requests", "10000") + hdrs.Add("X-Ratelimit-Limit-Tokens", "100") + hdrs.Add("X-Ratelimit-Reset-Tokens", "100") + hdrs.Add("X-Ratelimit-Reset-Requests", "10000") + hdrs.Add("X-Ratelimit-Remaining-Tokens", "100") + hdrs.Add("X-Ratelimit-Remaining-Requests", "10000") + hdrs.Add("Openai-Organization", "user-123") + cv := conv.Convert() + if cv.Input == "testError" { + mockRespErr := openai.EmbeddingResponse{} + hdrs.Add("Status", "404") + hdrs.Add("Error-Code", "404") + mockRespErr.SetHeader(hdrs) + return mockRespErr, errors.New("test error") + } + + MockResponse.SetHeader(hdrs) + + return MockResponse, m.MockCreateChatCompletionErr +} + +func (m *MockOpenAIClient) CreateChatCompletionStream(ctx context.Context, request openai.ChatCompletionRequest) (stream *openai.ChatCompletionStream, err error) { + if request.Messages[0].Content == "testError" { + return m.MockCreateChatCompletionStream, errors.New("test error") + } + return m.MockCreateChatCompletionStream, m.MockCreateChatCompletionErr +} + +func TestFormatAPIKey(t *testing.T) { + dummyAPIKey := "sk-12345678900abcdefghijklmnop" + formattedKey := FormatAPIKey(dummyAPIKey) + if formattedKey != "sk-mnop" { + t.Errorf("Formatted API key is incorrect: expected: %s actual: %s", "sk-mnop", formattedKey) + + } +} +func TestDefaultConfig(t *testing.T) { + dummyAPIKey := "sk-12345678900abcdefghijklmnop" + cfg := NRDefaultConfig(dummyAPIKey) + // Default Values + if cfg.LicenseKeyLastFour != "sk-mnop" { + t.Errorf("API Key is incorrect: expected: %s actual: %s", "sk-mnop", cfg.LicenseKeyLastFour) + } + if cfg.Config.OrgID != "" { + t.Errorf("OrgID is incorrect: expected: %s actual: %s", "", cfg.Config.OrgID) + } + // Default Value set by openai package + if cfg.Config.APIType != openai.APITypeOpenAI { + t.Errorf("API Type is incorrect: expected: %s actual: %s", openai.APITypeOpenAI, cfg.Config.APIType) + } +} + +func TestDefaultConfigAzure(t *testing.T) { + dummyAPIKey := "sk-12345678900abcdefghijklmnop" + baseURL := "https://azure-base-url.com" + cfg := NRDefaultAzureConfig(dummyAPIKey, baseURL) + // Default Values + if cfg.LicenseKeyLastFour != "sk-mnop" { + t.Errorf("API Key is incorrect: expected: %s actual: %s", "sk-mnop", cfg.LicenseKeyLastFour) + } + if cfg.Config.BaseURL != baseURL { + t.Errorf("baseURL is incorrect: expected: %s actual: %s", baseURL, cfg.Config.BaseURL) + } + // Default Value set by openai package + if cfg.Config.APIType != openai.APITypeAzure { + t.Errorf("API Type is incorrect: expected: %s actual: %s", openai.APITypeAzure, cfg.Config.APIType) + } +} + +func TestNRNewClient(t *testing.T) { + dummyAPIKey := "sk-12345678900abcdefghijklmnop" + client := NRNewClient(dummyAPIKey) + if client.LicenseKeyLastFour != "sk-mnop" { + t.Errorf("API Key is incorrect: expected: %s actual: %s", "sk-mnop", client.LicenseKeyLastFour) + } +} + +func TestNRNewClientWithConfigs(t *testing.T) { + // Regular Config + dummyAPIKey := "sk-12345678900abcdefghijklmnop" + cfg := NRDefaultConfig(dummyAPIKey) + client := NRNewClientWithConfig(cfg) + if client.LicenseKeyLastFour != "sk-mnop" { + t.Errorf("API Key is incorrect: expected: %s actual: %s", "sk-mnop", client.LicenseKeyLastFour) + } + // Azure Config + baseURL := "https://azure-base-url.com" + azureCfg := NRDefaultAzureConfig(dummyAPIKey, baseURL) + azureClient := NRNewClientWithConfig(azureCfg) + if azureClient.LicenseKeyLastFour != "sk-mnop" { + t.Errorf("API Key is incorrect: expected: %s actual: %s", "sk-mnop", azureClient.LicenseKeyLastFour) + } + if azureCfg.Config.BaseURL != baseURL { + t.Errorf("baseURL is incorrect: expected: %s actual: %s", baseURL, azureCfg.Config.BaseURL) + } + // Default Value set by openai package + if azureCfg.Config.APIType != openai.APITypeAzure { + t.Errorf("API Type is incorrect: expected: %s actual: %s", openai.APITypeAzure, azureCfg.Config.APIType) + } +} + +func TestAddCustomAttributes(t *testing.T) { + client := NRNewClient("sk-12345678900abcdefghijklmnop") + client.AddCustomAttributes(map[string]interface{}{ + "llm.foo": "bar", + }) + if client.CustomAttributes["llm.foo"] != "bar" { + t.Errorf("Custom attribute is incorrect: expected: %s actual: %s", "bar", client.CustomAttributes["llm.foo"]) + } +} +func TestAddCustomAttributesIncorrectPrefix(t *testing.T) { + client := NRNewClient("sk-12345678900abcdefghijklmnop") + client.AddCustomAttributes(map[string]interface{}{ + "msdwmdoawd.foo": "bar", + }) + if len(client.CustomAttributes) != 0 { + t.Errorf("Custom attribute is incorrect: expected: %d actual: %d", 0, len(client.CustomAttributes)) + } +} + +func TestNRCreateChatCompletion(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0, + MaxTokens: 150, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "What is 8*5", + }, + }, + } + app := integrationsupport.NewTestApp(nil, newrelic.ConfigAIMonitoringEnabled(true)) + resp, err := NRCreateChatCompletion(cw, req, app.Application) + if err != nil { + t.Error(err) + } + if resp.ChatCompletionResponse.Choices[0].Message.Content != "\n\nHello there, how may I assist you today?" { + t.Errorf("Chat completion response is incorrect: expected: %s actual: %s", "\n\nHello there, how may I assist you today?", resp.ChatCompletionResponse.Choices[0].Message.Content) + } + app.ExpectCustomEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "LlmChatCompletionSummary", + "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, + "response.headers.ratelimitRemainingTokens": "100", + "response.headers.ratelimitRemainingRequests": "10000", + "response.headers.ratelimitResetTokens": "100", + "response.headers.ratelimitResetRequests": "10000", + "response.headers.ratelimitLimitTokens": "100", + "response.headers.ratelimitLimitRequests": "10000", + }, + }, + { + Intrinsics: map[string]interface{}{ + "type": "LlmChatCompletionMessage", + "timestamp": internal.MatchAnything, + }, + UserAttributes: map[string]interface{}{ + "trace_id": internal.MatchAnything, + "transaction_id": internal.MatchAnything, + "span_id": internal.MatchAnything, + "id": "chatcmpl-123", + "sequence": 0, + "role": "assistant", + "content": "\n\nHello there, how may I assist you today?", + "request_id": "chatcmpl-123", + "vendor": "openai", + "ingest_source": "go", + "response.model": "gpt-3.5-turbo", + }, + AgentAttributes: map[string]interface{}{}, + }, + }) + +} + +func TestNRCreateChatCompletionAIMonitoringNotEnabled(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0, + MaxTokens: 150, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "What is 8*5", + }, + }, + } + app := integrationsupport.NewTestApp(nil) + resp, err := NRCreateChatCompletion(cw, req, app.Application) + if err != errAIMonitoringDisabled { + t.Error(err) + } + // If AI Monitoring is disabled, no events should be sent, but a response from OpenAI should still be returned + if resp.ChatCompletionResponse.Choices[0].Message.Content != "\n\nHello there, how may I assist you today?" { + t.Errorf("Chat completion response is incorrect: expected: %s actual: %s", "\n\nHello there, how may I assist you today?", resp.ChatCompletionResponse.Choices[0].Message.Content) + } + app.ExpectCustomEvents(t, []internal.WantEvent{}) + +} + +func TestNRCreateChatCompletionError(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0, + MaxTokens: 150, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "testError", + }, + }, + } + app := integrationsupport.NewTestApp(nil, newrelic.ConfigAIMonitoringEnabled(true)) + _, err := NRCreateChatCompletion(cw, req, app.Application) + if err != nil { + t.Error(err) + } + app.ExpectCustomEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "LlmChatCompletionSummary", + "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, + "response.headers.ratelimitRemainingTokens": "100", + "response.headers.ratelimitRemainingRequests": "10000", + "response.headers.ratelimitResetTokens": "100", + "response.headers.ratelimitResetRequests": "10000", + "response.headers.ratelimitLimitTokens": "100", + "response.headers.ratelimitLimitRequests": "10000", + }, + }, + }) + app.ExpectErrorEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "TransactionError", + "transactionName": "OtherTransaction/Go/OpenAIChatCompletion", + "guid": internal.MatchAnything, + "priority": internal.MatchAnything, + "sampled": internal.MatchAnything, + "traceId": internal.MatchAnything, + "error.class": "OpenAIError", + "error.message": "test error", + }, + UserAttributes: map[string]interface{}{ + "error.code": "404", + "http.status": "404", + "completion_id": internal.MatchAnything, + }, + }}) +} +func TestNRCreateEmbedding(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + embeddingReq := openai.EmbeddingRequest{ + Input: []string{ + "The food was delicious and the waiter", + "Other examples of embedding request", + }, + Model: openai.AdaEmbeddingV2, + EncodingFormat: openai.EmbeddingEncodingFormatFloat, + } + + app := integrationsupport.NewTestApp(nil, newrelic.ConfigAIMonitoringEnabled(true)) + + _, err := NRCreateEmbedding(cw, embeddingReq, app.Application) + if err != nil { + t.Error(err) + } + app.ExpectCustomEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "LlmEmbedding", + "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", + "response.headers.ratelimitRemainingTokens": "100", + "response.headers.ratelimitRemainingRequests": "10000", + "response.headers.ratelimitResetTokens": "100", + "response.headers.ratelimitResetRequests": "10000", + "response.headers.ratelimitLimitTokens": "100", + "response.headers.ratelimitLimitRequests": "10000", + }, + }, + }) + +} + +func TestNRCreateEmbeddingAIMonitoringNotEnabled(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + embeddingReq := openai.EmbeddingRequest{ + Input: []string{ + "The food was delicious and the waiter", + "Other examples of embedding request", + }, + Model: openai.AdaEmbeddingV2, + EncodingFormat: openai.EmbeddingEncodingFormatFloat, + } + + app := integrationsupport.NewTestApp(nil) + + resp, err := NRCreateEmbedding(cw, embeddingReq, app.Application) + if err != errAIMonitoringDisabled { + t.Error(err) + } + // If AI Monitoring is disabled, no events should be sent, but a response from OpenAI should still be returned + app.ExpectCustomEvents(t, []internal.WantEvent{}) + if resp.Data[0].Embedding[0] != 0.1 { + t.Errorf("Embedding response is incorrect: expected: %f actual: %f", 0.1, resp.Data[0].Embedding[0]) + } + +} +func TestNRCreateEmbeddingError(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + embeddingReq := openai.EmbeddingRequest{ + Input: "testError", + Model: openai.AdaEmbeddingV2, + EncodingFormat: openai.EmbeddingEncodingFormatFloat, + } + + app := integrationsupport.NewTestApp(nil, newrelic.ConfigAIMonitoringEnabled(true)) + + _, err := NRCreateEmbedding(cw, embeddingReq, app.Application) + if err != nil { + t.Error(err) + } + + app.ExpectCustomEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "LlmEmbedding", + "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", + "response.headers.ratelimitRemainingTokens": "100", + "response.headers.ratelimitRemainingRequests": "10000", + "response.headers.ratelimitResetTokens": "100", + "response.headers.ratelimitResetRequests": "10000", + "response.headers.ratelimitLimitTokens": "100", + "response.headers.ratelimitLimitRequests": "10000", + }, + }, + }) + + app.ExpectErrorEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "TransactionError", + "transactionName": "OtherTransaction/Go/OpenAIEmbedding", + "guid": internal.MatchAnything, + "priority": internal.MatchAnything, + "sampled": internal.MatchAnything, + "traceId": internal.MatchAnything, + "error.class": "OpenAIError", + "error.message": "test error", + }, + UserAttributes: map[string]interface{}{ + "error.code": "404", + "http.status": "404", + "embedding_id": internal.MatchAnything, + }, + }}) +} + +func TestNRCreateStream(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0, + MaxTokens: 1500, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "Say this is a test", + }, + }, + Stream: true, + } + app := integrationsupport.NewTestApp(nil, newrelic.ConfigAIMonitoringEnabled(true)) + _, err := NRCreateChatCompletionStream(cw, context.Background(), req, app.Application) + if err != nil { + t.Error(err) + } + app.ExpectCustomEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "LlmChatCompletionSummary", + "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, + "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, + }, + }, + }) +} + +func TestNRCreateStreamAIMonitoringNotEnabled(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0, + MaxTokens: 1500, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "Say this is a test", + }, + }, + Stream: true, + } + app := integrationsupport.NewTestApp(nil) + _, err := NRCreateChatCompletionStream(cw, context.Background(), req, app.Application) + if err != errAIMonitoringDisabled { + t.Error(err) + } + app.ExpectCustomEvents(t, []internal.WantEvent{}) + app.ExpectTxnEvents(t, []internal.WantEvent{}) +} + +func TestNRCreateStreamError(t *testing.T) { + mockClient := &MockOpenAIClient{} + cw := &ClientWrapper{ + Client: mockClient, + LicenseKeyLastFour: "sk-mnop", + } + req := openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Temperature: 0, + MaxTokens: 1500, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleUser, + Content: "testError", + }, + }, + Stream: true, + } + app := integrationsupport.NewTestApp(nil, newrelic.ConfigAIMonitoringEnabled(true)) + _, err := NRCreateChatCompletionStream(cw, context.Background(), req, app.Application) + if err.Error() != "test error" { + t.Error(err) + } + + app.ExpectErrorEvents(t, []internal.WantEvent{ + { + Intrinsics: map[string]interface{}{ + "type": "TransactionError", + "transactionName": "OtherTransaction/Go/OpenAIChatCompletionStream", + "guid": internal.MatchAnything, + "priority": internal.MatchAnything, + "sampled": internal.MatchAnything, + "traceId": internal.MatchAnything, + "error.class": "OpenAIError", + "error.message": "test error", + }, + }}) + +} diff --git a/v3/newrelic/application.go b/v3/newrelic/application.go index 0249c738f..f7bdf0058 100644 --- a/v3/newrelic/application.go +++ b/v3/newrelic/application.go @@ -48,6 +48,32 @@ func (app *Application) RecordCustomEvent(eventType string, params map[string]in } } +// RecordLlmFeedbackEvent adds a LLM Feedback event. +// An error is logged if eventType or params is invalid. +func (app *Application) RecordLLMFeedbackEvent(trace_id string, rating any, category string, message string, metadata map[string]interface{}) { + if app == nil || app.app == nil { + return + } + CustomEventData := map[string]interface{}{ + "trace_id": trace_id, + "rating": rating, + "category": category, + "message": message, + "ingest_source": "Go", + } + for k, v := range metadata { + CustomEventData[k] = v + } + // if rating is an int or string, record the event + err := app.app.RecordCustomEvent("LlmFeedbackMessage", CustomEventData) + if err != nil { + app.app.Error("unable to record custom event", map[string]interface{}{ + "event-type": "LlmFeedbackMessage", + "reason": err.Error(), + }) + } +} + // RecordCustomMetric records a custom metric. The metric name you // provide will be prefixed by "Custom/". Custom metrics are not // currently supported in serverless mode. @@ -136,7 +162,6 @@ func (app *Application) Shutdown(timeout time.Duration) { // a boolean true value is returned as the second return value. If it is // false, then the Config data returned is the standard default configuration. // This usually occurs if the Application is not yet fully initialized. -// func (app *Application) Config() (Config, bool) { if app == nil || app.app == nil { return defaultConfig(), false diff --git a/v3/newrelic/config.go b/v3/newrelic/config.go index 5d79bf8fa..d0461ca1c 100644 --- a/v3/newrelic/config.go +++ b/v3/newrelic/config.go @@ -235,6 +235,17 @@ type Config struct { DynoNamePrefixesToShorten []string } + // AIMonitoring controls the behavior of AI monitoring features. + AIMonitoring struct { + Enabled bool + // Indicates whether streams will be instrumented + Streaming struct { + Enabled bool + } + RecordContent struct { + Enabled bool + } + } // CrossApplicationTracer controls behavior relating to cross application // tracing (CAT). In the case where CrossApplicationTracer and // DistributedTracer are both enabled, DistributedTracer takes precedence. @@ -667,6 +678,9 @@ func defaultConfig() Config { c.Heroku.UseDynoNames = true c.Heroku.DynoNamePrefixesToShorten = []string{"scheduler", "run"} + c.AIMonitoring.Enabled = false + c.AIMonitoring.Streaming.Enabled = true + c.AIMonitoring.RecordContent.Enabled = true c.InfiniteTracing.TraceObserver.Port = 443 c.InfiniteTracing.SpanEvents.QueueSize = 10000 diff --git a/v3/newrelic/config_options.go b/v3/newrelic/config_options.go index 5b9261e17..91ba00b59 100644 --- a/v3/newrelic/config_options.go +++ b/v3/newrelic/config_options.go @@ -236,6 +236,22 @@ func ConfigAppLogDecoratingEnabled(enabled bool) ConfigOption { } } +func ConfigAIMonitoringEnabled(enabled bool) ConfigOption { + return func(cfg *Config) { + if enabled && !cfg.HighSecurity { + cfg.AIMonitoring.Enabled = true + } else { + cfg.AIMonitoring.Enabled = false + } + } +} + +func ConfigAIMonitoringRecordContentEnabled(enabled bool) ConfigOption { + return func(cfg *Config) { + cfg.AIMonitoring.RecordContent.Enabled = enabled + } +} + // ConfigAppLogMetricsEnabled enables or disables the collection of metrics // data for logs seen by an instrumented logging framework // default: true diff --git a/v3/newrelic/config_test.go b/v3/newrelic/config_test.go index 37eb88159..9400a4e20 100644 --- a/v3/newrelic/config_test.go +++ b/v3/newrelic/config_test.go @@ -130,6 +130,15 @@ func TestCopyConfigReferenceFieldsPresent(t *testing.T) { "agent_version":"0.2.2", "host":"my-hostname", "settings":{ + "AIMonitoring": { + "Enabled": false, + "RecordContent": { + "Enabled": true + }, + "Streaming": { + "Enabled": true + } + }, "AppName":"my appname", "ApplicationLogging": { "Enabled": true, @@ -326,6 +335,15 @@ func TestCopyConfigReferenceFieldsAbsent(t *testing.T) { "agent_version":"0.2.2", "host":"my-hostname", "settings":{ + "AIMonitoring": { + "Enabled": false, + "RecordContent": { + "Enabled": true + }, + "Streaming": { + "Enabled": true + } + }, "AppName":"my appname", "ApplicationLogging": { "Enabled": true, diff --git a/v3/newrelic/internal_test.go b/v3/newrelic/internal_test.go index 994ce2c2b..a1deb7878 100644 --- a/v3/newrelic/internal_test.go +++ b/v3/newrelic/internal_test.go @@ -289,6 +289,26 @@ func testApp(replyfn func(*internal.ConnectReply), cfgfn func(*Config), t testin } } +func TestRecordLLMFeedbackEventSuccess(t *testing.T) { + app := testApp(nil, nil, t) + app.RecordLLMFeedbackEvent("traceid", "5", "informative", "message", validParams) + app.expectNoLoggedErrors(t) + app.ExpectCustomEvents(t, []internal.WantEvent{{ + Intrinsics: map[string]interface{}{ + "type": "LlmFeedbackMessage", + "timestamp": internal.MatchAnything, + }, + UserAttributes: map[string]interface{}{ + "trace_id": "traceid", + "rating": "5", + "category": "informative", + "message": "message", + "ingest_source": "Go", + "zip": 1, + "zap": 2, + }, + }}) +} func TestRecordCustomEventSuccess(t *testing.T) { app := testApp(nil, nil, t) app.RecordCustomEvent("myType", validParams)