Skip to content

Commit 4886c38

Browse files
yahya-moumansabrenner
yahya-mouman
and
sabrenner
authored
Initial APM side for aws bedrock (#4937)
* Initial APM side for aws bedrock * add extract response tags * add extract response tags * remove hook * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime.js * added example test for invoke amazon * added example test for invoke amazon * update test with todos * update test with todos * Drop underscore in name * Update packages/datadog-plugin-aws-sdk/test/bedrock.spec.js Co-authored-by: Sam Brenner <[email protected]> * Constants normalization * Add Mistral AI * Add aws bedrock rec * remove file * added tests with mocked responses * added jamba support to AI21 lab * update bedrock version * Update tests * remove only * Update response extractions to only pick up first completion/generation * Update packages/datadog-plugin-aws-sdk/src/services/bedrockruntime.js Co-authored-by: Sam Brenner <[email protected]> * Change from constants to a struct for model provider * format * switch case * Add classes for generations and requestParams * Make constructors name object based. and stringify prompt if it's not a string * stringify message if it's not a string * es lint * fix bad variable name * add extra tags * camelCase and lint * camelCase and lint --------- Co-authored-by: Sam Brenner <[email protected]>
1 parent 71fc75f commit 4886c38

5 files changed

Lines changed: 540 additions & 1 deletion

File tree

‎packages/datadog-instrumentations/src/aws-sdk.js‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,8 @@ function getChannelSuffix (name) {
167167
'sns',
168168
'sqs',
169169
'states',
170-
'stepfunctions'
170+
'stepfunctions',
171+
'bedrock runtime'
171172
].includes(name)
172173
? name
173174
: 'default'

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

Lines changed: 295 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,295 @@
1+
'use strict'
2+
3+
const BaseAwsSdkPlugin = require('../base')
4+
const log = require('../../../dd-trace/src/log')
5+
6+
const PROVIDER = {
7+
AI21: 'AI21',
8+
AMAZON: 'AMAZON',
9+
ANTHROPIC: 'ANTHROPIC',
10+
COHERE: 'COHERE',
11+
META: 'META',
12+
STABILITY: 'STABILITY',
13+
MISTRAL: 'MISTRAL'
14+
}
15+
16+
const enabledOperations = ['invokeModel']
17+
18+
class BedrockRuntime extends BaseAwsSdkPlugin {
19+
static get id () { return 'bedrock runtime' }
20+
21+
isEnabled (request) {
22+
const operation = request.operation
23+
if (!enabledOperations.includes(operation)) {
24+
return false
25+
}
26+
27+
return super.isEnabled(request)
28+
}
29+
30+
generateTags (params, operation, response) {
31+
let tags = {}
32+
let modelName = ''
33+
let modelProvider = ''
34+
const modelMeta = params.modelId.split('.')
35+
if (modelMeta.length === 2) {
36+
[modelProvider, modelName] = modelMeta
37+
modelProvider = modelProvider.toUpperCase()
38+
} else {
39+
[, modelProvider, modelName] = modelMeta
40+
modelProvider = modelProvider.toUpperCase()
41+
}
42+
43+
const shouldSetChoiceIds = modelProvider === PROVIDER.COHERE && !modelName.includes('embed')
44+
45+
const requestParams = extractRequestParams(params, modelProvider)
46+
const textAndResponseReason = extractTextAndResponseReason(response, modelProvider, modelName, shouldSetChoiceIds)
47+
48+
tags = buildTagsFromParams(requestParams, textAndResponseReason, modelProvider, modelName, operation)
49+
50+
return tags
51+
}
52+
}
53+
54+
class Generation {
55+
constructor ({ message = '', finishReason = '', choiceId = '' } = {}) {
56+
// stringify message as it could be a single generated message as well as a list of embeddings
57+
this.message = typeof message === 'string' ? message : JSON.stringify(message) || ''
58+
this.finishReason = finishReason || ''
59+
this.choiceId = choiceId || undefined
60+
}
61+
}
62+
63+
class RequestParams {
64+
constructor ({
65+
prompt = '',
66+
temperature = undefined,
67+
topP = undefined,
68+
maxTokens = undefined,
69+
stopSequences = [],
70+
inputType = '',
71+
truncate = '',
72+
stream = '',
73+
n = undefined
74+
} = {}) {
75+
// TODO: set a truncation limit to prompt
76+
// stringify prompt as it could be a single prompt as well as a list of message objects
77+
this.prompt = typeof prompt === 'string' ? prompt : JSON.stringify(prompt) || ''
78+
this.temperature = temperature !== undefined ? temperature : undefined
79+
this.topP = topP !== undefined ? topP : undefined
80+
this.maxTokens = maxTokens !== undefined ? maxTokens : undefined
81+
this.stopSequences = stopSequences || []
82+
this.inputType = inputType || ''
83+
this.truncate = truncate || ''
84+
this.stream = stream || ''
85+
this.n = n !== undefined ? n : undefined
86+
}
87+
}
88+
89+
function extractRequestParams (params, provider) {
90+
const requestBody = JSON.parse(params.body)
91+
const modelId = params.modelId
92+
93+
switch (provider) {
94+
case PROVIDER.AI21: {
95+
let userPrompt = requestBody.prompt
96+
if (modelId.includes('jamba')) {
97+
for (const message of requestBody.messages) {
98+
if (message.role === 'user') {
99+
userPrompt = message.content // Return the content of the most recent user message
100+
}
101+
}
102+
}
103+
return new RequestParams({
104+
prompt: userPrompt,
105+
temperature: requestBody.temperature,
106+
topP: requestBody.top_p,
107+
maxTokens: requestBody.max_tokens,
108+
stopSequences: requestBody.stop_sequences
109+
})
110+
}
111+
case PROVIDER.AMAZON: {
112+
if (modelId.includes('embed')) {
113+
return new RequestParams({ prompt: requestBody.inputText })
114+
}
115+
const textGenerationConfig = requestBody.textGenerationConfig || {}
116+
return new RequestParams({
117+
prompt: requestBody.inputText,
118+
temperature: textGenerationConfig.temperature,
119+
topP: textGenerationConfig.topP,
120+
maxTokens: textGenerationConfig.maxTokenCount,
121+
stopSequences: textGenerationConfig.stopSequences
122+
})
123+
}
124+
case PROVIDER.ANTHROPIC: {
125+
const prompt = requestBody.prompt || requestBody.messages
126+
return new RequestParams({
127+
prompt,
128+
temperature: requestBody.temperature,
129+
topP: requestBody.top_p,
130+
maxTokens: requestBody.max_tokens_to_sample,
131+
stopSequences: requestBody.stop_sequences
132+
})
133+
}
134+
case PROVIDER.COHERE: {
135+
if (modelId.includes('embed')) {
136+
return new RequestParams({
137+
prompt: requestBody.texts,
138+
inputType: requestBody.input_type,
139+
truncate: requestBody.truncate
140+
})
141+
}
142+
return new RequestParams({
143+
prompt: requestBody.prompt,
144+
temperature: requestBody.temperature,
145+
topP: requestBody.p,
146+
maxTokens: requestBody.max_tokens,
147+
stopSequences: requestBody.stop_sequences,
148+
stream: requestBody.stream,
149+
n: requestBody.num_generations
150+
})
151+
}
152+
case PROVIDER.META: {
153+
return new RequestParams({
154+
prompt: requestBody.prompt,
155+
temperature: requestBody.temperature,
156+
topP: requestBody.top_p,
157+
maxTokens: requestBody.max_gen_len
158+
})
159+
}
160+
case PROVIDER.MISTRAL: {
161+
return new RequestParams({
162+
prompt: requestBody.prompt,
163+
temperature: requestBody.temperature,
164+
topP: requestBody.top_p,
165+
maxTokens: requestBody.max_tokens,
166+
stopSequences: requestBody.stop,
167+
topK: requestBody.top_k
168+
})
169+
}
170+
case PROVIDER.STABILITY: {
171+
return new RequestParams()
172+
}
173+
default: {
174+
return new RequestParams()
175+
}
176+
}
177+
}
178+
179+
function extractTextAndResponseReason (response, provider, modelName, shouldSetChoiceIds) {
180+
const body = JSON.parse(Buffer.from(response.body).toString('utf8'))
181+
182+
try {
183+
switch (provider) {
184+
case PROVIDER.AI21: {
185+
if (modelName.includes('jamba')) {
186+
const generations = body.choices || []
187+
if (generations.length > 0) {
188+
const generation = generations[0]
189+
return new Generation({
190+
message: generation.message,
191+
finishReason: generation.finish_reason,
192+
choiceId: shouldSetChoiceIds ? generation.id : undefined
193+
})
194+
}
195+
}
196+
const completions = body.completions || []
197+
if (completions.length > 0) {
198+
const completion = completions[0]
199+
return new Generation({
200+
message: completion.data?.text,
201+
finishReason: completion?.finishReason,
202+
choiceId: shouldSetChoiceIds ? completion?.id : undefined
203+
})
204+
}
205+
return new Generation()
206+
}
207+
case PROVIDER.AMAZON: {
208+
if (modelName.includes('embed')) {
209+
return new Generation({ message: body.embedding })
210+
}
211+
const results = body.results || []
212+
if (results.length > 0) {
213+
const result = results[0]
214+
return new Generation({ message: result.outputText, finishReason: result.completionReason })
215+
}
216+
break
217+
}
218+
case PROVIDER.ANTHROPIC: {
219+
return new Generation({ message: body.completion || body.content, finishReason: body.stop_reason })
220+
}
221+
case PROVIDER.COHERE: {
222+
if (modelName.includes('embed')) {
223+
const embeddings = body.embeddings || [[]]
224+
if (embeddings.length > 0) {
225+
return new Generation({ message: embeddings[0] })
226+
}
227+
}
228+
const generations = body.generations || []
229+
if (generations.length > 0) {
230+
const generation = generations[0]
231+
return new Generation({
232+
message: generation.text,
233+
finishReason: generation.finish_reason,
234+
choiceId: shouldSetChoiceIds ? generation.id : undefined
235+
})
236+
}
237+
break
238+
}
239+
case PROVIDER.META: {
240+
return new Generation({ message: body.generation, finishReason: body.stop_reason })
241+
}
242+
case PROVIDER.MISTRAL: {
243+
const mistralGenerations = body.outputs || []
244+
if (mistralGenerations.length > 0) {
245+
const generation = mistralGenerations[0]
246+
return new Generation({ message: generation.text, finishReason: generation.stop_reason })
247+
}
248+
break
249+
}
250+
case PROVIDER.STABILITY: {
251+
return new Generation()
252+
}
253+
default: {
254+
return new Generation()
255+
}
256+
}
257+
} catch (error) {
258+
log.warn('Unable to extract text/finishReason from response body. Defaulting to empty text/finishReason.')
259+
return new Generation()
260+
}
261+
262+
return new Generation()
263+
}
264+
265+
function buildTagsFromParams (requestParams, textAndResponseReason, modelProvider, modelName, operation) {
266+
const tags = {}
267+
268+
// add request tags
269+
tags['resource.name'] = operation
270+
tags['aws.bedrock.request.model'] = modelName
271+
tags['aws.bedrock.request.model_provider'] = modelProvider
272+
tags['aws.bedrock.request.prompt'] = requestParams.prompt
273+
tags['aws.bedrock.request.temperature'] = requestParams.temperature
274+
tags['aws.bedrock.request.top_p'] = requestParams.topP
275+
tags['aws.bedrock.request.max_tokens'] = requestParams.maxTokens
276+
tags['aws.bedrock.request.stop_sequences'] = requestParams.stopSequences
277+
tags['aws.bedrock.request.input_type'] = requestParams.inputType
278+
tags['aws.bedrock.request.truncate'] = requestParams.truncate
279+
tags['aws.bedrock.request.stream'] = requestParams.stream
280+
tags['aws.bedrock.request.n'] = requestParams.n
281+
282+
// add response tags
283+
if (modelName.includes('embed')) {
284+
tags['aws.bedrock.response.embedding_length'] = textAndResponseReason.message.length
285+
}
286+
if (textAndResponseReason.choiceId) {
287+
tags['aws.bedrock.response.choices.id'] = textAndResponseReason.choiceId
288+
}
289+
tags['aws.bedrock.response.choices.text'] = textAndResponseReason.message
290+
tags['aws.bedrock.response.choices.finish_reason'] = textAndResponseReason.finishReason
291+
292+
return tags
293+
}
294+
295+
module.exports = BedrockRuntime

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,5 @@ exports.sns = require('./sns')
1212
exports.sqs = require('./sqs')
1313
exports.states = require('./states')
1414
exports.stepfunctions = require('./stepfunctions')
15+
exports.bedrockruntime = require('./bedrockruntime')
1516
exports.default = require('./default')

0 commit comments

Comments
 (0)