diff --git a/.changeset/tidy-gemini-streams.md b/.changeset/tidy-gemini-streams.md new file mode 100644 index 000000000..6ef884bb0 --- /dev/null +++ b/.changeset/tidy-gemini-streams.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-gemini': minor +--- + +Add native structured-output streaming for the Gemini and experimental Gemini Interactions text adapters. diff --git a/packages/ai-gemini/src/adapters/text.ts b/packages/ai-gemini/src/adapters/text.ts index 1b7587bb6..a31141c26 100644 --- a/packages/ai-gemini/src/adapters/text.ts +++ b/packages/ai-gemini/src/adapters/text.ts @@ -222,6 +222,107 @@ export class GeminiTextAdapter< } } + /** Stream Gemini's native JSON response and emit the parsed object before completion. */ + async *structuredOutputStream( + options: StructuredOutputOptions, + ): AsyncIterable { + const { chatOptions: requestedChatOptions, outputSchema } = options + const chatOptions = { + ...requestedChatOptions, + runId: requestedChatOptions.runId ?? generateId(this.name), + } + const mappedOptions = this.mapCommonOptionsToGemini(chatOptions) + + try { + chatOptions.logger.request( + `activity=structuredOutputStream provider=gemini model=${this.model} messages=${chatOptions.messages.length}`, + { provider: 'gemini', model: this.model }, + ) + const result = await this.client.models.generateContentStream({ + ...mappedOptions, + config: { + ...mappedOptions.config, + responseMimeType: 'application/json', + responseSchema: outputSchema, + }, + }) + + let rawText = '' + let finished: + | Extract + | undefined + let failed = false + for await (const chunk of this.processStreamChunks( + result, + chatOptions, + chatOptions.logger, + )) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) { + rawText += chunk.delta + } + if (chunk.type === EventType.RUN_ERROR) failed = true + if (chunk.type === EventType.RUN_FINISHED) { + finished = chunk + } else { + yield chunk + } + } + + if (failed) return + if (!finished) { + yield structuredStreamError( + chatOptions, + 'Gemini structured-output stream ended without a terminal event', + 'truncated-stream', + ) + return + } + if (!rawText) { + yield structuredStreamError( + chatOptions, + 'Gemini structured-output stream contained no content', + 'empty-response', + ) + return + } + + let object: unknown + try { + object = JSON.parse(rawText) + } catch { + yield structuredStreamError( + chatOptions, + 'Failed to parse Gemini structured-output stream as JSON', + 'parse-error', + ) + return + } + + yield { + type: EventType.CUSTOM, + name: 'structured-output.complete', + value: { object, raw: rawText }, + model: chatOptions.model, + timestamp: Date.now(), + } + yield finished + } catch (error) { + const rawEvent = toRunErrorRawEvent(error) + const message = + error instanceof Error + ? error.message + : 'An unknown error occurred during structured output streaming.' + chatOptions.logger.errors('gemini.structuredOutputStream fatal', { + error, + source: 'gemini.structuredOutputStream', + }) + yield { + ...structuredStreamError(chatOptions, message, 'provider-error'), + ...(rawEvent !== undefined && { rawEvent }), + } + } + } + /** * Extract text content from a non-streaming response */ @@ -952,6 +1053,22 @@ export class GeminiTextAdapter< } } +function structuredStreamError( + options: TextOptions, + message: string, + code: string, +): Extract { + return { + type: EventType.RUN_ERROR, + runId: options.runId, + model: options.model, + timestamp: Date.now(), + message, + code, + error: { message, code }, + } +} + /** * Creates a Gemini text adapter with explicit API key. * Type resolution happens here at the call site. diff --git a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts index 9d1de8ffb..f28a9bf1a 100644 --- a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts +++ b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts @@ -438,6 +438,148 @@ export class GeminiTextInteractionsAdapter< ) } } + + async *structuredOutputStream( + options: StructuredOutputOptions, + ): AsyncIterable { + const { chatOptions, outputSchema } = options + const runId = chatOptions.runId ?? generateId(this.name) + const threadId = chatOptions.threadId ?? generateId(this.name) + const timestamp = Date.now() + const effectivePreviousInteractionId = + chatOptions.modelOptions?.previous_interaction_id ?? + this.interactionIdByThread.get(threadId) + const baseRequest = buildInteractionsRequest({ + ...chatOptions, + modelOptions: { + ...chatOptions.modelOptions, + previous_interaction_id: effectivePreviousInteractionId, + }, + }) + const request: GeminiInteractionsRequestBody = { + ...baseRequest, + stream: true, + response_format: { + type: 'text', + mime_type: 'application/json', + schema: outputSchema, + }, + } + + try { + chatOptions.logger.request( + `activity=structuredOutputStream provider=gemini-text-interactions model=${this.model} messages=${chatOptions.messages.length}`, + { provider: this.name, model: this.model, request }, + ) + const stream = (await this.client.interactions.create( + request as GeminiInteractionsRequestBody & + Parameters[0], + { signal: chatOptions.abortController?.signal }, + )) as AsyncIterable + + let rawText = '' + let finished: + | Extract + | undefined + let failed = false + for await (const chunk of translateInteractionEvents( + stream, + chatOptions.model, + runId, + threadId, + chatOptions.parentRunId, + timestamp, + this.name, + chatOptions.logger, + )) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) + rawText += chunk.delta + if ( + chunk.type === EventType.CUSTOM && + chunk.name === 'gemini.interactionId' + ) { + const value = + chunk.value as GeminiInteractionsCustomEventValue<'gemini.interactionId'> + this.interactionIdByThread.set(threadId, value.interactionId) + } + if (chunk.type === EventType.RUN_ERROR) failed = true + if (chunk.type === EventType.RUN_FINISHED) finished = chunk + else yield chunk + } + + if (failed) return + if (!finished) { + yield interactionsStructuredStreamError( + chatOptions, + runId, + 'Gemini Interactions structured-output stream ended without a terminal event', + 'truncated-stream', + ) + return + } + if (!rawText) { + yield interactionsStructuredStreamError( + chatOptions, + runId, + 'Gemini Interactions structured-output stream contained no content', + 'empty-response', + ) + return + } + let object: unknown + try { + object = JSON.parse(rawText) + } catch { + yield interactionsStructuredStreamError( + chatOptions, + runId, + 'Failed to parse Gemini Interactions structured-output stream as JSON', + 'parse-error', + ) + return + } + yield { + type: EventType.CUSTOM, + name: 'structured-output.complete', + value: { object, raw: rawText }, + model: chatOptions.model, + timestamp, + } + yield finished + } catch (error) { + const message = + error instanceof Error + ? error.message + : 'An unknown error occurred during structured output streaming.' + chatOptions.logger.errors( + 'gemini-text-interactions.structuredOutputStream fatal', + { error, source: 'gemini-text-interactions.structuredOutputStream' }, + ) + yield interactionsStructuredStreamError( + chatOptions, + runId, + message, + 'provider-error', + ) + } + } +} + +function interactionsStructuredStreamError( + options: TextOptions, + runId: string, + message: string, + code: string, +): Extract { + return { + type: EventType.RUN_ERROR, + runId, + model: options.model, + timestamp: Date.now(), + message, + code, + error: { message, code }, + } } /** @experimental Interactions API is in Beta. */ diff --git a/packages/ai-gemini/tests/gemini-adapter.test.ts b/packages/ai-gemini/tests/gemini-adapter.test.ts index 31d7c3246..3c675c65b 100644 --- a/packages/ai-gemini/tests/gemini-adapter.test.ts +++ b/packages/ai-gemini/tests/gemini-adapter.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, beforeEach, vi } from 'vitest' import { z } from 'zod' -import { chat, summarize } from '@tanstack/ai' +import { EventType, chat, summarize } from '@tanstack/ai' import type { Tool, StreamChunk } from '@tanstack/ai' import { Type, @@ -963,6 +963,67 @@ describe('GeminiAdapter through AI', () => { expect(adapter.supportsCombinedToolsAndSchema()).toBe(false) }) + it('streams structured output natively and completes before RUN_FINISHED', async () => { + mocks.generateContentStreamSpy.mockResolvedValue( + createStream([ + { + candidates: [{ content: { parts: [{ text: '{"city":' }] } }], + }, + { + candidates: [ + { + content: { parts: [{ text: '"Madrid"}' }] }, + finishReason: 'STOP', + }, + ], + usageMetadata: { totalTokenCount: 4 }, + }, + ]), + ) + + const events: Array = [] + for await (const event of createTextAdapter().structuredOutputStream({ + chatOptions: { + model: 'gemini-2.5-pro', + messages: [{ role: 'user', content: 'Return a city' }], + logger: { + request: vi.fn(), + provider: vi.fn(), + errors: vi.fn(), + } as any, + }, + outputSchema: { + type: 'object', + properties: { city: { type: 'string' } }, + }, + })) { + events.push(event) + } + + const payload = mocks.generateContentStreamSpy.mock.calls[0]![0] + expect(payload.config).toMatchObject({ + responseMimeType: 'application/json', + responseSchema: expect.objectContaining({ type: 'object' }), + }) + expect( + events.filter((event) => event.type === 'TEXT_MESSAGE_CONTENT'), + ).toHaveLength(2) + const completeIndex = events.findIndex( + (event) => + event.type === EventType.CUSTOM && + event.name === 'structured-output.complete', + ) + const finishedIndex = events.findIndex( + (event) => event.type === EventType.RUN_FINISHED, + ) + expect(completeIndex).toBeGreaterThan(-1) + expect(completeIndex).toBeLessThan(finishedIndex) + expect((events[completeIndex] as any).value).toEqual({ + object: { city: 'Madrid' }, + raw: '{"city":"Madrid"}', + }) + }) + it('routes summarize() through the gemini chat-stream path', async () => { const summaryText = 'Short and sweet.' const streamChunks = [ diff --git a/packages/ai-gemini/tests/text-interactions-adapter.test.ts b/packages/ai-gemini/tests/text-interactions-adapter.test.ts index c043efb63..977aa1d93 100644 --- a/packages/ai-gemini/tests/text-interactions-adapter.test.ts +++ b/packages/ai-gemini/tests/text-interactions-adapter.test.ts @@ -1,6 +1,6 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { z } from 'zod' -import { chat } from '@tanstack/ai' +import { EventType, chat } from '@tanstack/ai' import type { StreamChunk, Tool } from '@tanstack/ai' import { GeminiTextInteractionsAdapter } from '../src/experimental/text-interactions/adapter' import type { GeminiTextInteractionsProviderOptions } from '../src/experimental/text-interactions/adapter' @@ -1085,9 +1085,8 @@ describe('GeminiTextInteractionsAdapter', () => { }) it('structuredOutput parses JSON text from interaction.outputs', async () => { - // `chat({outputSchema, …no tools})` goes straight to structuredOutput; - // no agent-loop chatStream first. So a single mocked Interaction - // response is enough. + // Exercise the non-streaming adapter method directly; `chat()` prefers + // structuredOutputStream when the adapter provides native streaming. mocks.interactionsCreateSpy.mockResolvedValueOnce({ id: 'int_structured', status: 'completed', @@ -1100,13 +1099,23 @@ describe('GeminiTextInteractionsAdapter', () => { }) const adapter = createAdapter() - const result = await chat({ - adapter, - messages: [{ role: 'user', content: 'Give JSON' }], - outputSchema: z.object({ foo: z.string() }), + const result = await adapter.structuredOutput({ + chatOptions: { + model: 'gemini-2.5-flash', + messages: [{ role: 'user', content: 'Give JSON' }], + logger: { + request: vi.fn(), + provider: vi.fn(), + errors: vi.fn(), + } as any, + }, + outputSchema: { + type: 'object', + properties: { foo: { type: 'string' } }, + }, }) - expect(result).toEqual({ foo: 'bar' }) + expect(result.data).toEqual({ foo: 'bar' }) expect(mocks.interactionsCreateSpy).toHaveBeenCalledTimes(1) const structuredPayload = mocks.interactionsCreateSpy.mock.calls[0]![0] @@ -1119,6 +1128,81 @@ describe('GeminiTextInteractionsAdapter', () => { expect(structuredPayload.stream).toBeUndefined() }) + it('streams structured output natively and preserves terminal event ordering', async () => { + mocks.interactionsCreateSpy.mockResolvedValue( + mkStream([ + { + event_type: 'interaction.created', + interaction: { id: 'int_structured_stream', status: 'in_progress' }, + }, + { + event_type: 'step.delta', + index: 0, + delta: { type: 'text', text: '{"foo":' }, + }, + { + event_type: 'step.delta', + index: 0, + delta: { type: 'text', text: '"bar"}' }, + }, + { event_type: 'step.stop', index: 0 }, + { + event_type: 'interaction.completed', + interaction: { + id: 'int_structured_stream', + status: 'completed', + }, + }, + ]), + ) + + const adapter = createAdapter() + const events = await collectChunks( + adapter.structuredOutputStream({ + chatOptions: { + model: 'gemini-2.5-flash', + messages: [{ role: 'user', content: 'Give JSON' }], + logger: { + request: vi.fn(), + provider: vi.fn(), + errors: vi.fn(), + } as any, + }, + outputSchema: { + type: 'object', + properties: { foo: { type: 'string' } }, + }, + }), + ) + + const payload = mocks.interactionsCreateSpy.mock.calls[0]![0] + expect(payload.stream).toBe(true) + expect(payload.response_format).toMatchObject({ + type: 'text', + mime_type: 'application/json', + schema: expect.objectContaining({ type: 'object' }), + }) + expect( + events.filter((event) => event.type === 'TEXT_MESSAGE_CONTENT'), + ).toHaveLength(2) + const interactionIdIndex = events.findIndex( + (event) => + event.type === EventType.CUSTOM && + event.name === 'gemini.interactionId', + ) + const completeIndex = events.findIndex( + (event) => + event.type === EventType.CUSTOM && + event.name === 'structured-output.complete', + ) + const finishedIndex = events.findIndex( + (event) => event.type === EventType.RUN_FINISHED, + ) + expect(interactionIdIndex).toBeLessThan(completeIndex) + expect(completeIndex).toBeLessThan(finishedIndex) + expect((events[completeIndex] as any).value.object).toEqual({ foo: 'bar' }) + }) + // =========================== // interactionId map lifecycle // =========================== @@ -1507,16 +1591,24 @@ describe('GeminiTextInteractionsAdapter', () => { }, ]), ) - .mockResolvedValueOnce({ - id: 'int_struct', - status: 'completed', - steps: [ + .mockResolvedValueOnce( + mkStream([ { - type: 'model_output', - content: [{ type: 'text', text: '{"answer":"42"}' }], + event_type: 'interaction.created', + interaction: { id: 'int_struct', status: 'in_progress' }, }, - ], - }) + { + event_type: 'step.delta', + index: 0, + delta: { type: 'text', text: '{"answer":"42"}' }, + }, + { event_type: 'step.stop', index: 0 }, + { + event_type: 'interaction.completed', + interaction: { id: 'int_struct', status: 'completed' }, + }, + ]), + ) const adapter = createAdapter() diff --git a/testing/e2e/src/lib/feature-support.ts b/testing/e2e/src/lib/feature-support.ts index cd3e2203d..8a411d3d9 100644 --- a/testing/e2e/src/lib/feature-support.ts +++ b/testing/e2e/src/lib/feature-support.ts @@ -118,6 +118,7 @@ export const matrix: Record> = { // non-streaming `structuredOutput`) but aren't exercised by E2E yet. 'structured-output-stream': new Set([ 'openai', + 'gemini', 'groq', 'grok', 'bedrock',