OpenAI Integration (newrelic#860)

* OpenAI integration
nr-swilloughby · Mar 19, 2024 · 4da4344 · 4da4344
1 parent 67fecce
commit 4da4344
Show file tree

Hide file tree

Showing 13 changed files with 1,565 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 => ../..