Skip to content

Commit 73c8137

Browse files
ezhang6811AsakerMohdvastin
authored
Implementation and contract tests for GenAi attribute support (#137)
Co-authored-by: Mohamed Asaker <[email protected]> Co-authored-by: Vastin <[email protected]>
1 parent bb29ccf commit 73c8137

File tree

7 files changed

+820
-34
lines changed

7 files changed

+820
-34
lines changed

‎src/OpenTelemetry.Instrumentation.AWS/Implementation/AWSLlmModelProcessor.cs‎

Lines changed: 362 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,362 @@
1+
// Copyright The OpenTelemetry Authors
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using System.Diagnostics;
5+
using System.Text;
6+
using System.Text.Json;
7+
using Amazon.Runtime;
8+
9+
namespace OpenTelemetry.Instrumentation.AWS.Implementation;
10+
11+
internal class AWSLlmModelProcessor
12+
{
13+
internal static void ProcessGenAiAttributes<T>(Activity activity, T message, string modelName, bool isRequest)
14+
{
15+
// message can be either a request or a response. isRequest is used by the model-specific methods to determine
16+
// whether to extract the request or response attributes.
17+
18+
// Currently, the .NET SDK does not expose "X-Amzn-Bedrock-*" HTTP headers in the response metadata, as per
19+
// https://github.com/aws/aws-sdk-net/issues/3171. As a result, we can only extract attributes given what is in
20+
// the response body. For the Claude, Command, and Mistral models, the input and output tokens are not provided
21+
// in the response body, so we approximate their values by dividing the input and output lengths by 6, based on
22+
// the Bedrock documentation here: https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
23+
24+
var messageBodyProperty = message?.GetType()?.GetProperty("Body");
25+
if (messageBodyProperty != null)
26+
{
27+
var body = messageBodyProperty.GetValue(message) as MemoryStream;
28+
if (body != null)
29+
{
30+
try
31+
{
32+
var jsonString = Encoding.UTF8.GetString(body.ToArray());
33+
var jsonObject = JsonSerializer.Deserialize<Dictionary<string, JsonElement>>(jsonString);
34+
if (jsonObject == null)
35+
{
36+
return;
37+
}
38+
39+
// extract model specific attributes based on model name
40+
if (modelName.Contains("amazon.titan"))
41+
{
42+
ProcessTitanModelAttributes(activity, jsonObject, isRequest);
43+
}
44+
else if (modelName.Contains("anthropic.claude"))
45+
{
46+
ProcessClaudeModelAttributes(activity, jsonObject, isRequest);
47+
}
48+
else if (modelName.Contains("meta.llama3"))
49+
{
50+
ProcessLlamaModelAttributes(activity, jsonObject, isRequest);
51+
}
52+
else if (modelName.Contains("cohere.command"))
53+
{
54+
ProcessCommandModelAttributes(activity, jsonObject, isRequest);
55+
}
56+
else if (modelName.Contains("ai21.jamba"))
57+
{
58+
ProcessJambaModelAttributes(activity, jsonObject, isRequest);
59+
}
60+
else if (modelName.Contains("mistral.mistral"))
61+
{
62+
ProcessMistralModelAttributes(activity, jsonObject, isRequest);
63+
}
64+
}
65+
catch (Exception ex)
66+
{
67+
Console.WriteLine("Exception: " + ex.Message);
68+
}
69+
}
70+
}
71+
}
72+
73+
private static void ProcessTitanModelAttributes(Activity activity, Dictionary<string, JsonElement> jsonBody, bool isRequest)
74+
{
75+
try
76+
{
77+
if (isRequest)
78+
{
79+
if (jsonBody.TryGetValue("textGenerationConfig", out var textGenerationConfig))
80+
{
81+
if (textGenerationConfig.TryGetProperty("topP", out var topP))
82+
{
83+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTopP, topP.GetDouble());
84+
}
85+
86+
if (textGenerationConfig.TryGetProperty("temperature", out var temperature))
87+
{
88+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTemperature, temperature.GetDouble());
89+
}
90+
91+
if (textGenerationConfig.TryGetProperty("maxTokenCount", out var maxTokens))
92+
{
93+
activity.SetTag(AWSSemanticConventions.AttributeGenAiMaxTokens, maxTokens.GetInt32());
94+
}
95+
}
96+
}
97+
else
98+
{
99+
if (jsonBody.TryGetValue("inputTextTokenCount", out var inputTokens))
100+
{
101+
activity.SetTag(AWSSemanticConventions.AttributeGenAiInputTokens, inputTokens.GetInt32());
102+
}
103+
104+
if (jsonBody.TryGetValue("results", out var resultsArray))
105+
{
106+
var results = resultsArray[0];
107+
if (results.TryGetProperty("tokenCount", out var outputTokens))
108+
{
109+
activity.SetTag(AWSSemanticConventions.AttributeGenAiOutputTokens, outputTokens.GetInt32());
110+
}
111+
112+
if (results.TryGetProperty("completionReason", out var finishReasons))
113+
{
114+
activity.SetTag(AWSSemanticConventions.AttributeGenAiFinishReasons, new string[] { finishReasons.GetString() ?? string.Empty });
115+
}
116+
}
117+
}
118+
}
119+
catch (Exception ex)
120+
{
121+
Console.WriteLine("Exception: " + ex.Message);
122+
}
123+
}
124+
125+
private static void ProcessClaudeModelAttributes(Activity activity, Dictionary<string, JsonElement> jsonBody, bool isRequest)
126+
{
127+
try
128+
{
129+
if (isRequest)
130+
{
131+
if (jsonBody.TryGetValue("top_p", out var topP))
132+
{
133+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTopP, topP.GetDouble());
134+
}
135+
136+
if (jsonBody.TryGetValue("temperature", out var temperature))
137+
{
138+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTemperature, temperature.GetDouble());
139+
}
140+
141+
if (jsonBody.TryGetValue("max_tokens", out var maxTokens))
142+
{
143+
activity.SetTag(AWSSemanticConventions.AttributeGenAiMaxTokens, maxTokens.GetInt32());
144+
}
145+
}
146+
else
147+
{
148+
if (jsonBody.TryGetValue("usage", out var usage))
149+
{
150+
if (usage.TryGetProperty("input_tokens", out var inputTokens))
151+
{
152+
activity.SetTag(AWSSemanticConventions.AttributeGenAiInputTokens, inputTokens.GetInt32());
153+
}
154+
if (usage.TryGetProperty("output_tokens", out var outputTokens))
155+
{
156+
activity.SetTag(AWSSemanticConventions.AttributeGenAiOutputTokens, outputTokens.GetInt32());
157+
}
158+
}
159+
if (jsonBody.TryGetValue("stop_reason", out var finishReasons))
160+
{
161+
activity.SetTag(AWSSemanticConventions.AttributeGenAiFinishReasons, new string[] { finishReasons.GetString() ?? string.Empty });
162+
}
163+
}
164+
}
165+
catch (Exception ex)
166+
{
167+
Console.WriteLine("Exception: " + ex.Message);
168+
}
169+
}
170+
171+
private static void ProcessLlamaModelAttributes(Activity activity, Dictionary<string, JsonElement> jsonBody, bool isRequest)
172+
{
173+
try
174+
{
175+
if (isRequest)
176+
{
177+
if (jsonBody.TryGetValue("top_p", out var topP))
178+
{
179+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTopP, topP.GetDouble());
180+
}
181+
182+
if (jsonBody.TryGetValue("temperature", out var temperature))
183+
{
184+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTemperature, temperature.GetDouble());
185+
}
186+
187+
if (jsonBody.TryGetValue("max_gen_len", out var maxTokens))
188+
{
189+
activity.SetTag(AWSSemanticConventions.AttributeGenAiMaxTokens, maxTokens.GetInt32());
190+
}
191+
}
192+
else
193+
{
194+
if (jsonBody.TryGetValue("prompt_token_count", out var inputTokens))
195+
{
196+
activity.SetTag(AWSSemanticConventions.AttributeGenAiInputTokens, inputTokens.GetInt32());
197+
}
198+
199+
if (jsonBody.TryGetValue("generation_token_count", out var outputTokens))
200+
{
201+
activity.SetTag(AWSSemanticConventions.AttributeGenAiOutputTokens, outputTokens.GetInt32());
202+
}
203+
204+
if (jsonBody.TryGetValue("stop_reason", out var finishReasons))
205+
{
206+
activity.SetTag(AWSSemanticConventions.AttributeGenAiFinishReasons, new string[] { finishReasons.GetString() ?? string.Empty });
207+
}
208+
}
209+
}
210+
catch (Exception ex)
211+
{
212+
Console.WriteLine("Exception: " + ex.Message);
213+
}
214+
}
215+
216+
private static void ProcessCommandModelAttributes(Activity activity, Dictionary<string, JsonElement> jsonBody, bool isRequest)
217+
{
218+
try
219+
{
220+
if (isRequest)
221+
{
222+
if (jsonBody.TryGetValue("p", out var topP))
223+
{
224+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTopP, topP.GetDouble());
225+
}
226+
227+
if (jsonBody.TryGetValue("temperature", out var temperature))
228+
{
229+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTemperature, temperature.GetDouble());
230+
}
231+
232+
if (jsonBody.TryGetValue("max_tokens", out var maxTokens))
233+
{
234+
activity.SetTag(AWSSemanticConventions.AttributeGenAiMaxTokens, maxTokens.GetInt32());
235+
}
236+
237+
// input tokens not provided in Command response body, so we estimate the value based on input length
238+
if (jsonBody.TryGetValue("message", out var input))
239+
{
240+
activity.SetTag(AWSSemanticConventions.AttributeGenAiInputTokens, Convert.ToInt32(Math.Ceiling((double) (input.GetString()?.Length ?? 0) / 6)));
241+
}
242+
}
243+
else
244+
{
245+
if (jsonBody.TryGetValue("finish_reason", out var finishReasons))
246+
{
247+
activity.SetTag(AWSSemanticConventions.AttributeGenAiFinishReasons, new string[] { finishReasons.GetString() ?? string.Empty });
248+
}
249+
250+
// completion tokens not provided in Command response body, so we estimate the value based on output length
251+
if (jsonBody.TryGetValue("text", out var output))
252+
{
253+
activity.SetTag(AWSSemanticConventions.AttributeGenAiOutputTokens, Convert.ToInt32(Math.Ceiling((double) (output.GetString()?.Length ?? 0) / 6)));
254+
}
255+
}
256+
}
257+
catch (Exception ex)
258+
{
259+
Console.WriteLine("Exception: " + ex.Message);
260+
}
261+
}
262+
263+
private static void ProcessJambaModelAttributes(Activity activity, Dictionary<string, JsonElement> jsonBody, bool isRequest)
264+
{
265+
try
266+
{
267+
if (isRequest)
268+
{
269+
if (jsonBody.TryGetValue("top_p", out var topP))
270+
{
271+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTopP, topP.GetDouble());
272+
}
273+
274+
if (jsonBody.TryGetValue("temperature", out var temperature))
275+
{
276+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTemperature, temperature.GetDouble());
277+
}
278+
279+
if (jsonBody.TryGetValue("max_tokens", out var maxTokens))
280+
{
281+
activity.SetTag(AWSSemanticConventions.AttributeGenAiMaxTokens, maxTokens.GetInt32());
282+
}
283+
}
284+
else
285+
{
286+
if (jsonBody.TryGetValue("usage", out var usage))
287+
{
288+
if (usage.TryGetProperty("prompt_tokens", out var inputTokens))
289+
{
290+
activity.SetTag(AWSSemanticConventions.AttributeGenAiInputTokens, inputTokens.GetInt32());
291+
}
292+
if (usage.TryGetProperty("completion_tokens", out var outputTokens))
293+
{
294+
activity.SetTag(AWSSemanticConventions.AttributeGenAiOutputTokens, outputTokens.GetInt32());
295+
}
296+
}
297+
if (jsonBody.TryGetValue("choices", out var choices))
298+
{
299+
if (choices[0].TryGetProperty("finish_reason", out var finishReasons))
300+
{
301+
activity.SetTag(AWSSemanticConventions.AttributeGenAiFinishReasons, new string[] { finishReasons.GetString() ?? string.Empty });
302+
}
303+
}
304+
}
305+
}
306+
catch (Exception ex)
307+
{
308+
Console.WriteLine("Exception: " + ex.Message);
309+
}
310+
}
311+
312+
private static void ProcessMistralModelAttributes(Activity activity, Dictionary<string, JsonElement> jsonBody, bool isRequest)
313+
{
314+
try
315+
{
316+
if (isRequest)
317+
{
318+
if (jsonBody.TryGetValue("top_p", out var topP))
319+
{
320+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTopP, topP.GetDouble());
321+
}
322+
323+
if (jsonBody.TryGetValue("temperature", out var temperature))
324+
{
325+
activity.SetTag(AWSSemanticConventions.AttributeGenAiTemperature, temperature.GetDouble());
326+
}
327+
328+
if (jsonBody.TryGetValue("max_tokens", out var maxTokens))
329+
{
330+
activity.SetTag(AWSSemanticConventions.AttributeGenAiMaxTokens, maxTokens.GetInt32());
331+
}
332+
333+
// input tokens not provided in Mistral response body, so we estimate the value based on input length
334+
if (jsonBody.TryGetValue("prompt", out var input))
335+
{
336+
activity.SetTag(AWSSemanticConventions.AttributeGenAiInputTokens, Convert.ToInt32(Math.Ceiling((double) (input.GetString()?.Length ?? 0) / 6)));
337+
}
338+
}
339+
else
340+
{
341+
if (jsonBody.TryGetValue("outputs", out var outputsArray))
342+
{
343+
var output = outputsArray[0];
344+
if (output.TryGetProperty("stop_reason", out var finishReasons))
345+
{
346+
activity.SetTag(AWSSemanticConventions.AttributeGenAiFinishReasons, new string[] { finishReasons.GetString() ?? string.Empty });
347+
}
348+
349+
// output tokens not provided in Mistral response body, so we estimate the value based on output length
350+
if (output.TryGetProperty("text", out var text))
351+
{
352+
activity.SetTag(AWSSemanticConventions.AttributeGenAiOutputTokens, Convert.ToInt32(Math.Ceiling((double) (text.GetString()?.Length ?? 0) / 6)));
353+
}
354+
}
355+
}
356+
}
357+
catch (Exception ex)
358+
{
359+
Console.WriteLine("Exception: " + ex.Message);
360+
}
361+
}
362+
}

‎src/OpenTelemetry.Instrumentation.AWS/Implementation/AWSSemanticConventions.cs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,14 @@ internal static class AWSSemanticConventions
2323
public const string AttributeAWSBedrockDataSourceId = "aws.bedrock.data_source.id";
2424

2525
// should be global convention for Gen AI attributes
26-
public const string AttributeGenAiModelId = "gen_ai.request.model";
2726
public const string AttributeGenAiSystem = "gen_ai.system";
27+
public const string AttributeGenAiModelId = "gen_ai.request.model";
28+
public const string AttributeGenAiTopP = "gen_ai.request.top_p";
29+
public const string AttributeGenAiTemperature = "gen_ai.request.temperature";
30+
public const string AttributeGenAiMaxTokens = "gen_ai.request.max_tokens";
31+
public const string AttributeGenAiInputTokens = "gen_ai.usage.input_tokens";
32+
public const string AttributeGenAiOutputTokens = "gen_ai.usage.output_tokens";
33+
public const string AttributeGenAiFinishReasons = "gen_ai.response.finish_reasons";
2834

2935
public const string AttributeHttpStatusCode = "http.status_code";
3036
public const string AttributeHttpResponseContentLength = "http.response_content_length";

0 commit comments

Comments
 (0)