Skip to content

Commit 29c26b0

Browse files
yahya-moumansabrenner
yahya-mouman
and
sabrenner
authored
(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 <106700075+sabrenner@users.noreply.github.com> * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com> * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com> * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/utils.js Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com> * 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 <106700075+sabrenner@users.noreply.github.com> * Update packages/dd-trace/test/llmobs/plugins/bedrockruntime/bedrockruntime.spec.js Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com> * 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 <106700075+sabrenner@users.noreply.github.com> * Update packages/dd-trace/test/llmobs/plugins/bedrockruntime/bedrockruntime.spec.js Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com> * 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 <106700075+sabrenner@users.noreply.github.com>
1 parent f41f5f7 commit 29c26b0

6 files changed

Lines changed: 374 additions & 163 deletions

File tree

‎.github/workflows/llmobs.yml‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,3 +72,24 @@ jobs:
7272
uses: ./.github/actions/testagent/logs
7373
with:
7474
suffix: llmobs-${{ github.job }}
75+
76+
aws-sdk:
77+
runs-on: ubuntu-latest
78+
env:
79+
PLUGINS: aws-sdk
80+
steps:
81+
- uses: actions/checkout@v4
82+
- uses: ./.github/actions/testagent/start
83+
- uses: ./.github/actions/node/setup
84+
- uses: ./.github/actions/install
85+
- uses: ./.github/actions/node/18
86+
- run: yarn test:llmobs:plugins:ci
87+
shell: bash
88+
- uses: ./.github/actions/node/latest
89+
- run: yarn test:llmobs:plugins:ci
90+
shell: bash
91+
- uses: codecov/codecov-action@v3
92+
- if: always()
93+
uses: ./.github/actions/testagent/logs
94+
with:
95+
suffix: llmobs-${{ github.job }}

‎packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/index.js‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
const CompositePlugin = require('../../../../dd-trace/src/plugins/composite')
22
const BedrockRuntimeTracing = require('./tracing')
3+
const BedrockRuntimeLLMObsPlugin = require('../../../../dd-trace/src/llmobs/plugins/bedrockruntime')
34
class BedrockRuntimePlugin extends CompositePlugin {
45
static get id () {
56
return 'bedrockruntime'
67
}
78

89
static get plugins () {
910
return {
11+
llmobs: BedrockRuntimeLLMObsPlugin,
1012
tracing: BedrockRuntimeTracing
1113
}
1214
}

‎packages/datadog-plugin-aws-sdk/test/bedrockruntime.spec.js‎

Lines changed: 4 additions & 163 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,10 @@
33
const agent = require('../../dd-trace/test/plugins/agent')
44
const nock = require('nock')
55
const { setup } = require('./spec_helpers')
6+
const { models, modelConfig } = require('./fixtures/bedrockruntime')
67

78
const serviceName = 'bedrock-service-name-test'
89

9-
const PROVIDER = {
10-
AI21: 'AI21',
11-
AMAZON: 'AMAZON',
12-
ANTHROPIC: 'ANTHROPIC',
13-
COHERE: 'COHERE',
14-
META: 'META',
15-
MISTRAL: 'MISTRAL'
16-
}
17-
1810
describe('Plugin', () => {
1911
describe('aws-sdk (bedrockruntime)', function () {
2012
setup()
@@ -44,157 +36,6 @@ describe('Plugin', () => {
4436
return agent.close({ ritmReset: false })
4537
})
4638

47-
const prompt = 'What is the capital of France?'
48-
const temperature = 0.5
49-
const topP = 1
50-
const topK = 1
51-
const maxTokens = 512
52-
53-
const models = [
54-
{
55-
provider: PROVIDER.AMAZON,
56-
modelId: 'amazon.titan-text-lite-v1',
57-
userPrompt: prompt,
58-
requestBody: {
59-
inputText: prompt,
60-
textGenerationConfig: {
61-
temperature,
62-
topP,
63-
maxTokenCount: maxTokens
64-
}
65-
},
66-
response: {
67-
inputTextTokenCount: 7,
68-
results: {
69-
inputTextTokenCount: 7,
70-
results: [
71-
{
72-
tokenCount: 35,
73-
outputText: '\n' +
74-
'Paris is the capital of France. France is a country that is located in Western Europe. ' +
75-
'Paris is one of the most populous cities in the European Union. ',
76-
completionReason: 'FINISH'
77-
}
78-
]
79-
}
80-
}
81-
},
82-
{
83-
provider: PROVIDER.AI21,
84-
modelId: 'ai21.jamba-1-5-mini-v1',
85-
userPrompt: prompt,
86-
requestBody: {
87-
messages: [
88-
{
89-
role: 'user',
90-
content: prompt
91-
}
92-
],
93-
max_tokens: maxTokens,
94-
temperature,
95-
top_p: topP,
96-
top_k: topK
97-
},
98-
response: {
99-
id: 'req_0987654321',
100-
choices: [
101-
{
102-
index: 0,
103-
message: {
104-
role: 'assistant',
105-
content: 'The capital of France is Paris.'
106-
},
107-
finish_reason: 'stop'
108-
}
109-
],
110-
usage: {
111-
prompt_tokens: 10,
112-
completion_tokens: 7,
113-
total_tokens: 17
114-
}
115-
}
116-
},
117-
{
118-
provider: PROVIDER.ANTHROPIC,
119-
modelId: 'anthropic.claude-v2',
120-
userPrompt: `\n\nHuman:${prompt}\n\nAssistant:`,
121-
requestBody: {
122-
prompt: `\n\nHuman:${prompt}\n\nAssistant:`,
123-
temperature,
124-
top_p: topP,
125-
top_k: topK,
126-
max_tokens_to_sample: maxTokens
127-
},
128-
response: {
129-
type: 'completion',
130-
completion: ' Paris is the capital of France.',
131-
stop_reason: 'stop_sequence',
132-
stop: '\n\nHuman:'
133-
}
134-
},
135-
{
136-
provider: PROVIDER.COHERE,
137-
modelId: 'cohere.command-light-text-v14',
138-
userPrompt: prompt,
139-
requestBody: {
140-
prompt,
141-
temperature,
142-
p: topP,
143-
k: topK,
144-
max_tokens: maxTokens
145-
},
146-
response: {
147-
id: '91c65da4-e2cd-4930-a4a9-f5c68c8a137c',
148-
generations: [
149-
{
150-
id: 'c040d384-ad9c-4d15-8c2f-f36fbfb0eb55',
151-
text: ' The capital of France is Paris. \n',
152-
finish_reason: 'COMPLETE'
153-
}
154-
],
155-
prompt: 'What is the capital of France?'
156-
}
157-
158-
},
159-
{
160-
provider: PROVIDER.META,
161-
modelId: 'meta.llama3-70b-instruct-v1',
162-
userPrompt: prompt,
163-
requestBody: {
164-
prompt,
165-
temperature,
166-
top_p: topP,
167-
max_gen_len: maxTokens
168-
},
169-
response: {
170-
generation: '\n\nThe capital of France is Paris.',
171-
prompt_token_count: 10,
172-
generation_token_count: 7,
173-
stop_reason: 'stop'
174-
}
175-
},
176-
{
177-
provider: PROVIDER.MISTRAL,
178-
modelId: 'mistral.mistral-7b-instruct-v0',
179-
userPrompt: prompt,
180-
requestBody: {
181-
prompt,
182-
max_tokens: maxTokens,
183-
temperature,
184-
top_p: topP,
185-
top_k: topK
186-
},
187-
response: {
188-
outputs: [
189-
{
190-
text: 'The capital of France is Paris.',
191-
stop_reason: 'stop'
192-
}
193-
]
194-
}
195-
}
196-
]
197-
19839
models.forEach(model => {
19940
it(`should invoke model for provider:${model.provider}`, done => {
20041
const request = {
@@ -221,9 +62,9 @@ describe('Plugin', () => {
22162
'aws.bedrock.request.prompt': model.userPrompt
22263
})
22364
expect(span.metrics).to.include({
224-
'aws.bedrock.request.temperature': temperature,
225-
'aws.bedrock.request.top_p': topP,
226-
'aws.bedrock.request.max_tokens': maxTokens
65+
'aws.bedrock.request.temperature': modelConfig.temperature,
66+
'aws.bedrock.request.top_p': modelConfig.topP,
67+
'aws.bedrock.request.max_tokens': modelConfig.maxTokens
22768
})
22869
}).then(done).catch(done)
22970

0 commit comments

Comments
 (0)