Skip to content

Commit

Permalink
fixed AOT message
Browse files Browse the repository at this point in the history
  • Loading branch information
@AsakerMohd
AsakerMohd committed Jan 15, 2025
1 parent e49ddaf commit 3cc4599
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 61 deletions.
127 changes: 68 additions & 59 deletions src/OpenTelemetry.Instrumentation.AWS/Implementation/AWSLlmModelProcessor.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,89 +7,98 @@
using System.Diagnostics;
using System.Text;
using System.Text.Json;
using Amazon.Runtime;
using OpenTelemetry.AWS;

namespace OpenTelemetry.Instrumentation.AWS.Implementation;

internal class AWSLlmModelProcessor
{
internal static void ProcessGenAiAttributes(Activity activity, AmazonWebServiceRequest message, string modelName, AWSSemanticConventions awsSemanticConventions)
{
var messageType = message?.GetType();
var messageBodyProperty = messageType?.GetProperty("Body");
if (messageBodyProperty != null)
{
if (messageBodyProperty.GetValue(message) is MemoryStream body)
{
ProcessGenAiAttributes(activity, body, modelName, true, awsSemanticConventions);
}
}
}

internal static void ProcessGenAiAttributes(Activity activity, AmazonWebServiceResponse message, string modelName, AWSSemanticConventions awsSemanticConventions)
{
var messageType = message?.GetType();
var messageBodyProperty = messageType?.GetProperty("Body");
if (messageBodyProperty != null)
{
if (messageBodyProperty.GetValue(message) is MemoryStream body)
{
ProcessGenAiAttributes(activity, body, modelName, false, awsSemanticConventions);
}
}
}

#if NET
[UnconditionalSuppressMessage(
"Specify StringComparison for clarity",
"CA1307",
Justification = "Adding StringComparison only works for NET Core but not the framework.")]
internal static void ProcessGenAiAttributes<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicProperties)]T>(Activity activity, T message, string modelName, bool isRequest, AWSSemanticConventions awsSemanticConventions)
#else
internal static void ProcessGenAiAttributes<T>(Activity activity, T message, string modelName, bool isRequest, AWSSemanticConventions awsSemanticConventions)
#endif
private static void ProcessGenAiAttributes(Activity activity, MemoryStream messageBody, string modelName, bool isRequest, AWSSemanticConventions awsSemanticConventions)
{
// message can be either a request or a response. isRequest is used by the model-specific methods to determine
// whether to extract the request or response attributes.

// Currently, the .NET SDK does not expose "X-Amzn-Bedrock-*" HTTP headers in the response metadata, as per
// https://github.com/aws/aws-sdk-net/issues/3171. As a result, we can only extract attributes given what is in
// the response body. For the Claude, Command, and Mistral models, the input and output tokens are not provided
// in the response body, so we approximate their values by dividing the input and output lengths by 6, based on
// the Bedrock documentation here: https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html

if (message is null)
{
return;
}

var messageBodyProperty = typeof(T).GetProperty("Body");
if (messageBodyProperty != null)
try
{
if (messageBodyProperty.GetValue(message) is MemoryStream body)
{
try
{
var jsonString = Encoding.UTF8.GetString(body.ToArray());
var jsonString = Encoding.UTF8.GetString(messageBody.ToArray());
#if NET
var jsonObject = JsonSerializer.Deserialize(jsonString, SourceGenerationContext.Default.DictionaryStringJsonElement);
var jsonObject = JsonSerializer.Deserialize(jsonString, SourceGenerationContext.Default.DictionaryStringJsonElement);
#else
var jsonObject = JsonSerializer.Deserialize<Dictionary<string, JsonElement>>(jsonString);
var jsonObject = JsonSerializer.Deserialize<Dictionary<string, JsonElement>>(jsonString);
#endif
if (jsonObject == null)
{
return;
}
if (jsonObject == null)
{
return;
}

// extract model specific attributes based on model name
if (modelName.Contains("amazon.nova"))
{
ProcessNovaModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("amazon.titan"))
{
ProcessTitanModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("anthropic.claude"))
{
ProcessClaudeModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("meta.llama3"))
{
ProcessLlamaModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("cohere.command"))
{
ProcessCommandModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("ai21.jamba"))
{
ProcessJambaModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("mistral.mistral"))
{
ProcessMistralModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
}
catch (Exception ex)
{
AWSInstrumentationEventSource.Log.JsonParserException(nameof(AWSLlmModelProcessor), ex);
}
// extract model specific attributes based on model name
if (modelName.Contains("amazon.nova"))
{
ProcessNovaModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("amazon.titan"))
{
ProcessTitanModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("anthropic.claude"))
{
ProcessClaudeModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("meta.llama3"))
{
ProcessLlamaModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("cohere.command"))
{
ProcessCommandModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("ai21.jamba"))
{
ProcessJambaModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
else if (modelName.Contains("mistral.mistral"))
{
ProcessMistralModelAttributes(activity, jsonObject, isRequest, awsSemanticConventions);
}
}
catch (Exception ex)
{
AWSInstrumentationEventSource.Log.JsonParserException(nameof(AWSLlmModelProcessor), ex);
}
}

Expand Down
4 changes: 2 additions & 2 deletions src/OpenTelemetry.Instrumentation.AWS/Implementation/AWSTracingPipelineHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ private void AddResponseSpecificInformation(Activity activity, IExecutionContext
var modelString = model.ToString();
if (modelString != null)
{
AWSLlmModelProcessor.ProcessGenAiAttributes(activity, responseContext.Response, modelString, false, this.awsSemanticConventions);
AWSLlmModelProcessor.ProcessGenAiAttributes(activity, responseContext.Response, modelString, this.awsSemanticConventions);
}
}
}
Expand Down Expand Up @@ -183,7 +183,7 @@ private void AddRequestSpecificInformation(Activity activity, IRequestContext re
var modelString = model.ToString();
if (modelString != null)
{
AWSLlmModelProcessor.ProcessGenAiAttributes(activity, request, modelString, true, this.awsSemanticConventions);
AWSLlmModelProcessor.ProcessGenAiAttributes(activity, request, modelString, this.awsSemanticConventions);
}
}
}
Expand Down

0 comments on commit 3cc4599

Please sign in to comment.