fix: mistral ai model conditional

aws-observability · Nov 8, 2024 · c25d7de · c25d7de
1 parent a3355ea
commit c25d7de
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts
@@ -284,7 +284,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
         if (requestBody.top_p !== undefined) {
           spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
         }
-      } else if (modelId.includes('mistral.mistral')) {
+      } else if (modelId.includes('mistral')) {
         if (requestBody.prompt !== undefined) {
           // NOTE: We approximate the token count since this value is not directly available in the body
           // According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
@@ -386,7 +386,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
             responseBody.choices[0].finish_reason,
           ]);
         }
-      } else if (currentModelId.includes('mistral.mistral')) {
+      } else if (currentModelId.includes('mistral')) {
         if (responseBody.outputs?.[0]?.text !== undefined) {
           span.setAttribute(
             AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS,