From 2a3ef8fa23ef4223780e1c660dde0af15a80713b Mon Sep 17 00:00:00 2001 From: mirackara Date: Wed, 8 May 2024 12:07:28 -0500 Subject: [PATCH] OpenAI Comment Improvements --- v3/integrations/nropenai/nropenai.go | 61 +++++++++++++++++++++++++--- 1 file changed, 56 insertions(+), 5 deletions(-) diff --git a/v3/integrations/nropenai/nropenai.go b/v3/integrations/nropenai/nropenai.go index b858aa10b..5febb32b6 100644 --- a/v3/integrations/nropenai/nropenai.go +++ b/v3/integrations/nropenai/nropenai.go @@ -45,6 +45,7 @@ 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") ) +// OpenAIClient is any type that can invoke OpenAI model with a request. 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) @@ -157,7 +158,10 @@ func GetInput(any interface{}) any { } -// Wrapper for Recv() method that calls the underlying stream's Recv() method +// Wrapper for OpenAI Streaming Recv() method +// Captures the response messages as they are received in the wrapper +// Once the stream is closed, the Close() method is called and sends the captured +// data to New Relic func (w *ChatCompletionStreamWrapper) Recv() (openai.ChatCompletionStreamResponse, error) { response, err := w.stream.Recv() if err != nil { @@ -184,6 +188,7 @@ func (w *ChatCompletionStreamWrapper) Recv() (openai.ChatCompletionStreamRespons } +// Close the stream and send the event to New Relic func (w *ChatCompletionStreamWrapper) Close() { w.StreamingData["response.model"] = w.model NRCreateChatCompletionMessageStream(w.app, uuid.MustParse(w.uuid), w, w.cw, w.sequence) @@ -200,8 +205,11 @@ 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 +// NRCreateChatCompletionSummary captures the request data for a chat completion request +// A new segment is created for the chat completion request, and the response data is timed and captured +// Custom attributes are added to the event if they exist from client.AddCustomAttributes() +// After closing out the custom event for the chat completion summary, the function then calls +// NRCreateChatCompletionMessageInput/NRCreateChatCompletionMessage to capture the request messages func NRCreateChatCompletionSummary(txn *newrelic.Transaction, app *newrelic.Application, cw *ClientWrapper, req openai.ChatCompletionRequest) ChatCompletionResponseWrapper { // Start span txn.AddAttribute("llm", true) @@ -311,6 +319,9 @@ func NRCreateChatCompletionSummary(txn *newrelic.Transaction, app *newrelic.Appl } // Captures initial request messages and records a custom event in New Relic for each message +// similarly to NRCreateChatCompletionMessage, but only for the request messages +// Returns the sequence of the messages sent in the request +// which is used to calculate the sequence in the response messages func NRCreateChatCompletionMessageInput(txn *newrelic.Transaction, app *newrelic.Application, req openai.ChatCompletionRequest, inputuuid uuid.UUID, cw *ClientWrapper) int { sequence := 0 for i, message := range req.Messages { @@ -358,7 +369,13 @@ func NRCreateChatCompletionMessageInput(txn *newrelic.Transaction, app *newrelic } -// NRCreateChatCompletionMessage captures the completion response messages and records a custom event in New Relic for each message +// NRCreateChatCompletionMessage captures the completion response messages and records a custom event +// in New Relic for each message. The completion response messages are the responses from the model +// after the request messages have been sent and logged in NRCreateChatCompletionMessageInput. +// The sequence of the messages is calculated by logging each of the request messages first, then +// incrementing the sequence for each response message. +// The token count is calculated for each message and added to the custom event if the token count callback is set +// If not, no token count is added to the custom event func NRCreateChatCompletionMessage(txn *newrelic.Transaction, app *newrelic.Application, resp openai.ChatCompletionResponse, uuid uuid.UUID, cw *ClientWrapper, sequence int, req openai.ChatCompletionRequest) { spanID := txn.GetTraceMetadata().SpanID traceID := txn.GetTraceMetadata().TraceID @@ -412,6 +429,8 @@ func NRCreateChatCompletionMessage(txn *newrelic.Transaction, app *newrelic.Appl } } +// NRCreateChatCompletionMessageStream is identical to NRCreateChatCompletionMessage, but for streaming responses. +// Gets invoked only when the stream is closed func NRCreateChatCompletionMessageStream(app *newrelic.Application, uuid uuid.UUID, sw *ChatCompletionStreamWrapper, cw *ClientWrapper, sequence int) { spanID := sw.txn.GetTraceMetadata().SpanID @@ -476,6 +495,32 @@ func TokenCountingHelper(app *newrelic.Application, message openai.ChatCompletio return numTokens, (contentCounted && roleCounted) } +// Similar to NRCreateChatCompletionSummary, but for streaming responses +// Returns a custom wrapper with a stream that can be used to receive messages +// Example Usage: +/* + ctx := context.Background() + stream, err := nropenai.NRCreateChatCompletionStream(client, ctx, req, app) + if err != nil { + panic(err) + } + 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) + } + stream.Close() +*/ +// It is important to call stream.Close() after the stream has been used, as it will close the stream and send the event to New Relic. +// Additionally, custom attributes can be added to the client using client.AddCustomAttributes(map[string]interface{}) just like in NRCreateChatCompletionSummary func NRCreateChatCompletionStream(cw *ClientWrapper, ctx context.Context, req openai.ChatCompletionRequest, app *newrelic.Application) (*ChatCompletionStreamWrapper, error) { txn := app.StartTransaction("OpenAIChatCompletionStream") @@ -546,7 +591,13 @@ func NRCreateChatCompletionStream(cw *ClientWrapper, ctx context.Context, req op } // 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 +// If AI Monitoring is disabled, the wrapped function will still call the OpenAI CreateChatCompletion method +// and return the response with no New Relic instrumentation +// Calls NRCreateChatCompletionSummary to capture the request data and response data +// Returns a ChatCompletionResponseWrapper with the response and the TraceID of the transaction +// The trace ID is used to link the chat response with its feedback, with a call to SendFeedback() +// Otherwise, the response is the same as the OpenAI CreateChatCompletion method. It can be accessed +// by calling resp.ChatCompletionResponse func NRCreateChatCompletion(cw *ClientWrapper, req openai.ChatCompletionRequest, app *newrelic.Application) (ChatCompletionResponseWrapper, error) { config, _ := app.Config()