diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/aws-span-processing-util.ts b/aws-distro-opentelemetry-node-autoinstrumentation/src/aws-span-processing-util.ts index 46c8a6fd..6ad9f8c5 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/src/aws-span-processing-util.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/src/aws-span-processing-util.ts @@ -47,6 +47,12 @@ export class AwsSpanProcessingUtil { // TODO: Use Semantic Conventions once upgraded static GEN_AI_REQUEST_MODEL: string = 'gen_ai.request.model'; static GEN_AI_SYSTEM: string = 'gen_ai.system'; + static GEN_AI_REQUEST_MAX_TOKENS: string = 'gen_ai.request.max_tokens'; + static GEN_AI_REQUEST_TEMPERATURE: string = 'gen_ai.request.temperature'; + static GEN_AI_REQUEST_TOP_P: string = 'gen_ai.request.top_p'; + static GEN_AI_RESPONSE_FINISH_REASONS: string = 'gen_ai.response.finish_reasons'; + static GEN_AI_USAGE_INPUT_TOKENS: string = 'gen_ai.usage.input_tokens'; + static GEN_AI_USAGE_OUTPUT_TOKENS: string = 'gen_ai.usage.output_tokens'; static getDialectKeywords(): string[] { return SQL_DIALECT_KEYWORDS_JSON.keywords; 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 index ffeafacd..69a3a285 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts @@ -211,6 +211,79 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL] = modelId; } + if (request.commandInput?.body) { + const requestBody = JSON.parse(request.commandInput.body); + if (modelId.includes('amazon.titan')) { + if (requestBody.textGenerationConfig?.temperature !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = + requestBody.textGenerationConfig.temperature; + } + if (requestBody.textGenerationConfig?.topP !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.textGenerationConfig.topP; + } + if (requestBody.textGenerationConfig?.maxTokenCount !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = + requestBody.textGenerationConfig.maxTokenCount; + } + } else if (modelId.includes('anthropic.claude')) { + if (requestBody.max_tokens !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.top_p !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p; + } + } else if (modelId.includes('meta.llama')) { + if (requestBody.max_gen_len !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_gen_len; + } + if (requestBody.temperature !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.top_p !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p; + } + } else if (modelId.includes('cohere.command')) { + if (requestBody.max_tokens !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.p !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.p; + } + } else if (modelId.includes('ai21.jamba')) { + if (requestBody.max_tokens !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.top_p !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p; + } + } else if (modelId.includes('mistral.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. + // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html + spanAttributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.prompt.length / 6); + } + if (requestBody.max_tokens !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.top_p !== undefined) { + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p; + } + } + } + return { isIncoming, spanAttributes, @@ -218,4 +291,93 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { spanName, }; } + + responseHook(response: NormalizedResponse, span: Span, tracer: Tracer, config: AwsSdkInstrumentationConfig): void { + const currentModelId = response.request.commandInput?.modelId; + if (response.data?.body) { + const decodedResponseBody = new TextDecoder().decode(response.data.body); + const responseBody = JSON.parse(decodedResponseBody); + if (currentModelId.includes('amazon.titan')) { + if (responseBody.inputTextTokenCount !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.inputTextTokenCount); + } + if (responseBody.results?.[0]?.tokenCount !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.results[0].tokenCount); + } + if (responseBody.results?.[0]?.completionReason !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [ + responseBody.results[0].completionReason, + ]); + } + } else if (currentModelId.includes('anthropic.claude')) { + if (responseBody.usage?.input_tokens !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.input_tokens); + } + if (responseBody.usage?.output_tokens !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.output_tokens); + } + if (responseBody.stop_reason !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]); + } + } else if (currentModelId.includes('meta.llama')) { + if (responseBody.prompt_token_count !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.prompt_token_count); + } + if (responseBody.generation_token_count !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.generation_token_count); + } + if (responseBody.stop_reason !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]); + } + } else if (currentModelId.includes('cohere.command')) { + if (responseBody.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. + // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, Math.ceil(responseBody.prompt.length / 6)); + } + if (responseBody.generations?.[0]?.text !== undefined) { + span.setAttribute( + AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, + // 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. + // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html + Math.ceil(responseBody.generations[0].text.length / 6) + ); + } + if (responseBody.generations?.[0]?.finish_reason !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [ + responseBody.generations[0].finish_reason, + ]); + } + } else if (currentModelId.includes('ai21.jamba')) { + if (responseBody.usage?.prompt_tokens !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.prompt_tokens); + } + if (responseBody.usage?.completion_tokens !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.completion_tokens); + } + if (responseBody.choices?.[0]?.finish_reason !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [ + responseBody.choices[0].finish_reason, + ]); + } + } else if (currentModelId.includes('mistral.mistral')) { + if (responseBody.outputs?.[0]?.text !== undefined) { + span.setAttribute( + AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, + // 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. + // https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html + Math.ceil(responseBody.outputs[0].text.length / 6) + ); + } + if (responseBody.outputs?.[0]?.stop_reason !== undefined) { + span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [ + responseBody.outputs[0].stop_reason, + ]); + } + } + } + } } diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts index 23edc001..97c63a63 100644 --- a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts +++ b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts @@ -276,19 +276,102 @@ describe('BedrockRuntime', () => { }); describe('InvokeModel', () => { - it('adds modelId to span', async () => { - const dummyModelId: string = 'ABCDEFGH'; - const dummyBody: string = 'HGFEDCBA'; + it('Add AI21 Jamba model attributes to span', async () => { + const modelId: string = 'ai21.jamba-1-5-large-v1:0'; + const prompt: string = 'Describe the purpose of a compiler in one line.'; + const nativeRequest: any = { + messages: [ + { + role: 'user', + content: prompt, + }, + ], + top_p: 0.8, + temperature: 0.6, + max_tokens: 512, + }; + const mockRequestBody: string = JSON.stringify(nativeRequest); + const mockResponseBody: any = { + stop_reason: 'end_turn', + usage: { + prompt_tokens: 21, + completion_tokens: 24, + }, + choices: [ + { + finish_reason: 'stop', + }, + ], + request: { + commandInput: { + modelId: modelId, + }, + }, + }; + + nock(`https://bedrock-runtime.${region}.amazonaws.com`) + .post(`/model/${encodeURIComponent(modelId)}/invoke`) + .reply(200, mockResponseBody); - nock(`https://bedrock-runtime.${region}.amazonaws.com`).post(`/model/${dummyModelId}/invoke`).reply(200, { - modelId: dummyModelId, - body: dummyBody, + await bedrock + .invokeModel({ + modelId: modelId, + body: mockRequestBody, + }) + .catch((err: any) => { + console.log('error', err); + }); + + const testSpans: ReadableSpan[] = getTestSpans(); + const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => { + return s.name === 'BedrockRuntime.InvokeModel'; }); + expect(invokeModelSpans.length).toBe(1); + const invokeModelSpan = invokeModelSpans[0]; + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock'); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(512); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.6); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(0.8); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(21); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(24); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['stop']); + expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); + }); + it('Add Amazon Titan model attributes to span', async () => { + const modelId: string = 'amazon.titan-text-express-v1'; + const prompt: string = 'Complete this text. It was the best of times it was the worst...'; + const nativeRequest: any = { + inputText: prompt, + textGenerationConfig: { + maxTokenCount: 4096, + stopSequences: [], + temperature: 0, + topP: 1, + }, + }; + const mockRequestBody: string = JSON.stringify(nativeRequest); + const mockResponseBody: any = { + inputTextTokenCount: 15, + results: [ + { + tokenCount: 13, + completionReason: 'CONTENT_FILTERED', + }, + ], + }; + + nock(`https://bedrock-runtime.${region}.amazonaws.com`) + .post(`/model/${modelId}/invoke`) + .reply(200, mockResponseBody); await bedrock .invokeModel({ - modelId: dummyModelId, - body: dummyBody, + modelId: modelId, + body: mockRequestBody, }) .catch((err: any) => {}); @@ -301,7 +384,252 @@ describe('BedrockRuntime', () => { expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined(); expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined(); expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined(); - expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(dummyModelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock'); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(4096); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(1); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(15); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(13); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual([ + 'CONTENT_FILTERED', + ]); + expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); + }); + + it('Add Anthropic Claude model attributes to span', async () => { + const modelId: string = 'anthropic.claude-3-5-sonnet-20240620-v1:0'; + const prompt: string = 'Complete this text. It was the best of times it was the worst...'; + const nativeRequest: any = { + anthropic_version: 'bedrock-2023-05-31', + max_tokens: 1000, + temperature: 1.0, + top_p: 1, + messages: [ + { + role: 'user', + content: [{ type: 'text', text: prompt }], + }, + ], + }; + const mockRequestBody: string = JSON.stringify(nativeRequest); + const mockResponseBody: any = { + stop_reason: 'end_turn', + usage: { + input_tokens: 15, + output_tokens: 13, + }, + request: { + commandInput: { + modelId: modelId, + }, + }, + }; + + nock(`https://bedrock-runtime.${region}.amazonaws.com`) + .post(`/model/${encodeURIComponent(modelId)}/invoke`) + .reply(200, mockResponseBody); + + await bedrock + .invokeModel({ + modelId: modelId, + body: mockRequestBody, + }) + .catch((err: any) => { + console.log('error', err); + }); + + const testSpans: ReadableSpan[] = getTestSpans(); + const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => { + return s.name === 'BedrockRuntime.InvokeModel'; + }); + expect(invokeModelSpans.length).toBe(1); + const invokeModelSpan = invokeModelSpans[0]; + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock'); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(1000); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(1.0); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(1); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(15); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(13); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['end_turn']); + expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); + }); + + it('Add Cohere Command model attributes to span', async () => { + const modelId: string = 'cohere.command-light-text-v14'; + const prompt: string = "Describe the purpose of a 'hello world' program in one line"; + const nativeRequest: any = { + prompt: prompt, + max_tokens: 512, + temperature: 0.5, + p: 0.65, + }; + const mockRequestBody: string = JSON.stringify(nativeRequest); + const mockResponseBody: any = { + generations: [ + { + finish_reason: 'COMPLETE', + text: 'test-generation-text', + }, + ], + prompt: prompt, + request: { + commandInput: { + modelId: modelId, + }, + }, + }; + + nock(`https://bedrock-runtime.${region}.amazonaws.com`) + .post(`/model/${encodeURIComponent(modelId)}/invoke`) + .reply(200, mockResponseBody); + + await bedrock + .invokeModel({ + modelId: modelId, + body: mockRequestBody, + }) + .catch((err: any) => { + console.log('error', err); + }); + + const testSpans: ReadableSpan[] = getTestSpans(); + const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => { + return s.name === 'BedrockRuntime.InvokeModel'; + }); + expect(invokeModelSpans.length).toBe(1); + const invokeModelSpan = invokeModelSpans[0]; + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock'); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(512); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.5); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(0.65); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(10); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(4); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['COMPLETE']); + expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); + }); + + it('Add Meta Llama model attributes to span', async () => { + const modelId: string = 'meta.llama2-13b-chat-v1'; + const prompt: string = 'Describe the purpose of an interpreter program in one line.'; + const nativeRequest: any = { + prompt, + max_gen_len: 512, + temperature: 0.5, + top_p: 0.9, + }; + const mockRequestBody: string = JSON.stringify(nativeRequest); + const mockResponseBody: any = { + prompt_token_count: 31, + generation_token_count: 49, + stop_reason: 'stop', + request: { + commandInput: { + modelId: modelId, + }, + }, + }; + + nock(`https://bedrock-runtime.${region}.amazonaws.com`) + .post(`/model/${encodeURIComponent(modelId)}/invoke`) + .reply(200, mockResponseBody); + + await bedrock + .invokeModel({ + modelId: modelId, + body: mockRequestBody, + }) + .catch((err: any) => { + console.log('error', err); + }); + + const testSpans: ReadableSpan[] = getTestSpans(); + const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => { + return s.name === 'BedrockRuntime.InvokeModel'; + }); + expect(invokeModelSpans.length).toBe(1); + const invokeModelSpan = invokeModelSpans[0]; + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock'); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(512); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.5); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(0.9); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(31); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(49); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['stop']); + expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); + }); + + it('Add Mistral AI model attributes to span', async () => { + const modelId: string = 'mistral.mistral-7b-instruct-v0:2'; + const prompt: string = ` + [INST] + In Bash, how do I list all text files in the current directory + (excluding subdirectories) that have been modified in the last month? + [/INST] + `; + const nativeRequest: any = { + prompt: prompt, + max_tokens: 4096, + temperature: 0.75, + top_p: 1.0, + }; + const mockRequestBody: string = JSON.stringify(nativeRequest); + const mockResponseBody: any = { + outputs: [ + { + text: 'test-output-text', + stop_reason: 'stop', + }, + ], + request: { + commandInput: { + modelId: modelId, + }, + }, + }; + + nock(`https://bedrock-runtime.${region}.amazonaws.com`) + .post(`/model/${encodeURIComponent(modelId)}/invoke`) + .reply(200, mockResponseBody); + + await bedrock + .invokeModel({ + modelId: modelId, + body: mockRequestBody, + }) + .catch((err: any) => { + console.log('error', err); + }); + + const testSpans: ReadableSpan[] = getTestSpans(); + const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => { + return s.name === 'BedrockRuntime.InvokeModel'; + }); + expect(invokeModelSpans.length).toBe(1); + const invokeModelSpan = invokeModelSpans[0]; + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined(); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock'); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(4096); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.75); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(1.0); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(31); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(3); + expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['stop']); expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT); }); });