Skip to content

Commit eaa3c56

Browse files
isaacsisaacs
committed
fix(tracing): add system prompt, model to google genai (#18424)
Move the message reformatting into a separate util for google-genai, and add unit test coverage for that file. Add an integration test scenario to ensure that the system message will be included if provided in the config params. Related to getsentry/testing-ai-sdk-integrations#10 Fix JS-1218
1 parent 7752d60 commit eaa3c56

File tree

5 files changed

+170
-22
lines changed

5 files changed

+170
-22
lines changed

‎dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ async function run() {
6161
temperature: 0.8,
6262
topP: 0.9,
6363
maxOutputTokens: 150,
64+
systemInstruction: 'You are a friendly robot who likes to be funny.',
6465
},
6566
history: [
6667
{

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,9 @@ describe('Google GenAI integration', () => {
9494
'gen_ai.request.temperature': 0.8,
9595
'gen_ai.request.top_p': 0.9,
9696
'gen_ai.request.max_tokens': 150,
97-
'gen_ai.request.messages': expect.any(String), // Should include history when recordInputs: true
97+
'gen_ai.request.messages': expect.stringMatching(
98+
/\[\{"role":"system","content":"You are a friendly robot who likes to be funny."\},/,
99+
), // Should include history when recordInputs: true
98100
}),
99101
description: 'chat gemini-1.5-pro create',
100102
op: 'gen_ai.chat',

‎packages/core/src/tracing/google-genai/index.ts‎

Lines changed: 35 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,16 @@ import {
1616
GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE,
1717
GEN_AI_REQUEST_TOP_K_ATTRIBUTE,
1818
GEN_AI_REQUEST_TOP_P_ATTRIBUTE,
19+
GEN_AI_RESPONSE_MODEL_ATTRIBUTE,
1920
GEN_AI_RESPONSE_TEXT_ATTRIBUTE,
2021
GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE,
2122
GEN_AI_SYSTEM_ATTRIBUTE,
2223
GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE,
2324
GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
2425
GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
2526
} from '../ai/gen-ai-attributes';
26-
import { buildMethodPath, getFinalOperationName, getSpanOperation, getTruncatedJsonString } from '../ai/utils';
27+
import { truncateGenAiMessages } from '../ai/messageTruncation';
28+
import { buildMethodPath, getFinalOperationName, getSpanOperation } from '../ai/utils';
2729
import { CHAT_PATH, CHATS_CREATE_METHOD, GOOGLE_GENAI_SYSTEM_NAME } from './constants';
2830
import { instrumentStream } from './streaming';
2931
import type {
@@ -33,7 +35,8 @@ import type {
3335
GoogleGenAIOptions,
3436
GoogleGenAIResponse,
3537
} from './types';
36-
import { isStreamingMethod, shouldInstrument } from './utils';
38+
import type { ContentListUnion, ContentUnion, Message, PartListUnion } from './utils';
39+
import { contentUnionToMessages, isStreamingMethod, shouldInstrument } from './utils';
3740

3841
/**
3942
* Extract model from parameters or chat context object
@@ -134,26 +137,38 @@ function extractRequestAttributes(
134137
* Handles different parameter formats for different Google GenAI methods.
135138
*/
136139
function addPrivateRequestAttributes(span: Span, params: Record<string, unknown>): void {
137-
// For models.generateContent: ContentListUnion: Content | Content[] | PartUnion | PartUnion[]
140+
const messages: Message[] = [];
141+
142+
// config.systemInstruction: ContentUnion
143+
if (
144+
'config' in params &&
145+
params.config &&
146+
typeof params.config === 'object' &&
147+
'systemInstruction' in params.config &&
148+
params.config.systemInstruction
149+
) {
150+
messages.push(...contentUnionToMessages(params.config.systemInstruction as ContentUnion, 'system'));
151+
}
152+
153+
// For chats.create: history contains the conversation history
154+
if ('history' in params) {
155+
messages.push(...contentUnionToMessages(params.history as PartListUnion, 'user'));
156+
}
157+
158+
// For models.generateContent: ContentListUnion
138159
if ('contents' in params) {
139-
const contents = params.contents;
140-
// For models.generateContent: ContentListUnion: Content | Content[] | PartUnion | PartUnion[]
141-
const truncatedContents = getTruncatedJsonString(contents);
142-
span.setAttributes({ [GEN_AI_REQUEST_MESSAGES_ATTRIBUTE]: truncatedContents });
160+
messages.push(...contentUnionToMessages(params.contents as ContentListUnion, 'user'));
143161
}
144162

145-
// For chat.sendMessage: message can be string or Part[]
163+
// For chat.sendMessage: message can be PartListUnion
146164
if ('message' in params) {
147-
const message = params.message;
148-
const truncatedMessage = getTruncatedJsonString(message);
149-
span.setAttributes({ [GEN_AI_REQUEST_MESSAGES_ATTRIBUTE]: truncatedMessage });
165+
messages.push(...contentUnionToMessages(params.message as PartListUnion, 'user'));
150166
}
151167

152-
// For chats.create: history contains the conversation history
153-
if ('history' in params) {
154-
const history = params.history;
155-
const truncatedHistory = getTruncatedJsonString(history);
156-
span.setAttributes({ [GEN_AI_REQUEST_MESSAGES_ATTRIBUTE]: truncatedHistory });
168+
if (messages.length) {
169+
span.setAttributes({
170+
[GEN_AI_REQUEST_MESSAGES_ATTRIBUTE]: JSON.stringify(truncateGenAiMessages(messages)),
171+
});
157172
}
158173
}
159174

@@ -164,6 +179,10 @@ function addPrivateRequestAttributes(span: Span, params: Record<string, unknown>
164179
function addResponseAttributes(span: Span, response: GoogleGenAIResponse, recordOutputs?: boolean): void {
165180
if (!response || typeof response !== 'object') return;
166181

182+
if (response.modelVersion) {
183+
span.setAttribute(GEN_AI_RESPONSE_MODEL_ATTRIBUTE, response.modelVersion);
184+
}
185+
167186
// Add usage metadata if present
168187
if (response.usageMetadata && typeof response.usageMetadata === 'object') {
169188
const usage = response.usageMetadata;

‎packages/core/src/tracing/google-genai/utils.ts‎

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,50 @@ export function shouldInstrument(methodPath: string): methodPath is GoogleGenAII
1919
* Check if a method is a streaming method
2020
*/
2121
export function isStreamingMethod(methodPath: string): boolean {
22-
return (
23-
methodPath.includes('Stream') ||
24-
methodPath.endsWith('generateContentStream') ||
25-
methodPath.endsWith('sendMessageStream')
26-
);
22+
return methodPath.includes('Stream');
23+
}
24+
25+
// Copied from https://googleapis.github.io/js-genai/release_docs/index.html
26+
export type ContentListUnion = Content | Content[] | PartListUnion;
27+
export type ContentUnion = Content | PartUnion[] | PartUnion;
28+
export type Content = {
29+
parts?: Part[];
30+
role?: string;
31+
};
32+
export type PartUnion = Part | string;
33+
export type Part = Record<string, unknown> & {
34+
inlineData?: {
35+
data?: string;
36+
displayName?: string;
37+
mimeType?: string;
38+
};
39+
text?: string;
40+
};
41+
export type PartListUnion = PartUnion[] | PartUnion;
42+
43+
// our consistent span message shape
44+
export type Message = Record<string, unknown> & {
45+
role: string;
46+
content?: PartListUnion;
47+
parts?: PartListUnion;
48+
};
49+
50+
/**
51+
*
52+
*/
53+
export function contentUnionToMessages(content: ContentListUnion, role = 'user'): Message[] {
54+
if (typeof content === 'string') {
55+
return [{ role, content }];
56+
}
57+
if (Array.isArray(content)) {
58+
return content.flatMap(content => contentUnionToMessages(content, role));
59+
}
60+
if (typeof content !== 'object' || !content) return [];
61+
if ('role' in content && typeof content.role === 'string') {
62+
return [content as Message];
63+
}
64+
if ('parts' in content) {
65+
return [{ ...content, role } as Message];
66+
}
67+
return [{ role, content }];
2768
}

‎packages/core/test/lib/utils/google-genai-utils.test.ts‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import { describe, expect, it } from 'vitest';
2+
import type { ContentListUnion } from '../../../src/tracing/google-genai/utils';
3+
import { contentUnionToMessages, isStreamingMethod, shouldInstrument } from '../../../src/tracing/google-genai/utils';
4+
5+
describe('isStreamingMethod', () => {
6+
it('detects streaming methods', () => {
7+
expect(isStreamingMethod('messageStreamBlah')).toBe(true);
8+
expect(isStreamingMethod('blahblahblah generateContentStream')).toBe(true);
9+
expect(isStreamingMethod('blahblahblah sendMessageStream')).toBe(true);
10+
expect(isStreamingMethod('blahblahblah generateContentStream')).toBe(true);
11+
expect(isStreamingMethod('blahblahblah sendMessageStream')).toBe(true);
12+
expect(isStreamingMethod('blahblahblah generateContent')).toBe(false);
13+
expect(isStreamingMethod('blahblahblah sendMessage')).toBe(false);
14+
});
15+
});
16+
17+
describe('shouldInstrument', () => {
18+
it('detects which methods to instrument', () => {
19+
expect(shouldInstrument('models.generateContent')).toBe(true);
20+
expect(shouldInstrument('some.path.to.sendMessage')).toBe(true);
21+
expect(shouldInstrument('unknown')).toBe(false);
22+
});
23+
});
24+
25+
describe('convert google-genai messages to consistent message', () => {
26+
it('converts strings to messages', () => {
27+
expect(contentUnionToMessages('hello', 'system')).toStrictEqual([{ role: 'system', content: 'hello' }]);
28+
expect(contentUnionToMessages('hello')).toStrictEqual([{ role: 'user', content: 'hello' }]);
29+
});
30+
31+
it('converts arrays of strings to messages', () => {
32+
expect(contentUnionToMessages(['hello', 'goodbye'], 'system')).toStrictEqual([
33+
{ role: 'system', content: 'hello' },
34+
{ role: 'system', content: 'goodbye' },
35+
]);
36+
expect(contentUnionToMessages(['hello', 'goodbye'])).toStrictEqual([
37+
{ role: 'user', content: 'hello' },
38+
{ role: 'user', content: 'goodbye' },
39+
]);
40+
});
41+
42+
it('converts PartUnion to messages', () => {
43+
expect(contentUnionToMessages(['hello', { parts: ['i am here', { text: 'goodbye' }] }], 'system')).toStrictEqual([
44+
{ role: 'system', content: 'hello' },
45+
{ role: 'system', parts: ['i am here', { text: 'goodbye' }] },
46+
]);
47+
48+
expect(contentUnionToMessages(['hello', { parts: ['i am here', { text: 'goodbye' }] }])).toStrictEqual([
49+
{ role: 'user', content: 'hello' },
50+
{ role: 'user', parts: ['i am here', { text: 'goodbye' }] },
51+
]);
52+
});
53+
54+
it('converts ContentUnion to messages', () => {
55+
expect(
56+
contentUnionToMessages(
57+
{
58+
parts: ['hello', 'goodbye'],
59+
role: 'agent',
60+
},
61+
'user',
62+
),
63+
).toStrictEqual([{ parts: ['hello', 'goodbye'], role: 'agent' }]);
64+
});
65+
66+
it('handles unexpected formats safely', () => {
67+
expect(
68+
contentUnionToMessages(
69+
[
70+
{
71+
parts: ['hello', 'goodbye'],
72+
role: 'agent',
73+
},
74+
null,
75+
21345,
76+
{ data: 'this is content' },
77+
] as unknown as ContentListUnion,
78+
'user',
79+
),
80+
).toStrictEqual([
81+
{ parts: ['hello', 'goodbye'], role: 'agent' },
82+
{ role: 'user', content: { data: 'this is content' } },
83+
]);
84+
});
85+
});

0 commit comments

Comments
 (0)