Skip to content

Commit cdc2f6c

Browse files
tippmar-nrtippmar-nr
authored
feat: Add AI Monitoring instrumentation for AmazonBedrockRuntimeClient.ConverseAsync() (#2997)
1 parent 9994259 commit cdc2f6c

File tree

6 files changed

+504
-12
lines changed

6 files changed

+504
-12
lines changed

‎src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Bedrock/ConverseAsyncWrapper.cs

+253
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,253 @@
1+
// Copyright 2020 New Relic, Inc. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using System;
5+
using System.Collections.Concurrent;
6+
using System.Net;
7+
using System.Threading.Tasks;
8+
using NewRelic.Agent.Api;
9+
using NewRelic.Agent.Extensions.Helpers;
10+
using NewRelic.Agent.Extensions.Llm;
11+
using NewRelic.Agent.Extensions.Providers.Wrapper;
12+
using NewRelic.Reflection;
13+
14+
namespace NewRelic.Providers.Wrapper.Bedrock;
15+
16+
public class ConverseAsyncWrapper : IWrapper
17+
{
18+
public bool IsTransactionRequired => true; // part of spec, only create events for transactions.
19+
20+
private static ConcurrentDictionary<Type, Func<object, object>> _getResultFromGenericTask = new();
21+
private static ConcurrentDictionary<string, string> _libraryVersions = new();
22+
private const string WrapperName = "BedrockConverseAsync";
23+
private const string VendorName = "Bedrock";
24+
25+
public CanWrapResponse CanWrap(InstrumentedMethodInfo methodInfo)
26+
{
27+
return new CanWrapResponse(WrapperName.Equals(methodInfo.RequestedWrapperName));
28+
}
29+
30+
public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
31+
{
32+
// Don't do anything, including sending the version Supportability metric, if we're disabled
33+
if (!agent.Configuration.AiMonitoringEnabled)
34+
{
35+
return Delegates.NoOp;
36+
}
37+
38+
if (instrumentedMethodCall.IsAsync)
39+
{
40+
transaction.AttachToAsync();
41+
}
42+
43+
dynamic converseRequest = instrumentedMethodCall.MethodCall.MethodArguments[0];
44+
string modelId = converseRequest.ModelId.ToLower();
45+
46+
var operationType = "completion"; // Converse doesn't support embedding
47+
var segment = transaction.StartCustomSegment(instrumentedMethodCall.MethodCall, $"Llm/{operationType}/{VendorName}/{instrumentedMethodCall.MethodCall.Method.MethodName}");
48+
49+
// required per spec
50+
var version = GetOrAddLibraryVersion(instrumentedMethodCall.MethodCall.Method.Type.Assembly.ManifestModule.Assembly.FullName);
51+
agent.RecordSupportabilityMetric($"DotNet/ML/{VendorName}/{version}");
52+
53+
return Delegates.GetAsyncDelegateFor<Task>(
54+
agent,
55+
segment,
56+
false,
57+
TryProcessConverseResponse,
58+
TaskContinuationOptions.ExecuteSynchronously
59+
);
60+
61+
void TryProcessConverseResponse(Task responseTask)
62+
{
63+
// We need the duration, so we end the segment before creating the events.
64+
segment.End();
65+
66+
if (responseTask.IsFaulted)
67+
{
68+
HandleError(segment, converseRequest, responseTask, agent, modelId);
69+
return;
70+
}
71+
72+
dynamic converseResponse = GetTaskResult(responseTask);
73+
if (converseResponse == null || converseResponse.HttpStatusCode >= HttpStatusCode.MultipleChoices)
74+
{
75+
agent.Logger.Warn($"Error processing Converse response for model {modelId}: Response {(converseResponse == null ? "is null" : $"has non-success HttpStatusCode: {converseResponse.HttpStatusCode}")}");
76+
return;
77+
}
78+
79+
ProcessConverseResponse(segment, converseRequest, converseResponse, agent, modelId);
80+
}
81+
}
82+
83+
private void ProcessConverseResponse(ISegment segment, dynamic converseRequest, dynamic converseResponse, IAgent agent, string requestModelId)
84+
{
85+
// if request message content doesn't have a non-null Text property, we can't support instrumentation
86+
// last message is the current prompt
87+
var requestMessage = converseRequest?.Messages?[converseRequest.Messages.Count - 1];
88+
if (converseRequest == null || requestMessage == null || requestMessage.Content == null || requestMessage.Content.Count == 0 || requestMessage.Content[0].Text == null)
89+
{
90+
agent.Logger.Info($"Unable to process Converse response for model {requestModelId}: request was null or message content was not Text");
91+
return;
92+
}
93+
94+
if (converseResponse == null)
95+
{
96+
agent.Logger.Warn($"Error processing Converse response for model {requestModelId}: response was null");
97+
return;
98+
}
99+
100+
// if response message content doesn't have a non-null Text property, we can't support instrumentation
101+
var responseMessage = converseResponse.Output?.Message;
102+
if (responseMessage == null || responseMessage.Content == null || responseMessage.Content.Count == 0 || responseMessage.Content[0].Text == null)
103+
{
104+
agent.Logger.Info($"Unable to process Converse response for model {requestModelId}: response was null or message content was not Text");
105+
return;
106+
}
107+
108+
string requestRole = requestMessage.Role?.Value ?? "user";
109+
string promptText = requestMessage.Content?[0]?.Text ?? "";
110+
111+
string responseRole = responseMessage.Role?.Value ?? "assistant";
112+
string responseText = responseMessage.Content?[0]?.Text ?? "";
113+
string stopReason = converseResponse.StopReason?.Value;
114+
115+
string requestId = converseResponse.ResponseMetadata?.RequestId;
116+
int? requestMaxTokens = converseRequest.InferenceConfig?.MaxTokens;
117+
float? requestTemperature = converseRequest.InferenceConfig?.Temperature;
118+
119+
int? inputTokens = converseResponse.Usage?.InputTokens;
120+
int? outputTokens = converseResponse.Usage?.OutputTokens;
121+
122+
var completionId = EventHelper.CreateChatCompletionEvent(
123+
agent,
124+
segment,
125+
requestId,
126+
requestTemperature,
127+
requestMaxTokens,
128+
requestModelId,
129+
requestModelId,
130+
2, // one request, one response
131+
stopReason,
132+
VendorName,
133+
false,
134+
null, // not available in AWS
135+
null
136+
);
137+
138+
// Prompt
139+
EventHelper.CreateChatMessageEvent(
140+
agent,
141+
segment,
142+
requestId,
143+
null,
144+
requestModelId,
145+
promptText,
146+
requestRole,
147+
0,
148+
completionId,
149+
false,
150+
VendorName,
151+
inputTokens);
152+
153+
// Response
154+
EventHelper.CreateChatMessageEvent(
155+
agent,
156+
segment,
157+
requestId,
158+
null,
159+
requestModelId,
160+
responseText,
161+
responseRole,
162+
1,
163+
completionId,
164+
true,
165+
VendorName,
166+
outputTokens);
167+
}
168+
169+
private void HandleError(ISegment segment, dynamic converseRequest, Task responseTask, IAgent agent, string modelId)
170+
{
171+
agent.Logger.Info($"Error processing Converse response for model {modelId}: {responseTask.Exception!.Message}");
172+
173+
dynamic bedrockException = responseTask.Exception!.InnerException;
174+
if (bedrockException == null)
175+
{
176+
agent.Logger.Warn($"Error processing Converse response for model {modelId}: Task faulted but there was no inner exception");
177+
return;
178+
}
179+
180+
var requestMessage = converseRequest?.Messages?[converseRequest.Messages.Count - 1];
181+
182+
if (converseRequest == null || requestMessage == null)
183+
{
184+
agent.Logger.Warn($"Error processing Converse response for model {modelId}: request Message was null");
185+
return;
186+
}
187+
188+
HttpStatusCode statusCode = bedrockException.StatusCode;
189+
string errorCode = bedrockException.ErrorCode;
190+
string errorMessage = bedrockException.Message;
191+
string requestId = bedrockException.RequestId;
192+
193+
var errorData = new LlmErrorData
194+
{
195+
HttpStatusCode = ((int)statusCode).ToString(),
196+
ErrorCode = errorCode,
197+
ErrorParam = null, // not available in AWS
198+
ErrorMessage = errorMessage
199+
};
200+
201+
string requestRole = requestMessage.Role?.Value ?? "user";
202+
string promptText = requestMessage.Content?[0]?.Text ?? "";
203+
int? requestMaxTokens = converseRequest.InferenceConfig?.MaxTokens;
204+
float? requestTemperature = converseRequest.InferenceConfig?.Temperature;
205+
206+
207+
var completionId = EventHelper.CreateChatCompletionEvent(
208+
agent,
209+
segment,
210+
requestId,
211+
requestTemperature,
212+
requestMaxTokens,
213+
converseRequest.ModelId,
214+
null,
215+
0,
216+
null,
217+
VendorName,
218+
true,
219+
null,
220+
errorData);
221+
222+
// Prompt
223+
EventHelper.CreateChatMessageEvent(
224+
agent,
225+
segment,
226+
requestId,
227+
null,
228+
converseRequest.ModelId,
229+
promptText,
230+
requestRole,
231+
0,
232+
completionId,
233+
false,
234+
VendorName);
235+
}
236+
237+
238+
private string GetOrAddLibraryVersion(string assemblyFullName)
239+
{
240+
return _libraryVersions.GetOrAdd(assemblyFullName, VersionHelpers.GetLibraryVersion(assemblyFullName));
241+
}
242+
243+
private static object GetTaskResult(object task)
244+
{
245+
if (((Task)task).IsFaulted)
246+
{
247+
return null;
248+
}
249+
250+
var getResponse = _getResultFromGenericTask.GetOrAdd(task.GetType(), t => VisibilityBypasser.Instance.GeneratePropertyAccessor<object>(t, "Result"));
251+
return getResponse(task);
252+
}
253+
}

‎src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Bedrock/Instrumentation.xml

+9-1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,14 @@ SPDX-License-Identifier: Apache-2.0
1313
<exactMethodMatcher methodName="InvokeModelAsync" />
1414
</match>
1515
</tracerFactory>
16-
16+
17+
<tracerFactory name="BedrockConverseAsync">
18+
<!--
19+
public virtual Task<ConverseResponse> ConverseAsync(ConverseRequest request, System.Threading.CancellationToken cancellationToken = default(CancellationToken))
20+
-->
21+
<match assemblyName="AWSSDK.BedrockRuntime" className="Amazon.BedrockRuntime.AmazonBedrockRuntimeClient">
22+
<exactMethodMatcher methodName="ConverseAsync" />
23+
</match>
24+
</tracerFactory>
1725
</instrumentation>
1826
</extension>

0 commit comments

Comments
 (0)