(chore)LLMObs: instrument bedrock runtime invoke model (#5116)

* Update file structure * join spaces * add top_k * add index for plugins * add llmobs index * refactor apm tracing * add llmobs instrumentation for bedrock * add llmobs instrumentation for bedrock * rename * Update file structure * conflict * conflict * add index for plugins * add llmobs index * add llmobs instrumentation for bedrock * lint change * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/index.js Co-authored-by: Sam Brenner <[email protected]> * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js Co-authored-by: Sam Brenner <[email protected]> * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js Co-authored-by: Sam Brenner <[email protected]> * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/utils.js Co-authored-by: Sam Brenner <[email protected]> * CODEOWNERS * CODEOWNERS * remove shouldSetChoiceId override * remove shouldSetChoiceId override * remove shouldSetChoiceId override * remove shouldSetChoiceId override * lint * Update packages/dd-trace/src/llmobs/plugins/bedrockruntime.js Co-authored-by: Sam Brenner <[email protected]> * Update packages/dd-trace/test/llmobs/plugins/bedrockruntime/bedrockruntime.spec.js Co-authored-by: Sam Brenner <[email protected]> * resolving review comments * rename span * Add test fixtures and reduce redundant code * Add test to ci workflow * Update file structure * conflict * conflict * add index for plugins * add llmobs index * add llmobs instrumentation for bedrock * lint change * Update packages/dd-trace/src/llmobs/plugins/bedrockruntime.js Co-authored-by: Sam Brenner <[email protected]> * Update packages/dd-trace/test/llmobs/plugins/bedrockruntime/bedrockruntime.spec.js Co-authored-by: Sam Brenner <[email protected]> * resolving review comments * rename span * Add test fixtures and reduce redundant code * Add test to ci workflow * rename plugin name to aws sdk * update ci yml * add refactor for split --------- Co-authored-by: Sam Brenner <[email protected]>
DataDog · Jan 24, 2025 · 29c26b0 · 29c26b0
1 parent f41f5f7
commit 29c26b0
Show file tree

Hide file tree

Showing 6 changed files with 374 additions and 163 deletions.
diff --git a/.github/workflows/llmobs.yml b/.github/workflows/llmobs.yml
@@ -72,3 +72,24 @@ jobs:
         uses: ./.github/actions/testagent/logs
         with:
           suffix: llmobs-${{ github.job }}
+
+  aws-sdk:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: aws-sdk
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/testagent/start
+      - uses: ./.github/actions/node/setup
+      - uses: ./.github/actions/install
+      - uses: ./.github/actions/node/18
+      - run: yarn test:llmobs:plugins:ci
+        shell: bash
+      - uses: ./.github/actions/node/latest
+      - run: yarn test:llmobs:plugins:ci
+        shell: bash
+      - uses: codecov/codecov-action@v3
+      - if: always()
+        uses: ./.github/actions/testagent/logs
+        with:
+          suffix: llmobs-${{ github.job }}
diff --git a/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/index.js b/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/index.js
@@ -1,12 +1,14 @@
 const CompositePlugin = require('../../../../dd-trace/src/plugins/composite')
 const BedrockRuntimeTracing = require('./tracing')
+const BedrockRuntimeLLMObsPlugin = require('../../../../dd-trace/src/llmobs/plugins/bedrockruntime')
 class BedrockRuntimePlugin extends CompositePlugin {
   static get id () {
     return 'bedrockruntime'
   }
 
   static get plugins () {
     return {
+      llmobs: BedrockRuntimeLLMObsPlugin,
       tracing: BedrockRuntimeTracing
     }
   }

diff --git a/packages/datadog-plugin-aws-sdk/test/bedrockruntime.spec.js b/packages/datadog-plugin-aws-sdk/test/bedrockruntime.spec.js
@@ -3,18 +3,10 @@
 const agent = require('../../dd-trace/test/plugins/agent')
 const nock = require('nock')
 const { setup } = require('./spec_helpers')
+const { models, modelConfig } = require('./fixtures/bedrockruntime')
 
 const serviceName = 'bedrock-service-name-test'
 
-const PROVIDER = {
-  AI21: 'AI21',
-  AMAZON: 'AMAZON',
-  ANTHROPIC: 'ANTHROPIC',
-  COHERE: 'COHERE',
-  META: 'META',
-  MISTRAL: 'MISTRAL'
-}
-
 describe('Plugin', () => {
   describe('aws-sdk (bedrockruntime)', function () {
     setup()
@@ -44,157 +36,6 @@ describe('Plugin', () => {
           return agent.close({ ritmReset: false })
         })
 
-        const prompt = 'What is the capital of France?'
-        const temperature = 0.5
-        const topP = 1
-        const topK = 1
-        const maxTokens = 512
-
-        const models = [
-          {
-            provider: PROVIDER.AMAZON,
-            modelId: 'amazon.titan-text-lite-v1',
-            userPrompt: prompt,
-            requestBody: {
-              inputText: prompt,
-              textGenerationConfig: {
-                temperature,
-                topP,
-                maxTokenCount: maxTokens
-              }
-            },
-            response: {
-              inputTextTokenCount: 7,
-              results: {
-                inputTextTokenCount: 7,
-                results: [
-                  {
-                    tokenCount: 35,
-                    outputText: '\n' +
-                      'Paris is the capital of France. France is a country that is located in Western Europe. ' +
-                      'Paris is one of the most populous cities in the European Union. ',
-                    completionReason: 'FINISH'
-                  }
-                ]
-              }
-            }
-          },
-          {
-            provider: PROVIDER.AI21,
-            modelId: 'ai21.jamba-1-5-mini-v1',
-            userPrompt: prompt,
-            requestBody: {
-              messages: [
-                {
-                  role: 'user',
-                  content: prompt
-                }
-              ],
-              max_tokens: maxTokens,
-              temperature,
-              top_p: topP,
-              top_k: topK
-            },
-            response: {
-              id: 'req_0987654321',
-              choices: [
-                {
-                  index: 0,
-                  message: {
-                    role: 'assistant',
-                    content: 'The capital of France is Paris.'
-                  },
-                  finish_reason: 'stop'
-                }
-              ],
-              usage: {
-                prompt_tokens: 10,
-                completion_tokens: 7,
-                total_tokens: 17
-              }
-            }
-          },
-          {
-            provider: PROVIDER.ANTHROPIC,
-            modelId: 'anthropic.claude-v2',
-            userPrompt: `\n\nHuman:${prompt}\n\nAssistant:`,
-            requestBody: {
-              prompt: `\n\nHuman:${prompt}\n\nAssistant:`,
-              temperature,
-              top_p: topP,
-              top_k: topK,
-              max_tokens_to_sample: maxTokens
-            },
-            response: {
-              type: 'completion',
-              completion: ' Paris is the capital of France.',
-              stop_reason: 'stop_sequence',
-              stop: '\n\nHuman:'
-            }
-          },
-          {
-            provider: PROVIDER.COHERE,
-            modelId: 'cohere.command-light-text-v14',
-            userPrompt: prompt,
-            requestBody: {
-              prompt,
-              temperature,
-              p: topP,
-              k: topK,
-              max_tokens: maxTokens
-            },
-            response: {
-              id: '91c65da4-e2cd-4930-a4a9-f5c68c8a137c',
-              generations: [
-                {
-                  id: 'c040d384-ad9c-4d15-8c2f-f36fbfb0eb55',
-                  text: ' The capital of France is Paris. \n',
-                  finish_reason: 'COMPLETE'
-                }
-              ],
-              prompt: 'What is the capital of France?'
-            }
-
-          },
-          {
-            provider: PROVIDER.META,
-            modelId: 'meta.llama3-70b-instruct-v1',
-            userPrompt: prompt,
-            requestBody: {
-              prompt,
-              temperature,
-              top_p: topP,
-              max_gen_len: maxTokens
-            },
-            response: {
-              generation: '\n\nThe capital of France is Paris.',
-              prompt_token_count: 10,
-              generation_token_count: 7,
-              stop_reason: 'stop'
-            }
-          },
-          {
-            provider: PROVIDER.MISTRAL,
-            modelId: 'mistral.mistral-7b-instruct-v0',
-            userPrompt: prompt,
-            requestBody: {
-              prompt,
-              max_tokens: maxTokens,
-              temperature,
-              top_p: topP,
-              top_k: topK
-            },
-            response: {
-              outputs: [
-                {
-                  text: 'The capital of France is Paris.',
-                  stop_reason: 'stop'
-                }
-              ]
-            }
-          }
-        ]
-
         models.forEach(model => {
           it(`should invoke model for provider:${model.provider}`, done => {
             const request = {
@@ -221,9 +62,9 @@ describe('Plugin', () => {
                 'aws.bedrock.request.prompt': model.userPrompt
               })
               expect(span.metrics).to.include({
-                'aws.bedrock.request.temperature': temperature,
-                'aws.bedrock.request.top_p': topP,
-                'aws.bedrock.request.max_tokens': maxTokens
+                'aws.bedrock.request.temperature': modelConfig.temperature,
+                'aws.bedrock.request.top_p': modelConfig.topP,
+                'aws.bedrock.request.max_tokens': modelConfig.maxTokens
               })
             }).then(done).catch(done)