diff --git a/.changeset/keep-google-genai-stream-finish-reason.md b/.changeset/keep-google-genai-stream-finish-reason.md new file mode 100644 index 000000000..e40e602ee --- /dev/null +++ b/.changeset/keep-google-genai-stream-finish-reason.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(google-genai): Keep the finish reason of a streamed response that has no content diff --git a/js/src/instrumentation/plugins/google-genai-plugin.test.ts b/js/src/instrumentation/plugins/google-genai-plugin.test.ts index 4ef608163..177f1fbf7 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.test.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.test.ts @@ -542,6 +542,98 @@ describe("GoogleGenAIPlugin", () => { expect(JSON.stringify(span.log.mock.calls)).not.toContain("aGVsbG8="); expect(span.end).toHaveBeenCalledTimes(1); }); + + it("keeps the finish reason of a stream that ends without content", async () => { + plugin.enable(); + const handlers = subscribeSpy.mock.calls[1][0]; + const event: any = { + arguments: [ + { contents: "something blocked", model: "gemini-3.8-flash" }, + ], + }; + + async function* stream() { + yield { + candidates: [ + { + finishReason: "SAFETY", + safetyRatings: [ + { + category: "HARM_CATEGORY_DANGEROUS_CONTENT", + probability: "HIGH", + }, + ], + }, + ], + }; + yield { usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 } }; + } + + handlers.begin(event); + event.result = stream(); + handlers.resolve(event); + for await (const _chunk of event.result) { + // Consume the provider stream so the instrumentation finalizes it. + } + + const span = mockStartSpan.mock.results.at(-1)?.value as { + log: ReturnType; + }; + expect(span.log).toHaveBeenCalledWith( + expect.objectContaining({ + output: { + candidates: [ + { + finishReason: "SAFETY", + safetyRatings: [ + { + category: "HARM_CATEGORY_DANGEROUS_CONTENT", + probability: "HIGH", + }, + ], + }, + ], + usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 }, + }, + }), + ); + }); + + it("keeps the prompt feedback of a stream whose prompt was blocked", async () => { + plugin.enable(); + const handlers = subscribeSpy.mock.calls[1][0]; + const event: any = { + arguments: [ + { contents: "something blocked", model: "gemini-3.8-flash" }, + ], + }; + + async function* stream() { + yield { + promptFeedback: { blockReason: "SAFETY" }, + usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 }, + }; + } + + handlers.begin(event); + event.result = stream(); + handlers.resolve(event); + for await (const _chunk of event.result) { + // Consume the provider stream so the instrumentation finalizes it. + } + + const span = mockStartSpan.mock.results.at(-1)?.value as { + log: ReturnType; + }; + expect(span.log).toHaveBeenCalledWith( + expect.objectContaining({ + output: { + promptFeedback: { blockReason: "SAFETY" }, + usageMetadata: { promptTokenCount: 3, totalTokenCount: 3 }, + }, + }), + ); + }); }); describe("generateImages channel interception", () => { diff --git a/js/src/instrumentation/plugins/google-genai-plugin.ts b/js/src/instrumentation/plugins/google-genai-plugin.ts index d87bb6b0e..eeb4aebac 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.ts @@ -2001,19 +2001,22 @@ function aggregateGenerateContentChunks( const otherParts: Record[] = []; let groundingMetadata: unknown = undefined; let usageMetadata: GoogleGenAIUsageMetadata | null = null; - let lastResponse: GoogleGenAIGenerateContentResponse | null = null; + let promptFeedback: GoogleGenAIGenerateContentResponse["promptFeedback"]; + let lastCandidates: GoogleGenAIGenerateContentResponse["candidates"]; for (const chunk of chunks) { - lastResponse = chunk; - if (chunk.usageMetadata) { usageMetadata = chunk.usageMetadata; } + if (chunk.promptFeedback) { + promptFeedback = chunk.promptFeedback; + } if (chunk.groundingMetadata !== undefined) { groundingMetadata = chunk.groundingMetadata; } if (chunk.candidates && Array.isArray(chunk.candidates)) { + lastCandidates = chunk.candidates; for (const candidate of chunk.candidates) { if (candidate.content?.parts) { for (const part of candidate.content.parts) { @@ -2054,15 +2057,15 @@ function aggregateGenerateContentChunks( } parts.push(...otherParts); - if (parts.length > 0 && lastResponse?.candidates) { + if (lastCandidates) { const candidates: Record[] = []; - for (const candidate of lastResponse.candidates) { - const candidateDict: Record = { - content: { - parts, - role: "model", - }, - }; + for (const candidate of lastCandidates) { + // A stream can end without content (for example a SAFETY block), but its + // finish reason and ratings still belong in the output. + const candidateDict: Record = {}; + if (parts.length > 0) { + candidateDict.content = { parts, role: "model" }; + } if (candidate.finishReason !== undefined) { candidateDict.finishReason = candidate.finishReason; @@ -2077,15 +2080,22 @@ function aggregateGenerateContentChunks( candidateDict.safetyRatings = candidate.safetyRatings; } - candidates.push(candidateDict); + if (Object.keys(candidateDict).length > 0) { + candidates.push(candidateDict); + } + } + if (candidates.length > 0) { + aggregated.candidates = candidates; } - aggregated.candidates = candidates; } if (usageMetadata) { aggregated.usageMetadata = usageMetadata; populateUsageMetrics(metrics, usageMetadata); } + if (promptFeedback) { + aggregated.promptFeedback = promptFeedback; + } if (groundingMetadata !== undefined) { aggregated.groundingMetadata = groundingMetadata; } diff --git a/js/src/vendor-sdk-types/google-genai.ts b/js/src/vendor-sdk-types/google-genai.ts index eb525bcdf..c6287334a 100644 --- a/js/src/vendor-sdk-types/google-genai.ts +++ b/js/src/vendor-sdk-types/google-genai.ts @@ -237,6 +237,7 @@ export interface GoogleGenAIGenerateContentResponse { safetyRatings?: Record[]; }[]; groundingMetadata?: GoogleGenAIGroundingMetadata; + promptFeedback?: Record; usageMetadata?: GoogleGenAIUsageMetadata; text?: string; [key: string]: unknown;