diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-streaming.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-streaming.mjs new file mode 100644 index 000000000000..be5c75638694 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-streaming.mjs @@ -0,0 +1,237 @@ +import { GoogleGenAI } from '@google/genai'; +import * as Sentry from '@sentry/node'; +import express from 'express'; + +function startMockGoogleGenAIServer() { + const app = express(); + app.use(express.json()); + + // Streaming endpoint for models.generateContentStream and chat.sendMessageStream + app.post('/v1beta/models/:model\\:streamGenerateContent', (req, res) => { + const model = req.params.model; + + if (model === 'error-model') { + res.status(404).set('x-request-id', 'mock-request-123').end('Model not found'); + return; + } + + // Set headers for streaming response + res.setHeader('Content-Type', 'application/json'); + res.setHeader('Transfer-Encoding', 'chunked'); + + // Create a mock stream + const mockStream = createMockStream(model); + + // Send chunks + const sendChunk = async () => { + const { value, done } = await mockStream.next(); + if (done) { + res.end(); + return; + } + + res.write(`data: ${JSON.stringify(value)}\n\n`); + setTimeout(sendChunk, 10); // Small delay between chunks + }; + + sendChunk(); + }); + + return new Promise(resolve => { + const server = app.listen(0, () => { + resolve(server); + }); + }); +} + +// Helper function to create mock stream +async function* createMockStream(model) { + if (model === 'blocked-model') { + // First chunk: Contains promptFeedback with blockReason + yield { + promptFeedback: { + blockReason: 'SAFETY', + blockReasonMessage: 'The prompt was blocked due to safety concerns', + }, + responseId: 'mock-blocked-response-streaming-id', + modelVersion: 'gemini-1.5-pro', + }; + + // Note: In a real blocked scenario, there would typically be no more chunks + // But we'll add one more to test that processing stops after the error + yield { + candidates: [ + { + content: { + parts: [{ text: 'This should not be processed' }], + role: 'model', + }, + index: 0, + }, + ], + }; + return; + } + + // First chunk: Start of response with initial text + yield { + candidates: [ + { + content: { + parts: [{ text: 'Hello! ' }], + role: 'model', + }, + index: 0, + }, + ], + responseId: 'mock-response-streaming-id', + modelVersion: 'gemini-1.5-pro', + }; + + // Second chunk: More text content + yield { + candidates: [ + { + content: { + parts: [{ text: 'This is a streaming ' }], + role: 'model', + }, + index: 0, + }, + ], + }; + + // Third chunk: Final text content + yield { + candidates: [ + { + content: { + parts: [{ text: 'response from Google GenAI!' }], + role: 'model', + }, + index: 0, + }, + ], + }; + + // Final chunk: End with finish reason and usage metadata + yield { + candidates: [ + { + content: { + parts: [{ text: '' }], // Empty text in final chunk + role: 'model', + }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 12, + totalTokenCount: 22, + }, + }; +} + +async function run() { + const server = await startMockGoogleGenAIServer(); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const client = new GoogleGenAI({ + apiKey: 'mock-api-key', + httpOptions: { baseUrl: `http://localhost:${server.address().port}` }, + }); + + // Test 1: models.generateContentStream (streaming) + const streamResponse = await client.models.generateContentStream({ + model: 'gemini-1.5-flash', + config: { + temperature: 0.7, + topP: 0.9, + maxOutputTokens: 100, + }, + contents: [ + { + role: 'user', + parts: [{ text: 'Tell me about streaming' }], + }, + ], + }); + + // Consume the stream + for await (const _ of streamResponse) { + void _; + } + + // Test 2: chat.sendMessageStream (streaming) + const streamingChat = client.chats.create({ + model: 'gemini-1.5-pro', + config: { + temperature: 0.8, + topP: 0.9, + maxOutputTokens: 150, + }, + }); + + const chatStreamResponse = await streamingChat.sendMessageStream({ + message: 'Tell me a streaming joke', + }); + + // Consume the chat stream + for await (const _ of chatStreamResponse) { + void _; + } + + // Test 3: Blocked content streaming (should trigger error handling) + try { + const blockedStreamResponse = await client.models.generateContentStream({ + model: 'blocked-model', + config: { + temperature: 0.7, + }, + contents: [ + { + role: 'user', + parts: [{ text: 'This should be blocked' }], + }, + ], + }); + + // Consume the blocked stream + for await (const _ of blockedStreamResponse) { + void _; + } + } catch { + // Expected: The stream should be processed, but the span should be marked with error status + // The error handling happens in the streaming instrumentation, not as a thrown error + } + + // Test 4: Error handling for streaming + try { + const errorStreamResponse = await client.models.generateContentStream({ + model: 'error-model', + config: { + temperature: 0.7, + }, + contents: [ + { + role: 'user', + parts: [{ text: 'This will fail' }], + }, + ], + }); + + // Consume the error stream + for await (const _ of errorStreamResponse) { + void _; + } + } catch { + // Expected error + } + }); + + server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-tools.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-tools.mjs new file mode 100644 index 000000000000..97984f2eb1ed --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-tools.mjs @@ -0,0 +1,307 @@ +import { GoogleGenAI } from '@google/genai'; +import * as Sentry from '@sentry/node'; +import express from 'express'; + +function startMockGoogleGenAIServer() { + const app = express(); + app.use(express.json()); + + // Non-streaming endpoint for models.generateContent + app.post('/v1beta/models/:model\\:generateContent', (req, res) => { + const { tools } = req.body; + + // Check if tools are provided to return function call response + if (tools && tools.length > 0) { + const response = { + candidates: [ + { + content: { + parts: [ + { + text: 'I need to check the light status first.', + }, + { + functionCall: { + id: 'call_light_control_1', + name: 'controlLight', + args: { + brightness: 0.3, + colorTemperature: 'warm', + }, + }, + }, + ], + role: 'model', + }, + finishReason: 'stop', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 15, + candidatesTokenCount: 8, + totalTokenCount: 23, + }, + }; + + // Add functionCalls getter, this should exist in the response object + Object.defineProperty(response, 'functionCalls', { + get: function () { + return [ + { + id: 'call_light_control_1', + name: 'controlLight', + args: { + brightness: 0.3, + colorTemperature: 'warm', + }, + }, + ]; + }, + }); + + res.send(response); + return; + } + + // Regular response without tools + res.send({ + candidates: [ + { + content: { + parts: [ + { + text: 'Mock response from Google GenAI without tools!', + }, + ], + role: 'model', + }, + finishReason: 'stop', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 8, + candidatesTokenCount: 12, + totalTokenCount: 20, + }, + }); + }); + + // Streaming endpoint for models.generateContentStream + // And chat.sendMessageStream + app.post('/v1beta/models/:model\\:streamGenerateContent', (req, res) => { + const { tools } = req.body; + + // Set headers for streaming response + res.setHeader('Content-Type', 'application/json'); + res.setHeader('Transfer-Encoding', 'chunked'); + + // Create a mock stream + const mockStream = createMockToolsStream({ tools }); + + // Send chunks + const sendChunk = async () => { + // Testing .next() works as expected + const { value, done } = await mockStream.next(); + if (done) { + res.end(); + return; + } + + res.write(`data: ${JSON.stringify(value)}\n\n`); + setTimeout(sendChunk, 10); // Small delay between chunks + }; + + sendChunk(); + }); + + return new Promise(resolve => { + const server = app.listen(0, () => { + resolve(server); + }); + }); +} + +// Helper function to create mock stream +async function* createMockToolsStream({ tools }) { + // Check if tools are provided to return function call response + if (tools && tools.length > 0) { + // First chunk: Text response + yield { + candidates: [ + { + content: { + parts: [{ text: 'Let me control the lights for you.' }], + role: 'model', + }, + index: 0, + }, + ], + responseId: 'mock-response-tools-id', + modelVersion: 'gemini-2.0-flash-001', + }; + + // Second chunk: Function call + yield { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_light_stream_1', + name: 'controlLight', + args: { + brightness: 0.5, + colorTemperature: 'cool', + }, + }, + }, + ], + role: 'model', + }, + index: 0, + }, + ], + }; + + // Final chunk: End with finish reason and usage metadata + yield { + candidates: [ + { + content: { + parts: [{ text: ' Done!' }], // Additional text in final chunk + role: 'model', + }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 12, + candidatesTokenCount: 10, + totalTokenCount: 22, + }, + }; + return; + } + + // Regular stream without tools + yield { + candidates: [ + { + content: { + parts: [{ text: 'Mock streaming response' }], + role: 'model', + }, + index: 0, + }, + ], + responseId: 'mock-response-tools-id', + modelVersion: 'gemini-2.0-flash-001', + }; + + // Final chunk + yield { + candidates: [ + { + content: { + parts: [{ text: ' from Google GenAI!' }], + role: 'model', + }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 12, + totalTokenCount: 22, + }, + }; +} + +async function run() { + const server = await startMockGoogleGenAIServer(); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const client = new GoogleGenAI({ + apiKey: 'mock-api-key', + httpOptions: { baseUrl: `http://localhost:${server.address().port}` }, + }); + + // Test 1: Non-streaming with tools + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Dim the lights so the room feels cozy and warm.', + config: { + tools: [ + { + functionDeclarations: [ + { + name: 'controlLight', + parametersJsonSchema: { + type: 'object', + properties: { + brightness: { + type: 'number', + }, + colorTemperature: { + type: 'string', + }, + }, + required: ['brightness', 'colorTemperature'], + }, + }, + ], + }, + ], + }, + }); + + // Test 2: Streaming with tools + const stream = await client.models.generateContentStream({ + model: 'gemini-2.0-flash-001', + contents: 'Turn on the lights with medium brightness.', + config: { + tools: [ + { + functionDeclarations: [ + { + name: 'controlLight', + parametersJsonSchema: { + type: 'object', + properties: { + brightness: { + type: 'number', + }, + colorTemperature: { + type: 'string', + }, + }, + required: ['brightness', 'colorTemperature'], + }, + }, + ], + }, + ], + }, + }); + + // Consume the stream to trigger instrumentation + for await (const _ of stream) { + void _; + } + + // Test 3: Without tools for comparison + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Tell me about the weather.', + }); + }); + + server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs index ddb9e16b8254..91c75886e410 100644 --- a/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs @@ -54,6 +54,7 @@ async function run() { }); // Test 1: chats.create and sendMessage flow + // This should generate two spans: one for chats.create and one for sendMessage const chat = client.chats.create({ model: 'gemini-1.5-pro', config: { diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts b/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts index 9aa5523c61d7..92d669c7e10f 100644 --- a/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts @@ -169,6 +169,7 @@ describe('Google GenAI integration', () => { 'gen_ai.request.messages': expect.any(String), // Should include messages when recordInputs: true 'gen_ai.response.text': expect.any(String), // Should include response text when recordOutputs: true }), + description: expect.not.stringContaining('stream-response'), // Non-streaming span }), ]), }; @@ -202,4 +203,287 @@ describe('Google GenAI integration', () => { .completed(); }); }); + + const EXPECTED_TRANSACTION_TOOLS = { + transaction: 'main', + spans: expect.arrayContaining([ + // Non-streaming with tools + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-2.0-flash-001', + 'gen_ai.request.available_tools': expect.any(String), // Should include tools + 'gen_ai.request.messages': expect.any(String), // Should include contents + 'gen_ai.response.text': expect.any(String), // Should include response text + 'gen_ai.response.tool_calls': expect.any(String), // Should include tool calls + 'gen_ai.usage.input_tokens': 15, + 'gen_ai.usage.output_tokens': 8, + 'gen_ai.usage.total_tokens': 23, + }), + description: 'models gemini-2.0-flash-001', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Streaming with tools + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-2.0-flash-001', + 'gen_ai.request.available_tools': expect.any(String), // Should include tools + 'gen_ai.request.messages': expect.any(String), // Should include contents + 'gen_ai.response.streaming': true, + 'gen_ai.response.text': expect.any(String), // Should include response text + 'gen_ai.response.tool_calls': expect.any(String), // Should include tool calls + 'gen_ai.response.id': 'mock-response-tools-id', + 'gen_ai.response.model': 'gemini-2.0-flash-001', + 'gen_ai.usage.input_tokens': 12, + 'gen_ai.usage.output_tokens': 10, + 'gen_ai.usage.total_tokens': 22, + }), + description: 'models gemini-2.0-flash-001 stream-response', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Without tools for comparison + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-2.0-flash-001', + 'gen_ai.request.messages': expect.any(String), // Should include contents + 'gen_ai.response.text': expect.any(String), // Should include response text + 'gen_ai.usage.input_tokens': 8, + 'gen_ai.usage.output_tokens': 12, + 'gen_ai.usage.total_tokens': 20, + }), + description: 'models gemini-2.0-flash-001', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + ]), + }; + + createEsmAndCjsTests(__dirname, 'scenario-tools.mjs', 'instrument-with-options.mjs', (createRunner, test) => { + test('creates google genai related spans with tool calls', async () => { + await createRunner().ignore('event').expect({ transaction: EXPECTED_TRANSACTION_TOOLS }).start().completed(); + }); + }); + + const EXPECTED_TRANSACTION_STREAMING = { + transaction: 'main', + spans: expect.arrayContaining([ + // First span - models.generateContentStream (streaming) + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-1.5-flash', + 'gen_ai.request.temperature': 0.7, + 'gen_ai.request.top_p': 0.9, + 'gen_ai.request.max_tokens': 100, + 'gen_ai.response.streaming': true, + 'gen_ai.response.id': 'mock-response-streaming-id', + 'gen_ai.response.model': 'gemini-1.5-pro', + 'gen_ai.response.finish_reasons': '["STOP"]', + 'gen_ai.usage.input_tokens': 10, + 'gen_ai.usage.output_tokens': 12, + 'gen_ai.usage.total_tokens': 22, + }), + description: 'models gemini-1.5-flash stream-response', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Second span - chat.create + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'chat', + 'sentry.op': 'gen_ai.chat', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-1.5-pro', + 'gen_ai.request.temperature': 0.8, + 'gen_ai.request.top_p': 0.9, + 'gen_ai.request.max_tokens': 150, + }), + description: 'chat gemini-1.5-pro create', + op: 'gen_ai.chat', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Third span - chat.sendMessageStream (streaming) + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'chat', + 'sentry.op': 'gen_ai.chat', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-1.5-pro', + 'gen_ai.response.streaming': true, + 'gen_ai.response.id': 'mock-response-streaming-id', + 'gen_ai.response.model': 'gemini-1.5-pro', + }), + description: 'chat gemini-1.5-pro stream-response', + op: 'gen_ai.chat', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Fourth span - blocked content streaming + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + }), + description: 'models blocked-model stream-response', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'unknown_error', + }), + // Fifth span - error handling for streaming + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + }), + description: 'models error-model stream-response', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'internal_error', + }), + ]), + }; + + const EXPECTED_TRANSACTION_STREAMING_PII_TRUE = { + transaction: 'main', + spans: expect.arrayContaining([ + // First span - models.generateContentStream (streaming) with PII + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-1.5-flash', + 'gen_ai.request.temperature': 0.7, + 'gen_ai.request.top_p': 0.9, + 'gen_ai.request.max_tokens': 100, + 'gen_ai.request.messages': expect.any(String), // Should include contents when recordInputs: true + 'gen_ai.response.streaming': true, + 'gen_ai.response.id': 'mock-response-streaming-id', + 'gen_ai.response.model': 'gemini-1.5-pro', + 'gen_ai.response.finish_reasons': '["STOP"]', + 'gen_ai.usage.input_tokens': 10, + 'gen_ai.usage.output_tokens': 12, + 'gen_ai.usage.total_tokens': 22, + }), + description: 'models gemini-1.5-flash stream-response', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Second span - chat.create + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'chat', + 'sentry.op': 'gen_ai.chat', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-1.5-pro', + 'gen_ai.request.temperature': 0.8, + 'gen_ai.request.top_p': 0.9, + 'gen_ai.request.max_tokens': 150, + }), + description: 'chat gemini-1.5-pro create', + op: 'gen_ai.chat', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Third span - chat.sendMessageStream (streaming) with PII + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'chat', + 'sentry.op': 'gen_ai.chat', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'gemini-1.5-pro', + 'gen_ai.request.messages': expect.any(String), // Should include message when recordInputs: true + 'gen_ai.response.streaming': true, + 'gen_ai.response.id': 'mock-response-streaming-id', + 'gen_ai.response.model': 'gemini-1.5-pro', + 'gen_ai.response.finish_reasons': '["STOP"]', + 'gen_ai.usage.input_tokens': 10, + 'gen_ai.usage.output_tokens': 12, + 'gen_ai.usage.total_tokens': 22, + }), + description: 'chat gemini-1.5-pro stream-response', + op: 'gen_ai.chat', + origin: 'auto.ai.google_genai', + status: 'ok', + }), + // Fourth span - blocked content stream with PII + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'blocked-model', + 'gen_ai.request.temperature': 0.7, + 'gen_ai.request.messages': expect.any(String), // Should include contents when recordInputs: true + 'gen_ai.response.streaming': true, + }), + description: 'models blocked-model stream-response', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'unknown_error', + }), + // Fifth span - error handling for streaming with PII + expect.objectContaining({ + data: expect.objectContaining({ + 'gen_ai.operation.name': 'models', + 'sentry.op': 'gen_ai.models', + 'sentry.origin': 'auto.ai.google_genai', + 'gen_ai.system': 'google_genai', + 'gen_ai.request.model': 'error-model', + 'gen_ai.request.temperature': 0.7, + 'gen_ai.request.messages': expect.any(String), // Should include contents when recordInputs: true + }), + description: 'models error-model stream-response', + op: 'gen_ai.models', + origin: 'auto.ai.google_genai', + status: 'internal_error', + }), + ]), + }; + + createEsmAndCjsTests(__dirname, 'scenario-streaming.mjs', 'instrument.mjs', (createRunner, test) => { + test('creates google genai streaming spans with sendDefaultPii: false', async () => { + await createRunner().ignore('event').expect({ transaction: EXPECTED_TRANSACTION_STREAMING }).start().completed(); + }); + }); + + createEsmAndCjsTests(__dirname, 'scenario-streaming.mjs', 'instrument-with-pii.mjs', (createRunner, test) => { + test('creates google genai streaming spans with sendDefaultPii: true', async () => { + await createRunner() + .ignore('event') + .expect({ transaction: EXPECTED_TRANSACTION_STREAMING_PII_TRUE }) + .start() + .completed(); + }); + }); }); diff --git a/packages/core/src/utils/google-genai/constants.ts b/packages/core/src/utils/google-genai/constants.ts index 8617460482c6..b06e46e18755 100644 --- a/packages/core/src/utils/google-genai/constants.ts +++ b/packages/core/src/utils/google-genai/constants.ts @@ -2,7 +2,15 @@ export const GOOGLE_GENAI_INTEGRATION_NAME = 'Google_GenAI'; // https://ai.google.dev/api/rest/v1/models/generateContent // https://ai.google.dev/api/rest/v1/chats/sendMessage -export const GOOGLE_GENAI_INSTRUMENTED_METHODS = ['models.generateContent', 'chats.create', 'sendMessage'] as const; +// https://googleapis.github.io/js-genai/release_docs/classes/models.Models.html#generatecontentstream +// https://googleapis.github.io/js-genai/release_docs/classes/chats.Chat.html#sendmessagestream +export const GOOGLE_GENAI_INSTRUMENTED_METHODS = [ + 'models.generateContent', + 'models.generateContentStream', + 'chats.create', + 'sendMessage', + 'sendMessageStream', +] as const; // Constants for internal use export const GOOGLE_GENAI_SYSTEM_NAME = 'google_genai'; diff --git a/packages/core/src/utils/google-genai/index.ts b/packages/core/src/utils/google-genai/index.ts index 58d7e2e6b5e6..20e6e2a53606 100644 --- a/packages/core/src/utils/google-genai/index.ts +++ b/packages/core/src/utils/google-genai/index.ts @@ -1,10 +1,12 @@ import { getClient } from '../../currentScopes'; import { captureException } from '../../exports'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes'; -import { startSpan } from '../../tracing/trace'; +import { SPAN_STATUS_ERROR } from '../../tracing'; +import { startSpan, startSpanManual } from '../../tracing/trace'; import type { Span, SpanAttributeValue } from '../../types-hoist/span'; import { GEN_AI_OPERATION_NAME_ATTRIBUTE, + GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE, GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE, GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE, GEN_AI_REQUEST_MESSAGES_ATTRIBUTE, @@ -14,6 +16,7 @@ import { GEN_AI_REQUEST_TOP_K_ATTRIBUTE, GEN_AI_REQUEST_TOP_P_ATTRIBUTE, GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, GEN_AI_SYSTEM_ATTRIBUTE, GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, @@ -22,6 +25,7 @@ import { import { buildMethodPath, getFinalOperationName, getSpanOperation } from '../ai/utils'; import { handleCallbackErrors } from '../handleCallbackErrors'; import { CHAT_PATH, CHATS_CREATE_METHOD, GOOGLE_GENAI_SYSTEM_NAME } from './constants'; +import { instrumentStream } from './streaming'; import type { Candidate, ContentPart, @@ -29,7 +33,7 @@ import type { GoogleGenAIOptions, GoogleGenAIResponse, } from './types'; -import { shouldInstrument } from './utils'; +import { isStreamingMethod, shouldInstrument } from './utils'; /** * Extract model from parameters or chat context object @@ -91,8 +95,8 @@ function extractConfigAttributes(config: Record): Record, context?: unknown, ): Record { const attributes: Record = { @@ -101,14 +105,21 @@ function extractRequestAttributes( [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai', }; - if (args.length > 0 && typeof args[0] === 'object' && args[0] !== null) { - const params = args[0] as Record; - + if (params) { attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = extractModel(params, context); // Extract generation config parameters if ('config' in params && typeof params.config === 'object' && params.config) { - Object.assign(attributes, extractConfigAttributes(params.config as Record)); + const config = params.config as Record; + Object.assign(attributes, extractConfigAttributes(config)); + + // Extract available tools from config + if ('tools' in config && Array.isArray(config.tools)) { + const functionDeclarations = config.tools.map( + (tool: { functionDeclarations: unknown[] }) => tool.functionDeclarations, + ); + attributes[GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE] = JSON.stringify(functionDeclarations); + } } } else { attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = extractModel({}, context); @@ -186,6 +197,16 @@ function addResponseAttributes(span: Span, response: GoogleGenAIResponse, record }); } } + + // Add tool calls if recordOutputs is enabled + if (recordOutputs && response.functionCalls) { + const functionCalls = response.functionCalls; + if (Array.isArray(functionCalls) && functionCalls.length > 0) { + span.setAttributes({ + [GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]: JSON.stringify(functionCalls), + }); + } + } } /** @@ -201,43 +222,75 @@ function instrumentMethod( ): (...args: T) => R | Promise { const isSyncCreate = methodPath === CHATS_CREATE_METHOD; - const run = (...args: T): R | Promise => { - const requestAttributes = extractRequestAttributes(args, methodPath, context); - const model = requestAttributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] ?? 'unknown'; - const operationName = getFinalOperationName(methodPath); - - // Single span for both sync and async operations - return startSpan( - { - name: isSyncCreate ? `${operationName} ${model} create` : `${operationName} ${model}`, - op: getSpanOperation(methodPath), - attributes: requestAttributes, - }, - (span: Span) => { - if (options.recordInputs && args[0] && typeof args[0] === 'object') { - addPrivateRequestAttributes(span, args[0] as Record); - } - - return handleCallbackErrors( - () => originalMethod.apply(context, args), - error => { - captureException(error, { - mechanism: { handled: false, type: 'auto.ai.google_genai', data: { function: methodPath } }, - }); + return new Proxy(originalMethod, { + apply(target, _, args: T): R | Promise { + const params = args[0] as Record | undefined; + const requestAttributes = extractRequestAttributes(methodPath, params, context); + const model = requestAttributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] ?? 'unknown'; + const operationName = getFinalOperationName(methodPath); + + // Check if this is a streaming method + if (isStreamingMethod(methodPath)) { + // Use startSpanManual for streaming methods to control span lifecycle + return startSpanManual( + { + name: `${operationName} ${model} stream-response`, + op: getSpanOperation(methodPath), + attributes: requestAttributes, }, - () => {}, - result => { - // Only add response attributes for content-producing methods, not for chats.create - if (!isSyncCreate) { - addResponseAttributes(span, result, options.recordOutputs); + async (span: Span) => { + try { + if (options.recordInputs && params) { + addPrivateRequestAttributes(span, params); + } + const stream = await target.apply(context, args); + return instrumentStream(stream, span, Boolean(options.recordOutputs)) as R; + } catch (error) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + captureException(error, { + mechanism: { + handled: false, + type: 'auto.ai.google_genai', + data: { function: methodPath }, + }, + }); + span.end(); + throw error; } }, ); - }, - ); - }; - - return run; + } + // Single span for both sync and async operations + return startSpan( + { + name: isSyncCreate ? `${operationName} ${model} create` : `${operationName} ${model}`, + op: getSpanOperation(methodPath), + attributes: requestAttributes, + }, + (span: Span) => { + if (options.recordInputs && params) { + addPrivateRequestAttributes(span, params); + } + + return handleCallbackErrors( + () => target.apply(context, args), + error => { + captureException(error, { + mechanism: { handled: false, type: 'auto.ai.google_genai', data: { function: methodPath } }, + }); + }, + () => {}, + result => { + // Only add response attributes for content-producing methods, not for chats.create + if (!isSyncCreate) { + addResponseAttributes(span, result, options.recordOutputs); + } + }, + ); + }, + ); + }, + }) as (...args: T) => R | Promise; } /** diff --git a/packages/core/src/utils/google-genai/streaming.ts b/packages/core/src/utils/google-genai/streaming.ts new file mode 100644 index 000000000000..b9462e8c90dd --- /dev/null +++ b/packages/core/src/utils/google-genai/streaming.ts @@ -0,0 +1,163 @@ +import { captureException } from '../../exports'; +import { SPAN_STATUS_ERROR } from '../../tracing'; +import type { Span, SpanAttributeValue } from '../../types-hoist/span'; +import { + GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE, + GEN_AI_RESPONSE_ID_ATTRIBUTE, + GEN_AI_RESPONSE_MODEL_ATTRIBUTE, + GEN_AI_RESPONSE_STREAMING_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, + GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, +} from '../ai/gen-ai-attributes'; +import type { GoogleGenAIResponse } from './types'; + +/** + * State object used to accumulate information from a stream of Google GenAI events. + */ +interface StreamingState { + /** Collected response text fragments (for output recording). */ + responseTexts: string[]; + /** Reasons for finishing the response, as reported by the API. */ + finishReasons: string[]; + /** The response ID. */ + responseId?: string; + /** The model name. */ + responseModel?: string; + /** Number of prompt/input tokens used. */ + promptTokens?: number; + /** Number of completion/output tokens used. */ + completionTokens?: number; + /** Number of total tokens used. */ + totalTokens?: number; + /** Accumulated tool calls (finalized) */ + toolCalls: Array>; +} + +/** + * Checks if a response chunk contains an error + * @param chunk - The response chunk to check + * @param span - The span to update if error is found + * @returns Whether an error occurred + */ +function isErrorChunk(chunk: GoogleGenAIResponse, span: Span): boolean { + const feedback = chunk?.promptFeedback; + if (feedback?.blockReason) { + const message = feedback.blockReasonMessage ?? feedback.blockReason; + span.setStatus({ code: SPAN_STATUS_ERROR, message: `Content blocked: ${message}` }); + captureException(`Content blocked: ${message}`, { + mechanism: { handled: false, type: 'auto.ai.google_genai' }, + }); + return true; + } + return false; +} + +/** + * Processes response metadata from a chunk + * @param chunk - The response chunk to process + * @param state - The state of the streaming process + */ +function handleResponseMetadata(chunk: GoogleGenAIResponse, state: StreamingState): void { + if (typeof chunk.responseId === 'string') state.responseId = chunk.responseId; + if (typeof chunk.modelVersion === 'string') state.responseModel = chunk.modelVersion; + + const usage = chunk.usageMetadata; + if (usage) { + if (typeof usage.promptTokenCount === 'number') state.promptTokens = usage.promptTokenCount; + if (typeof usage.candidatesTokenCount === 'number') state.completionTokens = usage.candidatesTokenCount; + if (typeof usage.totalTokenCount === 'number') state.totalTokens = usage.totalTokenCount; + } +} + +/** + * Processes candidate content from a response chunk + * @param chunk - The response chunk to process + * @param state - The state of the streaming process + * @param recordOutputs - Whether to record outputs + */ +function handleCandidateContent(chunk: GoogleGenAIResponse, state: StreamingState, recordOutputs: boolean): void { + if (Array.isArray(chunk.functionCalls)) { + state.toolCalls.push(...chunk.functionCalls); + } + + for (const candidate of chunk.candidates ?? []) { + if (candidate?.finishReason && !state.finishReasons.includes(candidate.finishReason)) { + state.finishReasons.push(candidate.finishReason); + } + + for (const part of candidate?.content?.parts ?? []) { + if (recordOutputs && part.text) state.responseTexts.push(part.text); + if (part.functionCall) { + state.toolCalls.push({ + type: 'function', + id: part.functionCall.id, + name: part.functionCall.name, + arguments: part.functionCall.args, + }); + } + } + } +} + +/** + * Processes a single chunk from the Google GenAI stream + * @param chunk - The chunk to process + * @param state - The state of the streaming process + * @param recordOutputs - Whether to record outputs + * @param span - The span to update + */ +function processChunk(chunk: GoogleGenAIResponse, state: StreamingState, recordOutputs: boolean, span: Span): void { + if (!chunk || isErrorChunk(chunk, span)) return; + handleResponseMetadata(chunk, state); + handleCandidateContent(chunk, state, recordOutputs); +} + +/** + * Instruments an async iterable stream of Google GenAI response chunks, updates the span with + * streaming attributes and (optionally) the aggregated output text, and yields + * each chunk from the input stream unchanged. + */ +export async function* instrumentStream( + stream: AsyncIterable, + span: Span, + recordOutputs: boolean, +): AsyncGenerator { + const state: StreamingState = { + responseTexts: [], + finishReasons: [], + toolCalls: [], + }; + + try { + for await (const chunk of stream) { + processChunk(chunk, state, recordOutputs, span); + yield chunk; + } + } finally { + const attrs: Record = { + [GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]: true, + }; + + if (state.responseId) attrs[GEN_AI_RESPONSE_ID_ATTRIBUTE] = state.responseId; + if (state.responseModel) attrs[GEN_AI_RESPONSE_MODEL_ATTRIBUTE] = state.responseModel; + if (state.promptTokens !== undefined) attrs[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = state.promptTokens; + if (state.completionTokens !== undefined) attrs[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = state.completionTokens; + if (state.totalTokens !== undefined) attrs[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE] = state.totalTokens; + + if (state.finishReasons.length) { + attrs[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE] = JSON.stringify(state.finishReasons); + } + if (recordOutputs && state.responseTexts.length) { + attrs[GEN_AI_RESPONSE_TEXT_ATTRIBUTE] = state.responseTexts.join(''); + } + if (recordOutputs && state.toolCalls.length) { + attrs[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE] = JSON.stringify(state.toolCalls); + } + + span.setAttributes(attrs); + span.end(); + } +} diff --git a/packages/core/src/utils/google-genai/utils.ts b/packages/core/src/utils/google-genai/utils.ts index c7a18477c7dd..a394ed64a1bb 100644 --- a/packages/core/src/utils/google-genai/utils.ts +++ b/packages/core/src/utils/google-genai/utils.ts @@ -14,3 +14,14 @@ export function shouldInstrument(methodPath: string): methodPath is GoogleGenAII const methodName = methodPath.split('.').pop(); return GOOGLE_GENAI_INSTRUMENTED_METHODS.includes(methodName as GoogleGenAIIstrumentedMethod); } + +/** + * Check if a method is a streaming method + */ +export function isStreamingMethod(methodPath: string): boolean { + return ( + methodPath.includes('Stream') || + methodPath.endsWith('generateContentStream') || + methodPath.endsWith('sendMessageStream') + ); +}