Skip to content

Commit 32493bf

Browse files
RulaKhaledRulaKhaled
committed
update with tests
1 parent 858c820 commit 32493bf

File tree

3 files changed

+231
-0
lines changed

3 files changed

+231
-0
lines changed

‎dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-tools.mjs

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
import { instrumentAnthropicAiClient } from '@sentry/core';
2+
import * as Sentry from '@sentry/node';
3+
4+
function createMockStreamEvents(model = 'claude-3-haiku-20240307') {
5+
async function* generator() {
6+
// initial message metadata with id/model and input tokens
7+
yield {
8+
type: 'content_block_start',
9+
message: {
10+
id: 'msg_stream_tool_1',
11+
type: 'message',
12+
role: 'assistant',
13+
model,
14+
content: [],
15+
stop_reason: 'end_turn',
16+
usage: { input_tokens: 11 },
17+
},
18+
};
19+
20+
// streamed text
21+
yield { type: 'content_block_delta', delta: { text: 'Starting tool...' } };
22+
23+
// tool_use streamed via partial json
24+
yield { type: 'content_block_start', index: 0, content_block: { type: 'tool_use', id: 'tool_weather_2', name: 'weather' } };
25+
yield { type: 'content_block_delta', index: 0, delta: { partial_json: '{"city":' } };
26+
yield { type: 'content_block_delta', index: 0, delta: { partial_json: '"Paris"}' } };
27+
yield { type: 'content_block_stop', index: 0 };
28+
29+
// more text
30+
yield { type: 'content_block_delta', delta: { text: 'Done.' } };
31+
32+
// final usage
33+
yield { type: 'message_delta', usage: { output_tokens: 9 } };
34+
}
35+
return generator();
36+
}
37+
38+
class MockAnthropic {
39+
constructor(config) {
40+
this.apiKey = config.apiKey;
41+
this.messages = {
42+
create: this._messagesCreate.bind(this),
43+
stream: this._messagesStream.bind(this),
44+
};
45+
}
46+
47+
async _messagesCreate(params) {
48+
await new Promise(resolve => setTimeout(resolve, 5));
49+
if (params?.stream) {
50+
return createMockStreamEvents(params.model);
51+
}
52+
return {
53+
id: 'msg_mock_no_stream',
54+
type: 'message',
55+
model: params.model,
56+
role: 'assistant',
57+
content: [{ type: 'text', text: 'No stream' }],
58+
usage: { input_tokens: 2, output_tokens: 3 },
59+
};
60+
}
61+
62+
async _messagesStream(params) {
63+
await new Promise(resolve => setTimeout(resolve, 5));
64+
return createMockStreamEvents(params?.model);
65+
}
66+
}
67+
68+
async function run() {
69+
await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
70+
const mockClient = new MockAnthropic({ apiKey: 'mock-api-key' });
71+
const client = instrumentAnthropicAiClient(mockClient);
72+
73+
// stream via create(stream:true)
74+
const stream1 = await client.messages.create({
75+
model: 'claude-3-haiku-20240307',
76+
messages: [{ role: 'user', content: 'Need the weather' }],
77+
tools: [
78+
{
79+
name: 'weather',
80+
description: 'Get weather',
81+
input_schema: { type: 'object', properties: { city: { type: 'string' } }, required: ['city'] },
82+
},
83+
],
84+
stream: true,
85+
});
86+
for await (const _ of stream1) {
87+
void _;
88+
}
89+
90+
// stream via messages.stream
91+
const stream2 = await client.messages.stream({
92+
model: 'claude-3-haiku-20240307',
93+
messages: [{ role: 'user', content: 'Need the weather' }],
94+
tools: [
95+
{
96+
name: 'weather',
97+
description: 'Get weather',
98+
input_schema: { type: 'object', properties: { city: { type: 'string' } }, required: ['city'] },
99+
},
100+
],
101+
});
102+
for await (const _ of stream2) {
103+
void _;
104+
}
105+
});
106+
}
107+
108+
run();
109+
110+

