Added Configuration Setting 'AIMonitoring'

v3/integrations/nropenai/examples/chatcompletion/chatcompletion_example.go

@@ -16,6 +16,9 @@ func main() {
 		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)
@@ -52,8 +55,13 @@ func main() {
 		},
 	}
 	// Create Chat Completion
-	resp := nropenai.NRCreateChatCompletion(client, req, app)
-	fmt.Println(resp.ChatCompletionResponse.Choices[0].Message.Content)
+	resp, err := nropenai.NRCreateChatCompletion(client, req, app)
+
+	if err != nil {
+		fmt.Println("Unable to create chat completion: ", err)
+	} else {
+		fmt.Println(resp.ChatCompletionResponse.Choices[0].Message.Content)
+	}
 
 	// Shutdown Application
 	app.Shutdown(5 * time.Second)

v3/integrations/nropenai/examples/chatcompletionfeedback/chatcompletionfeedback.go

@@ -31,6 +31,7 @@ func main() {
 		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)
@@ -53,11 +54,16 @@ func main() {
 		},
 	}
 	// Create Chat Completion
-	resp := nropenai.NRCreateChatCompletion(client, req, app)
+	resp, err := nropenai.NRCreateChatCompletion(client, req, app)
+
+	if err != nil {
+		fmt.Println("Unable to create chat completion: ", err)
+	} else {
+		// Print the contents of the message
+		fmt.Println("Message Response: ", resp.ChatCompletionResponse.Choices[0].Message.Content)
+		SendFeedback(app, resp)
+	}
 
-	// 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)
 }

v3/integrations/nropenai/examples/embeddings/embeddings_example.go

@@ -16,6 +16,8 @@ func main() {
 		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)
@@ -45,7 +47,13 @@ func main() {
 		Model:          openai.AdaEmbeddingV2,
 		EncodingFormat: openai.EmbeddingEncodingFormatFloat,
 	}
-	nropenai.NRCreateEmbedding(client, embeddingReq, app)
+	resp, err := nropenai.NRCreateEmbedding(client, embeddingReq, app)
+	if err != nil {
+		fmt.Println("Unable to create embedding: ", err)
+	} else {
+		fmt.Println("Embedding Created!")
+		fmt.Println(resp.Usage.PromptTokens)
+	}
 	// Shutdown Application
 	app.Shutdown(5 * time.Second)
 }

v3/integrations/nropenai/nropenai.go

@@ -5,13 +5,18 @@ package nropenai
 
 import (
 	"context"
+	"errors"
 	"reflect"
 
 	"github.com/google/uuid"
 	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/sashabaranov/go-openai"
 )
 
+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.")
+)
+
 // Wrapper for OpenAI Configuration
 type ConfigWrapper struct {
 	Config             *openai.ClientConfig
@@ -222,12 +227,18 @@ func NRCreateChatCompletionMessage(txn *newrelic.Transaction, app *newrelic.Appl
 	chatCompletionMessageSpan.End()
 }
 
-func NRCreateChatCompletion(cw *ClientWrapper, req openai.ChatCompletionRequest, app *newrelic.Application) ChatCompletionResponseWrapper {
+func NRCreateChatCompletion(cw *ClientWrapper, req openai.ChatCompletionRequest, app *newrelic.Application) (ChatCompletionResponseWrapper, error) {
+	config, _ := app.Config()
+	resp := ChatCompletionResponseWrapper{}
+	// If AI Monitoring is disabled, do not start a transaction
+	if !config.AIMonitoring.Enabled {
+		return resp, errAIMonitoringDisabled
+	}
 	// Start NR Transaction
 	txn := app.StartTransaction("OpenAIChatCompletion")
-	resp := NRCreateChatCompletionSummary(txn, app, cw, req)
+	resp = NRCreateChatCompletionSummary(txn, app, cw, req)
 
-	return resp
+	return resp, nil
 }
 
 // If multiple messages are sent, only the first message is used
@@ -244,16 +255,24 @@ func GetInput(any interface{}) any {
 	return any
 
 }
-func NRCreateEmbedding(cw *ClientWrapper, req openai.EmbeddingRequest, app *newrelic.Application) openai.EmbeddingResponse {
-	EmbeddingsData := map[string]interface{}{}
-	uuid := uuid.New()
+func NRCreateEmbedding(cw *ClientWrapper, req openai.EmbeddingRequest, app *newrelic.Application) (openai.EmbeddingResponse, error) {
+	config, _ := app.Config()
+
+	resp := openai.EmbeddingResponse{}
+
+	// If AI Monitoring is disabled, do not start a transaction
+	if !config.AIMonitoring.Enabled {
+		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")
 	resp, err := cw.Client.CreateEmbeddings(context.Background(), req)
@@ -304,5 +323,5 @@ func NRCreateEmbedding(cw *ClientWrapper, req openai.EmbeddingRequest, app *newr
 
 	app.RecordCustomEvent("LlmEmbedding", EmbeddingsData)
 	txn.End()
-	return resp
+	return resp, nil
 }

v3/newrelic/config.go

@@ -235,6 +235,14 @@ 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
+		}
+	}
 	// CrossApplicationTracer controls behavior relating to cross application
 	// tracing (CAT).  In the case where CrossApplicationTracer and
 	// DistributedTracer are both enabled, DistributedTracer takes precedence.
@@ -667,6 +675,9 @@ func defaultConfig() Config {
 	c.Heroku.UseDynoNames = true
 	c.Heroku.DynoNamePrefixesToShorten = []string{"scheduler", "run"}
 
+	c.AIMonitoring.Enabled = false
+	c.AIMonitoring.Streaming.Enabled = true
+
 	c.InfiniteTracing.TraceObserver.Port = 443
 	c.InfiniteTracing.SpanEvents.QueueSize = 10000
 

v3/newrelic/config_options.go

@@ -236,6 +236,16 @@ 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
+		}
+	}
+}
+
 // ConfigAppLogMetricsEnabled enables or disables the collection of metrics
 // data for logs seen by an instrumented logging framework
 // default: true

v3/newrelic/config_test.go

@@ -130,6 +130,12 @@ func TestCopyConfigReferenceFieldsPresent(t *testing.T) {
 		"agent_version":"0.2.2",
 		"host":"my-hostname",
 		"settings":{
+			"AIMonitoring": {
+				"Enabled": false,
+				"Streaming": {
+					"Enabled": true
+				}
+			},
 			"AppName":"my appname",
 			"ApplicationLogging": {
 				"Enabled": true,
@@ -326,6 +332,12 @@ func TestCopyConfigReferenceFieldsAbsent(t *testing.T) {
 		"agent_version":"0.2.2",
 		"host":"my-hostname",
 		"settings":{
+			"AIMonitoring": {
+				"Enabled": false,
+				"Streaming": {
+					"Enabled": true
+				}
+			},
 			"AppName":"my appname",
 			"ApplicationLogging": {
 				"Enabled": true,