Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit f12c439

Browse files
sabrennersabrenner
authored
fix(llmobs): properly parse "newer" anthropic models, cohere models from bedrock-runtime calls (#6383)
* utils fixes * update testagent version * update fixtures * generate cassettes * test changes * use fixed testagent version * add in break line * add in nullish coalesce for anthropic input messages
1 parent e6dcda4 commit f12c439

12 files changed

+577
-145
lines changed

‎docker-compose.yml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ services:
177177
- LDAP_PASSWORDS=password1,password2
178178

179179
testagent:
180-
image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.31.1
180+
image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.33.1
181181
ports:
182182
- "127.0.0.1:9126:9126"
183183
environment:
@@ -186,6 +186,7 @@ services:
186186
- ENABLED_CHECKS=trace_stall,meta_tracer_version_header,trace_count_header,trace_peer_service
187187
- PORT=9126
188188
- VCR_CASSETTES_DIRECTORY=/vcr-cassettes
189+
# - AWS_SECRET_ACCESS_KEY # uncomment this line for generating aws service cassettes
189190
volumes:
190191
# when there are other products not using the cassette feature from the test agent,
191192
# we can move this directory to a shared location here, and change the above

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

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -149,12 +149,24 @@ function extractRequestParams (params, provider) {
149149
})
150150
}
151151
case PROVIDER.ANTHROPIC: {
152-
const prompt = requestBody.prompt || requestBody.messages
152+
let prompt = requestBody.prompt
153+
if (Array.isArray(requestBody.messages)) { // newer claude models
154+
for (let idx = requestBody.messages.length - 1; idx >= 0; idx--) {
155+
const message = requestBody.messages[idx]
156+
if (message.role === 'user') {
157+
prompt = message.content?.filter(block => block.type === 'text')
158+
.map(block => block.text)
159+
.join('')
160+
break
161+
}
162+
}
163+
}
164+
153165
return new RequestParams({
154166
prompt,
155167
temperature: requestBody.temperature,
156168
topP: requestBody.top_p,
157-
maxTokens: requestBody.max_tokens_to_sample,
169+
maxTokens: requestBody.max_tokens_to_sample ?? requestBody.max_tokens,
158170
stopSequences: requestBody.stop_sequences
159171
})
160172
}
@@ -253,7 +265,13 @@ function extractTextAndResponseReason (response, provider, modelName) {
253265
break
254266
}
255267
case PROVIDER.ANTHROPIC: {
256-
return new Generation({ message: body.completion || body.content, finishReason: body.stop_reason })
268+
let message = body.completion
269+
if (Array.isArray(body.content)) { // newer claude models
270+
message = body.content.find(item => item.type === 'text')?.text ?? body.content
271+
} else if (body.content) {
272+
message = body.content
273+
}
274+
return new Generation({ message, finishReason: body.stop_reason })
257275
}
258276
case PROVIDER.COHERE: {
259277
if (modelName.includes('embed')) {
@@ -262,6 +280,15 @@ function extractTextAndResponseReason (response, provider, modelName) {
262280
return new Generation({ message: embeddings[0] })
263281
}
264282
}
283+
284+
if (body.text) {
285+
return new Generation({
286+
message: body.text,
287+
finishReason: body.finish_reason,
288+
choiceId: shouldSetChoiceIds ? body.response_id : undefined
289+
})
290+
}
291+
265292
const generations = body.generations || []
266293
if (generations.length > 0) {
267294
const generation = generations[0]

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

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
const { expect } = require('chai')
44
const { describe, it, before, after } = require('mocha')
5-
const nock = require('nock')
65

76
const agent = require('../../dd-trace/test/plugins/agent')
87
const { setup } = require('./spec_helpers')
@@ -26,7 +25,7 @@ describe('Plugin', () => {
2625
return agent.load('aws-sdk')
2726
})
2827

29-
before(done => {
28+
before(() => {
3029
const requireVersion = version === '3.0.0' ? '3.422.0' : '>=3.422.0'
3130
AWS = require(`../../../versions/${bedrockRuntimeClientName}@${requireVersion}`).get()
3231
const NodeHttpHandler =
@@ -35,45 +34,41 @@ describe('Plugin', () => {
3534
.NodeHttpHandler
3635

3736
bedrockRuntimeClient = new AWS.BedrockRuntimeClient(
38-
{ endpoint: 'http://127.0.0.1:4566', region: 'us-east-1', ServiceId: serviceName, requestHandler: new NodeHttpHandler() }
37+
{
38+
endpoint: { url: 'http://127.0.0.1:9126/vcr/bedrock-runtime' },
39+
region: 'us-east-1',
40+
ServiceId: serviceName,
41+
requestHandler: new NodeHttpHandler()
42+
}
3943
)
40-
done()
4144
})
4245

4346
after(async () => {
44-
nock.cleanAll()
4547
return agent.close({ ritmReset: false })
4648
})
4749

4850
models.forEach(model => {
49-
it(`should invoke model for provider:${model.provider}`, done => {
51+
it(`should invoke model for provider: ${model.provider} (ModelId: ${model.modelId})`, async () => {
5052
const request = {
5153
body: JSON.stringify(model.requestBody),
5254
contentType: 'application/json',
5355
accept: 'application/json',
5456
modelId: model.modelId
5557
}
5658

57-
const response = JSON.stringify(model.response)
58-
59-
nock('http://127.0.0.1:4566')
60-
.post(`/model/${model.modelId}/invoke`)
61-
.reply(200, response)
62-
6359
const command = new AWS.InvokeModelCommand(request)
6460

65-
agent.assertSomeTraces(traces => {
61+
const tracesPromise = agent.assertSomeTraces(traces => {
6662
const span = traces[0][0]
6763
expect(span.meta).to.include({
6864
'aws.operation': 'invokeModel',
6965
'aws.bedrock.request.model': model.modelId.split('.')[1],
7066
'aws.bedrock.request.model_provider': model.provider.toLowerCase(),
7167
})
72-
}).then(done).catch(done)
73-
74-
bedrockRuntimeClient.send(command, (err) => {
75-
if (err) return done(err)
7668
})
69+
70+
await bedrockRuntimeClient.send(command)
71+
await tracesPromise
7772
})
7873
})
7974
})

0 commit comments

Comments
 (0)