‎dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-tools.mjs

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
import { instrumentAnthropicAiClient } from '@sentry/core';
2+
import * as Sentry from '@sentry/node';
3+
4+
class MockAnthropic {
5+
constructor(config) {
6+
this.apiKey = config.apiKey;
7+
8+
this.messages = {
9+
create: this._messagesCreate.bind(this),
10+
};
11+
}
12+
13+
async _messagesCreate(params) {
14+
await new Promise(resolve => setTimeout(resolve, 5));
15+
16+
return {
17+
id: 'msg_mock_tool_1',
18+
type: 'message',
19+
model: params.model,
20+
role: 'assistant',
21+
content: [
22+
{ type: 'text', text: 'Let me check the weather.' },
23+
{
24+
type: 'tool_use',
25+
id: 'tool_weather_1',
26+
name: 'weather',
27+
input: { city: 'Paris' },
28+
},
29+
{ type: 'text', text: 'It is sunny.' },
30+
],
31+
stop_reason: 'end_turn',
32+
stop_sequence: null,
33+
usage: {
34+
input_tokens: 5,
35+
output_tokens: 7,
36+
},
37+
};
38+
}
39+
}
40+
41+
async function run() {
42+
await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
43+
const mockClient = new MockAnthropic({ apiKey: 'mock-api-key' });
44+
const client = instrumentAnthropicAiClient(mockClient);
45+
46+
await client.messages.create({
47+
model: 'claude-3-haiku-20240307',
48+
messages: [{ role: 'user', content: 'What is the weather?' }],
49+
tools: [
50+
{
51+
name: 'weather',
52+
description: 'Get the weather by city',
53+
input_schema: {
54+
type: 'object',
55+
properties: { city: { type: 'string' } },
56+
required: ['city'],
57+
},
58+
},
59+
],
60+
});
61+
});
62+
}
63+
64+
run();
65+
66+

‎dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,4 +293,59 @@ describe('Anthropic integration', () => {
293293
await createRunner().ignore('event').expect({ transaction: EXPECTED_STREAM_SPANS_PII_TRUE }).start().completed();
294294
});
295295
});
296+
297+
// Non-streaming tool calls + available tools (PII true)
298+
createEsmAndCjsTests(__dirname, 'scenario-tools.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
299+
test('non-streaming sets available tools and tool calls with PII', async () => {
300+
const EXPECTED_TOOLS_JSON =
301+
'[{"name":"weather","description":"Get the weather by city","input_schema":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}]';
302+
const EXPECTED_TOOL_CALLS_JSON =
303+
'[{"type":"tool_use","id":"tool_weather_1","name":"weather","input":{"city":"Paris"}}]';
304+
await createRunner()
305+
.ignore('event')
306+
.expect({
307+
transaction: {
308+
spans: expect.arrayContaining([
309+
expect.objectContaining({
310+
op: 'gen_ai.messages',
311+
data: expect.objectContaining({
312+
'gen_ai.request.available_tools': EXPECTED_TOOLS_JSON,
313+
'gen_ai.response.tool_calls': EXPECTED_TOOL_CALLS_JSON,
314+
}),
315+
}),
316+
]),
317+
},
318+
})
319+
.start()
320+
.completed();
321+
});
322+
});
323+
324+
// Streaming tool calls + available tools (PII true)
325+
createEsmAndCjsTests(__dirname, 'scenario-stream-tools.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
326+
test('streaming sets available tools and tool calls with PII', async () => {
327+
const EXPECTED_TOOLS_JSON =
328+
'[{"name":"weather","description":"Get weather","input_schema":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}]';
329+
const EXPECTED_TOOL_CALLS_JSON =
330+
'[{"type":"tool_use","id":"tool_weather_2","name":"weather","input":{"city":"Paris"}}]';
331+
await createRunner()
332+
.ignore('event')
333+
.expect({
334+
transaction: {
335+
spans: expect.arrayContaining([
336+
expect.objectContaining({
337+
description: expect.stringContaining('stream-response'),
338+
op: 'gen_ai.messages',
339+
data: expect.objectContaining({
340+
'gen_ai.request.available_tools': EXPECTED_TOOLS_JSON,
341+
'gen_ai.response.tool_calls': EXPECTED_TOOL_CALLS_JSON,
342+
}),
343+
}),
344+
]),
345+
},
346+
})
347+
.start()
348+
.completed();
349+
});
350+
});
296351
});

0 commit comments

Comments
 (0)