33const agent = require ( '../../dd-trace/test/plugins/agent' )
44const nock = require ( 'nock' )
55const { setup } = require ( './spec_helpers' )
6+ const { models, modelConfig } = require ( './fixtures/bedrockruntime' )
67
78const 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-
1810describe ( '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 } ,
121- requestBody : {
122- prompt : `\n\nHuman:${ prompt } ,
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