@tanstack/ai-gemini 0.3.1 → 0.4.0

package/src/adapters/text.ts

@@ -20,6 +20,7 @@ import type {
   GenerateContentResponse,
   GoogleGenAI,
   Part,
+  ThinkingLevel,
 } from '@google/genai'
 import type {
   ContentPart,
@@ -29,13 +30,7 @@ import type {
   TextOptions,
 } from '@tanstack/ai'
 import type { ExternalTextProviderOptions } from '../text/text-provider-options'
-import type {
-  GeminiAudioMetadata,
-  GeminiDocumentMetadata,
-  GeminiImageMetadata,
-  GeminiMessageMetadataByModality,
-  GeminiVideoMetadata,
-} from '../message-types'
+import type { GeminiMessageMetadataByModality } from '../message-types'
 import type { GeminiClientConfig } from '../utils'
 
 /**
@@ -114,8 +109,7 @@ export class GeminiTextAdapter<
     } catch (error) {
       const timestamp = Date.now()
       yield {
-        type: 'error',
-        id: generateId(this.name),
+        type: 'RUN_ERROR',
         model: options.model,
         timestamp,
         error: {
@@ -202,25 +196,79 @@ export class GeminiTextAdapter<
     let accumulatedContent = ''
     const toolCallMap = new Map<
       string,
-      { name: string; args: string; index: number }
+      { name: string; args: string; index: number; started: boolean }
     >()
     let nextToolIndex = 0
 
+    // AG-UI lifecycle tracking
+    const runId = generateId(this.name)
+    const messageId = generateId(this.name)
+    let stepId: string | null = null
+    let hasEmittedRunStarted = false
+    let hasEmittedTextMessageStart = false
+    let hasEmittedStepStarted = false
+
     for await (const chunk of result) {
+      // Emit RUN_STARTED on first chunk
+      if (!hasEmittedRunStarted) {
+        hasEmittedRunStarted = true
+        yield {
+          type: 'RUN_STARTED',
+          runId,
+          model,
+          timestamp,
+        }
+      }
+
       if (chunk.candidates?.[0]?.content?.parts) {
         const parts = chunk.candidates[0].content.parts
 
         for (const part of parts) {
           if (part.text) {
-            accumulatedContent += part.text
-            yield {
-              type: 'content',
-              id: generateId(this.name),
-              model,
-              timestamp,
-              delta: part.text,
-              content: accumulatedContent,
-              role: 'assistant',
+            if (part.thought) {
+              // Emit STEP_STARTED on first thinking content
+              if (!hasEmittedStepStarted) {
+                hasEmittedStepStarted = true
+                stepId = generateId(this.name)
+                yield {
+                  type: 'STEP_STARTED',
+                  stepId,
+                  model,
+                  timestamp,
+                  stepType: 'thinking',
+                }
+              }
+
+              yield {
+                type: 'STEP_FINISHED',
+                stepId: stepId || generateId(this.name),
+                model,
+                timestamp,
+                delta: part.text,
+                content: part.text,
+              }
+            } else {
+              // Emit TEXT_MESSAGE_START on first text content
+              if (!hasEmittedTextMessageStart) {
+                hasEmittedTextMessageStart = true
+                yield {
+                  type: 'TEXT_MESSAGE_START',
+                  messageId,
+                  model,
+                  timestamp,
+                  role: 'assistant',
+                }
+              }
+
+              accumulatedContent += part.text
+              yield {
+                type: 'TEXT_MESSAGE_CONTENT',
+                messageId,
+                model,
+                timestamp,
+                delta: part.text,
+                content: accumulatedContent,
+              }
             }
           }
 
@@ -240,6 +288,7 @@ export class GeminiTextAdapter<
                     ? functionArgs
                     : JSON.stringify(functionArgs),
                 index: nextToolIndex++,
+                started: false,
               }
               toolCallMap.set(toolCallId, toolCallData)
             } else {
@@ -259,33 +308,51 @@ export class GeminiTextAdapter<
               }
             }
 
+            // Emit TOOL_CALL_START if not already started
+            if (!toolCallData.started) {
+              toolCallData.started = true
+              yield {
+                type: 'TOOL_CALL_START',
+                toolCallId,
+                toolName: toolCallData.name,
+                model,
+                timestamp,
+                index: toolCallData.index,
+              }
+            }
+
+            // Emit TOOL_CALL_ARGS
             yield {
-              type: 'tool_call',
-              id: generateId(this.name),
+              type: 'TOOL_CALL_ARGS',
+              toolCallId,
               model,
               timestamp,
-              toolCall: {
-                id: toolCallId,
-                type: 'function',
-                function: {
-                  name: toolCallData.name,
-                  arguments: toolCallData.args,
-                },
-              },
-              index: toolCallData.index,
+              delta: toolCallData.args,
+              args: toolCallData.args,
             }
           }
         }
       } else if (chunk.data) {
+        // Emit TEXT_MESSAGE_START on first text content
+        if (!hasEmittedTextMessageStart) {
+          hasEmittedTextMessageStart = true
+          yield {
+            type: 'TEXT_MESSAGE_START',
+            messageId,
+            model,
+            timestamp,
+            role: 'assistant',
+          }
+        }
+
         accumulatedContent += chunk.data
         yield {
-          type: 'content',
-          id: generateId(this.name),
+          type: 'TEXT_MESSAGE_CONTENT',
+          messageId,
           model,
           timestamp,
           delta: chunk.data,
           content: accumulatedContent,
-          role: 'assistant',
         }
       }
 
@@ -302,53 +369,98 @@ export class GeminiTextAdapter<
                   `${functionCall.name}_${Date.now()}_${nextToolIndex}`
                 const functionArgs = functionCall.args || {}
 
+                const argsString =
+                  typeof functionArgs === 'string'
+                    ? functionArgs
+                    : JSON.stringify(functionArgs)
+
                 toolCallMap.set(toolCallId, {
                   name: functionCall.name || '',
-                  args:
-                    typeof functionArgs === 'string'
-                      ? functionArgs
-                      : JSON.stringify(functionArgs),
+                  args: argsString,
                   index: nextToolIndex++,
+                  started: true,
                 })
 
+                // Emit TOOL_CALL_START
                 yield {
-                  type: 'tool_call',
-                  id: generateId(this.name),
+                  type: 'TOOL_CALL_START',
+                  toolCallId,
+                  toolName: functionCall.name || '',
                   model,
                   timestamp,
-                  toolCall: {
-                    id: toolCallId,
-                    type: 'function',
-                    function: {
-                      name: functionCall.name || '',
-                      arguments:
-                        typeof functionArgs === 'string'
-                          ? functionArgs
-                          : JSON.stringify(functionArgs),
-                    },
-                  },
                   index: nextToolIndex - 1,
                 }
+
+                // Emit TOOL_CALL_END with parsed input
+                let parsedInput: unknown = {}
+                try {
+                  parsedInput =
+                    typeof functionArgs === 'string'
+                      ? JSON.parse(functionArgs)
+                      : functionArgs
+                } catch {
+                  parsedInput = {}
+                }
+
+                yield {
+                  type: 'TOOL_CALL_END',
+                  toolCallId,
+                  toolName: functionCall.name || '',
+                  model,
+                  timestamp,
+                  input: parsedInput,
+                }
               }
             }
           }
         }
+
+        // Emit TOOL_CALL_END for all tracked tool calls
+        for (const [toolCallId, toolCallData] of toolCallMap.entries()) {
+          let parsedInput: unknown = {}
+          try {
+            parsedInput = JSON.parse(toolCallData.args)
+          } catch {
+            parsedInput = {}
+          }
+
+          yield {
+            type: 'TOOL_CALL_END',
+            toolCallId,
+            toolName: toolCallData.name,
+            model,
+            timestamp,
+            input: parsedInput,
+          }
+        }
+
         if (finishReason === FinishReason.MAX_TOKENS) {
           yield {
-            type: 'error',
-            id: generateId(this.name),
+            type: 'RUN_ERROR',
+            runId,
             model,
             timestamp,
             error: {
               message:
                 'The response was cut off because the maximum token limit was reached.',
+              code: 'max_tokens',
             },
           }
         }
 
+        // Emit TEXT_MESSAGE_END if we had text content
+        if (hasEmittedTextMessageStart) {
+          yield {
+            type: 'TEXT_MESSAGE_END',
+            messageId,
+            model,
+            timestamp,
+          }
+        }
+
         yield {
-          type: 'done',
-          id: generateId(this.name),
+          type: 'RUN_FINISHED',
+          runId,
           model,
           timestamp,
           finishReason: toolCallMap.size > 0 ? 'tool_calls' : 'stop',
@@ -365,20 +477,6 @@ export class GeminiTextAdapter<
   }
 
   private convertContentPartToGemini(part: ContentPart): Part {
-    const getDefaultFileType = (
-      part: 'image' | 'audio' | 'video' | 'document',
-    ) => {
-      switch (part) {
-        case 'image':
-          return 'image/jpeg'
-        case 'audio':
-          return 'audio/mp3'
-        case 'video':
-          return 'video/mp4'
-        case 'document':
-          return 'application/pdf'
-      }
-    }
     switch (part.type) {
       case 'text':
         return { text: part.content }
@@ -386,24 +484,26 @@ export class GeminiTextAdapter<
       case 'audio':
       case 'video':
       case 'document': {
-        const metadata = part.metadata as
-          | GeminiDocumentMetadata
-          | GeminiImageMetadata
-          | GeminiVideoMetadata
-          | GeminiAudioMetadata
-          | undefined
         if (part.source.type === 'data') {
           return {
             inlineData: {
               data: part.source.value,
-              mimeType: metadata?.mimeType ?? getDefaultFileType(part.type),
+              mimeType: part.source.mimeType,
             },
           }
         } else {
+          // For URL sources, use provided mimeType or fall back to reasonable defaults
+          const defaultMimeType = {
+            image: 'image/jpeg',
+            audio: 'audio/mp3',
+            video: 'video/mp4',
+            document: 'application/pdf',
+          }[part.type]
+
           return {
             fileData: {
               fileUri: part.source.value,
-              mimeType: metadata?.mimeType ?? getDefaultFileType(part.type),
+              mimeType: part.source.mimeType ?? defaultMimeType,
             },
           }
         }
@@ -476,19 +576,29 @@ export class GeminiTextAdapter<
     })
   }
 
-  private mapCommonOptionsToGemini(options: TextOptions) {
-    const providerOpts = options.modelOptions
+  private mapCommonOptionsToGemini(
+    options: TextOptions<GeminiTextProviderOptions>,
+  ) {
+    const modelOpts = options.modelOptions
+    const thinkingConfig = modelOpts?.thinkingConfig
     const requestOptions: GenerateContentParameters = {
       model: options.model,
       contents: this.formatMessages(options.messages),
       config: {
-        ...providerOpts,
+        ...modelOpts,
         temperature: options.temperature,
         topP: options.topP,
         maxOutputTokens: options.maxTokens,
+        thinkingConfig: thinkingConfig
+          ? {
+              ...thinkingConfig,
+              thinkingLevel: thinkingConfig.thinkingLevel
+                ? // Enum is provided by the SDK, we use it for the type but cast it to string constants, here we just cast them back
+                  (thinkingConfig.thinkingLevel as ThinkingLevel)
+                : undefined,
+            }
+          : undefined,
         systemInstruction: options.systemPrompts?.join('\n'),
-        ...((providerOpts as Record<string, unknown> | undefined)
-          ?.generationConfig as Record<string, unknown> | undefined),
         tools: convertToolsToProviderFormat(options.tools),
       },
     }

package/src/index.ts

@@ -79,9 +79,5 @@ export type {
   GeminiAudioMetadata,
   GeminiVideoMetadata,
   GeminiDocumentMetadata,
-  GeminiImageMimeType,
-  GeminiAudioMimeType,
-  GeminiVideoMimeType,
-  GeminiDocumentMimeType,
   GeminiMessageMetadataByModality,
 } from './message-types'

package/src/model-meta.ts

@@ -1,8 +1,9 @@
 import type {
   GeminiCachedContentOptions,
-  GeminiGenerationConfigOptions,
+  GeminiCommonConfigOptions,
   GeminiSafetyOptions,
   GeminiStructuredOutputOptions,
+  GeminiThinkingAdvancedOptions,
   GeminiThinkingOptions,
   GeminiToolConfigOptions,
 } from './text/text-provider-options'
@@ -77,10 +78,11 @@ const GEMINI_3_PRO = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
-    GeminiThinkingOptions
+    GeminiThinkingOptions &
+    GeminiThinkingAdvancedOptions
 >
 
 const GEMINI_3_FLASH = {
@@ -114,10 +116,11 @@ const GEMINI_3_FLASH = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
-    GeminiThinkingOptions
+    GeminiThinkingOptions &
+    GeminiThinkingAdvancedOptions
 >
 
 const GEMINI_3_PRO_IMAGE = {
@@ -147,10 +150,11 @@ const GEMINI_3_PRO_IMAGE = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
-    GeminiThinkingOptions
+    GeminiThinkingOptions &
+    GeminiThinkingAdvancedOptions
 >
 
 const GEMINI_2_5_PRO = {
@@ -185,7 +189,7 @@ const GEMINI_2_5_PRO = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
@@ -212,7 +216,7 @@ const GEMINI_2_5_PRO_TTS = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 
@@ -248,7 +252,7 @@ const GEMINI_2_5_FLASH = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
@@ -285,7 +289,7 @@ const GEMINI_2_5_FLASH_PREVIEW = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
@@ -318,7 +322,7 @@ const GEMINI_2_5_FLASH_IMAGE = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 /**
@@ -377,7 +381,7 @@ const GEMINI_2_5_FLASH_TTS = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 
@@ -412,7 +416,7 @@ const GEMINI_2_5_FLASH_LITE = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
@@ -448,7 +452,7 @@ const GEMINI_2_5_FLASH_LITE_PREVIEW = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
@@ -484,7 +488,7 @@ const GEMINI_2_FLASH = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions
 >
@@ -515,7 +519,7 @@ const GEMINI_2_FLASH_IMAGE = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 /* 
@@ -579,7 +583,7 @@ const GEMINI_2_FLASH_LITE = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions
 >
@@ -603,7 +607,7 @@ const IMAGEN_4_GENERATE = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 
@@ -626,7 +630,7 @@ const IMAGEN_4_GENERATE_ULTRA = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 
@@ -649,7 +653,7 @@ const IMAGEN_4_GENERATE_FAST = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 
@@ -671,7 +675,7 @@ const IMAGEN_3 = {
 } as const satisfies ModelMeta<
   GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions
 >
 /** 
@@ -909,55 +913,57 @@ export type GeminiChatModelProviderOptionsByName = {
   // Models with thinking and structured output support
   [GEMINI_3_PRO.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
-    GeminiThinkingOptions
+    GeminiThinkingOptions &
+    GeminiThinkingAdvancedOptions
   [GEMINI_3_FLASH.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
-    GeminiThinkingOptions
+    GeminiThinkingOptions &
+    GeminiThinkingAdvancedOptions
   [GEMINI_2_5_PRO.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
   [GEMINI_2_5_FLASH.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
   [GEMINI_2_5_FLASH_PREVIEW.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
   [GEMINI_2_5_FLASH_LITE.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
   [GEMINI_2_5_FLASH_LITE_PREVIEW.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions &
     GeminiThinkingOptions
   // Models with structured output but no thinking support
   [GEMINI_2_FLASH.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions
   [GEMINI_2_FLASH_LITE.name]: GeminiToolConfigOptions &
     GeminiSafetyOptions &
-    GeminiGenerationConfigOptions &
+    GeminiCommonConfigOptions &
     GeminiCachedContentOptions &
     GeminiStructuredOutputOptions
 }