@tanstack/ai-gemini 0.0.3 → 0.1.0

package/src/{gemini-adapter.ts → adapters/text.ts}

@@ -1,75 +1,109 @@
-import { FinishReason, GoogleGenAI } from '@google/genai'
-import { BaseAdapter } from '@tanstack/ai'
-import { GEMINI_EMBEDDING_MODELS, GEMINI_MODELS } from './model-meta'
-import { convertToolsToProviderFormat } from './tools/tool-converter'
-import type {
-  AIAdapterConfig,
-  ChatOptions,
-  ContentPart,
-  EmbeddingOptions,
-  EmbeddingResult,
-  ModelMessage,
-  StreamChunk,
-  SummarizationOptions,
-  SummarizationResult,
-} from '@tanstack/ai'
+import { FinishReason } from '@google/genai'
+import { BaseTextAdapter } from '@tanstack/ai/adapters'
+import { convertToolsToProviderFormat } from '../tools/tool-converter'
+import {
+  createGeminiClient,
+  generateId,
+  getGeminiApiKeyFromEnv,
+} from '../utils'
 import type {
+  GEMINI_MODELS,
   GeminiChatModelProviderOptionsByName,
   GeminiModelInputModalitiesByName,
-} from './model-meta'
-import type { ExternalTextProviderOptions } from './text/text-provider-options'
+} from '../model-meta'
+import type {
+  StructuredOutputOptions,
+  StructuredOutputResult,
+} from '@tanstack/ai/adapters'
 import type {
   GenerateContentParameters,
   GenerateContentResponse,
+  GoogleGenAI,
   Part,
 } from '@google/genai'
+import type {
+  ContentPart,
+  Modality,
+  ModelMessage,
+  StreamChunk,
+  TextOptions,
+} from '@tanstack/ai'
+import type { ExternalTextProviderOptions } from '../text/text-provider-options'
 import type {
   GeminiAudioMetadata,
   GeminiDocumentMetadata,
   GeminiImageMetadata,
   GeminiMessageMetadataByModality,
   GeminiVideoMetadata,
-} from './message-types'
+} from '../message-types'
+import type { GeminiClientConfig } from '../utils'
 
-export interface GeminiAdapterConfig extends AIAdapterConfig {
-  apiKey: string
-}
+/**
+ * Configuration for Gemini text adapter
+ */
+export interface GeminiTextConfig extends GeminiClientConfig {}
 
 /**
- * Gemini-specific provider options
- * Based on Google Generative AI SDK
- * @see https://ai.google.dev/api/rest/v1/GenerationConfig
+ * Gemini-specific provider options for text/chat
  */
-export type GeminiProviderOptions = ExternalTextProviderOptions
+export type GeminiTextProviderOptions = ExternalTextProviderOptions
 
-export class GeminiAdapter extends BaseAdapter<
-  typeof GEMINI_MODELS,
-  typeof GEMINI_EMBEDDING_MODELS,
-  GeminiProviderOptions,
-  Record<string, any>,
-  GeminiChatModelProviderOptionsByName,
-  GeminiModelInputModalitiesByName,
+// ===========================
+// Type Resolution Helpers
+// ===========================
+
+/**
+ * Resolve provider options for a specific model.
+ * If the model has explicit options in the map, use those; otherwise use base options.
+ */
+type ResolveProviderOptions<TModel extends string> =
+  TModel extends keyof GeminiChatModelProviderOptionsByName
+    ? GeminiChatModelProviderOptionsByName[TModel]
+    : GeminiTextProviderOptions
+
+/**
+ * Resolve input modalities for a specific model.
+ * If the model has explicit modalities in the map, use those; otherwise use all modalities.
+ */
+type ResolveInputModalities<TModel extends string> =
+  TModel extends keyof GeminiModelInputModalitiesByName
+    ? GeminiModelInputModalitiesByName[TModel]
+    : readonly ['text', 'image', 'audio', 'video', 'document']
+
+// ===========================
+// Adapter Implementation
+// ===========================
+
+/**
+ * Gemini Text (Chat) Adapter
+ *
+ * Tree-shakeable adapter for Gemini chat/text completion functionality.
+ * Import only what you need for smaller bundle sizes.
+ */
+export class GeminiTextAdapter<
+  TModel extends (typeof GEMINI_MODELS)[number],
+  TProviderOptions extends object = ResolveProviderOptions<TModel>,
+  TInputModalities extends ReadonlyArray<Modality> =
+    ResolveInputModalities<TModel>,
+> extends BaseTextAdapter<
+  TModel,
+  TProviderOptions,
+  TInputModalities,
   GeminiMessageMetadataByModality
 > {
-  name = 'gemini'
-  models = GEMINI_MODELS
-  embeddingModels = GEMINI_EMBEDDING_MODELS
-  declare _modelProviderOptionsByName: GeminiChatModelProviderOptionsByName
-  declare _modelInputModalitiesByName: GeminiModelInputModalitiesByName
-  declare _messageMetadataByModality: GeminiMessageMetadataByModality
+  readonly kind = 'text' as const
+  readonly name = 'gemini' as const
+
   private client: GoogleGenAI
 
-  constructor(config: GeminiAdapterConfig) {
-    super(config)
-    this.client = new GoogleGenAI({
-      apiKey: config.apiKey,
-    })
+  constructor(config: GeminiTextConfig, model: TModel) {
+    super({}, model)
+    this.client = createGeminiClient(config)
   }
 
   async *chatStream(
-    options: ChatOptions<string, GeminiProviderOptions>,
+    options: TextOptions<GeminiTextProviderOptions>,
   ): AsyncIterable<StreamChunk> {
-    // Map common options to Gemini format
     const mappedOptions = this.mapCommonOptionsToGemini(options)
 
     try {
@@ -81,7 +115,7 @@ export class GeminiAdapter extends BaseAdapter<
       const timestamp = Date.now()
       yield {
         type: 'error',
-        id: this.generateId(),
+        id: generateId(this.name),
         model: options.model,
         timestamp,
         error: {
@@ -94,125 +128,74 @@ export class GeminiAdapter extends BaseAdapter<
     }
   }
 
-  async summarize(options: SummarizationOptions): Promise<SummarizationResult> {
-    const prompt = this.buildSummarizationPrompt(options, options.text)
+  /**
+   * Generate structured output using Gemini's native JSON response format.
+   * Uses responseMimeType: 'application/json' and responseSchema for structured output.
+   * The outputSchema is already JSON Schema (converted in the ai layer).
+   */
+  async structuredOutput(
+    options: StructuredOutputOptions<GeminiTextProviderOptions>,
+  ): Promise<StructuredOutputResult<unknown>> {
+    const { chatOptions, outputSchema } = options
 
-    // Use models API like chatCompletion
-    const result = await this.client.models.generateContent({
-      model: options.model,
-      contents: [{ role: 'user', parts: [{ text: prompt }] }],
-      config: {
-        temperature: 0.3,
-        maxOutputTokens: options.maxLength || 500,
-      },
-    })
+    const mappedOptions = this.mapCommonOptionsToGemini(chatOptions)
 
-    // Extract text from candidates or use .text() method
-    let summary = ''
-    if (result.candidates?.[0]?.content?.parts) {
-      const parts = result.candidates[0].content.parts
-      for (const part of parts) {
-        if (part.text) {
-          summary += part.text
-        }
-      }
-    }
+    try {
+      // Add structured output configuration
+      const result = await this.client.models.generateContent({
+        ...mappedOptions,
+        config: {
+          ...mappedOptions.config,
+          responseMimeType: 'application/json',
+          responseSchema: outputSchema,
+        },
+      })
 
-    if (!summary && typeof result.text === 'string') {
-      summary = result.text
-    }
+      // Extract text content from the response
+      const rawText = this.extractTextFromResponse(result)
 
-    const promptTokens = this.estimateTokens(prompt)
-    const completionTokens = this.estimateTokens(summary)
+      // Parse the JSON response
+      let parsed: unknown
+      try {
+        parsed = JSON.parse(rawText)
+      } catch {
+        throw new Error(
+          `Failed to parse structured output as JSON. Content: ${rawText.slice(0, 200)}${rawText.length > 200 ? '...' : ''}`,
+        )
+      }
 
-    return {
-      id: this.generateId(),
-      model: options.model,
-      summary,
-      usage: {
-        promptTokens,
-        completionTokens,
-        totalTokens: promptTokens + completionTokens,
-      },
+      return {
+        data: parsed,
+        rawText,
+      }
+    } catch (error) {
+      throw new Error(
+        error instanceof Error
+          ? error.message
+          : 'An unknown error occurred during structured output generation.',
+      )
     }
   }
 
-  async createEmbeddings(options: EmbeddingOptions): Promise<EmbeddingResult> {
-    const inputs = Array.isArray(options.input)
-      ? options.input
-      : [options.input]
-
-    // According to docs: contents can be a string or array of strings
-    // Response has embeddings (plural) array with values property
-    const result = await this.client.models.embedContent({
-      model: options.model,
-      contents: inputs,
-    })
+  /**
+   * Extract text content from a non-streaming response
+   */
+  private extractTextFromResponse(response: GenerateContentResponse): string {
+    let textContent = ''
 
-    // Extract embeddings from result.embeddings array
-    const embeddings: Array<Array<number>> = []
-    if (result.embeddings && Array.isArray(result.embeddings)) {
-      for (const embedding of result.embeddings) {
-        if (embedding.values && Array.isArray(embedding.values)) {
-          embeddings.push(embedding.values)
-        } else if (Array.isArray(embedding)) {
-          embeddings.push(embedding)
+    if (response.candidates?.[0]?.content?.parts) {
+      for (const part of response.candidates[0].content.parts) {
+        if (part.text) {
+          textContent += part.text
         }
       }
     }
 
-    const promptTokens = inputs.reduce(
-      (sum, input) => sum + this.estimateTokens(input),
-      0,
-    )
-
-    return {
-      id: this.generateId(),
-      model: options.model || 'gemini-embedding-001',
-      embeddings,
-      usage: {
-        promptTokens,
-        totalTokens: promptTokens,
-      },
-    }
-  }
-
-  private buildSummarizationPrompt(
-    options: SummarizationOptions,
-    text: string,
-  ): string {
-    let prompt = 'You are a professional summarizer. '
-
-    switch (options.style) {
-      case 'bullet-points':
-        prompt += 'Provide a summary in bullet point format. '
-        break
-      case 'paragraph':
-        prompt += 'Provide a summary in paragraph format. '
-        break
-      case 'concise':
-        prompt += 'Provide a very concise summary in 1-2 sentences. '
-        break
-      default:
-        prompt += 'Provide a clear and concise summary. '
-    }
-
-    if (options.focus && options.focus.length > 0) {
-      prompt += `Focus on the following aspects: ${options.focus.join(', ')}. `
-    }
-
-    prompt += `\n\nText to summarize:\n${text}\n\nSummary:`
-
-    return prompt
-  }
-
-  private estimateTokens(text: string): number {
-    // Rough approximation: 1 token ≈ 4 characters
-    return Math.ceil(text.length / 4)
+    return textContent
   }
 
   private async *processStreamChunks(
-    result: AsyncGenerator<GenerateContentResponse, any, any>,
+    result: AsyncGenerator<GenerateContentResponse, unknown, unknown>,
     model: string,
   ): AsyncIterable<StreamChunk> {
     const timestamp = Date.now()
@@ -222,20 +205,17 @@ export class GeminiAdapter extends BaseAdapter<
       { name: string; args: string; index: number }
     >()
     let nextToolIndex = 0
-    // Iterate over the stream result (it's already an AsyncGenerator)
+
     for await (const chunk of result) {
-      // Extract content from candidates[0].content.parts
-      // Parts can contain text or functionCall
       if (chunk.candidates?.[0]?.content?.parts) {
         const parts = chunk.candidates[0].content.parts
 
         for (const part of parts) {
-          // Handle text content
           if (part.text) {
             accumulatedContent += part.text
             yield {
               type: 'content',
-              id: this.generateId(),
+              id: generateId(this.name),
               model,
               timestamp,
               delta: part.text,
@@ -244,15 +224,12 @@ export class GeminiAdapter extends BaseAdapter<
             }
           }
 
-          // Handle function calls (tool calls)
-          // Check both camelCase (SDK) and snake_case (direct API) formats
           const functionCall = part.functionCall
           if (functionCall) {
             const toolCallId =
               functionCall.name || `call_${Date.now()}_${nextToolIndex}`
             const functionArgs = functionCall.args || {}
 
-            // Check if we've seen this tool call before (for streaming args)
             let toolCallData = toolCallMap.get(toolCallId)
             if (!toolCallData) {
               toolCallData = {
@@ -265,8 +242,6 @@ export class GeminiAdapter extends BaseAdapter<
               }
               toolCallMap.set(toolCallId, toolCallData)
             } else {
-              // Merge arguments if streaming
-
               try {
                 const existingArgs = JSON.parse(toolCallData.args)
                 const newArgs =
@@ -276,7 +251,6 @@ export class GeminiAdapter extends BaseAdapter<
                 const mergedArgs = { ...existingArgs, ...newArgs }
                 toolCallData.args = JSON.stringify(mergedArgs)
               } catch {
-                // If parsing fails, use new args
                 toolCallData.args =
                   typeof functionArgs === 'string'
                     ? functionArgs
@@ -286,7 +260,7 @@ export class GeminiAdapter extends BaseAdapter<
 
             yield {
               type: 'tool_call',
-              id: this.generateId(),
+              id: generateId(this.name),
               model,
               timestamp,
               toolCall: {
@@ -302,11 +276,10 @@ export class GeminiAdapter extends BaseAdapter<
           }
         }
       } else if (chunk.data) {
-        // Fallback to chunk.data if available
         accumulatedContent += chunk.data
         yield {
           type: 'content',
-          id: this.generateId(),
+          id: generateId(this.name),
           model,
           timestamp,
           delta: chunk.data,
@@ -315,20 +288,14 @@ export class GeminiAdapter extends BaseAdapter<
         }
       }
 
-      // Check for finish reason
       if (chunk.candidates?.[0]?.finishReason) {
         const finishReason = chunk.candidates[0].finishReason
 
-        // UNEXPECTED_TOOL_CALL means Gemini tried to call a function but it wasn't properly declared
-        // This typically means there's an issue with the tool declaration format
-        // We should map it to tool_calls to try to process it anyway
         if (finishReason === FinishReason.UNEXPECTED_TOOL_CALL) {
-          // Try to extract function call from content.parts if available
           if (chunk.candidates[0].content?.parts) {
             for (const part of chunk.candidates[0].content.parts) {
               const functionCall = part.functionCall
               if (functionCall) {
-                // We found a function call - process it
                 const toolCallId =
                   functionCall.name || `call_${Date.now()}_${nextToolIndex}`
                 const functionArgs = functionCall.args || {}
@@ -344,7 +311,7 @@ export class GeminiAdapter extends BaseAdapter<
 
                 yield {
                   type: 'tool_call',
-                  id: this.generateId(),
+                  id: generateId(this.name),
                   model,
                   timestamp,
                   toolCall: {
@@ -367,7 +334,7 @@ export class GeminiAdapter extends BaseAdapter<
         if (finishReason === FinishReason.MAX_TOKENS) {
           yield {
             type: 'error',
-            id: this.generateId(),
+            id: generateId(this.name),
             model,
             timestamp,
             error: {
@@ -379,14 +346,14 @@ export class GeminiAdapter extends BaseAdapter<
 
         yield {
           type: 'done',
-          id: this.generateId(),
+          id: generateId(this.name),
           model,
           timestamp,
           finishReason: toolCallMap.size > 0 ? 'tool_calls' : 'stop',
           usage: chunk.usageMetadata
             ? {
                 promptTokens: chunk.usageMetadata.promptTokenCount ?? 0,
-                completionTokens: chunk.usageMetadata.thoughtsTokenCount ?? 0,
+                completionTokens: chunk.usageMetadata.candidatesTokenCount ?? 0,
                 totalTokens: chunk.usageMetadata.totalTokenCount ?? 0,
               }
             : undefined,
@@ -396,6 +363,20 @@ export class GeminiAdapter extends BaseAdapter<
   }
 
   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 }
@@ -409,25 +390,23 @@ export class GeminiAdapter extends BaseAdapter<
           | GeminiVideoMetadata
           | GeminiAudioMetadata
           | undefined
-        // Gemini uses inlineData for base64 and fileData for URLs
         if (part.source.type === 'data') {
           return {
             inlineData: {
               data: part.source.value,
-              mimeType: metadata?.mimeType ?? 'image/jpeg',
+              mimeType: metadata?.mimeType ?? getDefaultFileType(part.type),
             },
           }
         } else {
           return {
             fileData: {
               fileUri: part.source.value,
-              mimeType: metadata?.mimeType ?? 'image/jpeg',
+              mimeType: metadata?.mimeType ?? getDefaultFileType(part.type),
             },
           }
         }
       }
       default: {
-        // Exhaustive check - this should never happen with known types
         const _exhaustiveCheck: never = part
         throw new Error(
           `Unsupported content part type: ${(_exhaustiveCheck as ContentPart).type}`,
@@ -443,26 +422,29 @@ export class GeminiAdapter extends BaseAdapter<
       const role: 'user' | 'model' = msg.role === 'assistant' ? 'model' : 'user'
       const parts: Array<Part> = []
 
-      // Handle multimodal content (array of ContentPart)
       if (Array.isArray(msg.content)) {
         for (const contentPart of msg.content) {
           parts.push(this.convertContentPartToGemini(contentPart))
         }
       } else if (msg.content) {
-        // Handle string content (backward compatibility)
         parts.push({ text: msg.content })
       }
 
-      // Handle tool calls (from assistant)
       if (msg.role === 'assistant' && msg.toolCalls?.length) {
         for (const toolCall of msg.toolCalls) {
           let parsedArgs: Record<string, unknown> = {}
           try {
             parsedArgs = toolCall.function.arguments
-              ? JSON.parse(toolCall.function.arguments)
+              ? (JSON.parse(toolCall.function.arguments) as Record<
+                  string,
+                  unknown
+                >)
               : {}
           } catch {
-            parsedArgs = toolCall.function.arguments as any
+            parsedArgs = toolCall.function.arguments as unknown as Record<
+              string,
+              unknown
+            >
           }
 
           parts.push({
@@ -474,11 +456,10 @@ export class GeminiAdapter extends BaseAdapter<
         }
       }
 
-      // Handle tool results (from tool role)
       if (msg.role === 'tool' && msg.toolCallId) {
         parts.push({
           functionResponse: {
-            name: msg.toolCallId, // Gemini uses function name here
+            name: msg.toolCallId,
             response: {
               content: msg.content || '',
             },
@@ -492,22 +473,20 @@ export class GeminiAdapter extends BaseAdapter<
       }
     })
   }
-  /**
-   * Maps common options to Gemini-specific format
-   * Handles translation of normalized options to Gemini's API format
-   */
-  private mapCommonOptionsToGemini(options: ChatOptions) {
-    const providerOpts = options.providerOptions
+
+  private mapCommonOptionsToGemini(options: TextOptions) {
+    const providerOpts = options.modelOptions
     const requestOptions: GenerateContentParameters = {
       model: options.model,
       contents: this.formatMessages(options.messages),
       config: {
         ...providerOpts,
-        temperature: options.options?.temperature,
-        topP: options.options?.topP,
-        maxOutputTokens: options.options?.maxTokens,
+        temperature: options.temperature,
+        topP: options.topP,
+        maxOutputTokens: options.maxTokens,
         systemInstruction: options.systemPrompts?.join('\n'),
-        ...providerOpts?.generationConfig,
+        ...((providerOpts as Record<string, unknown> | undefined)
+          ?.generationConfig as Record<string, unknown> | undefined),
         tools: convertToolsToProviderFormat(options.tools),
       },
     }
@@ -517,61 +496,33 @@ export class GeminiAdapter extends BaseAdapter<
 }
 
 /**
- * Creates a Gemini adapter with simplified configuration
- * @param apiKey - Your Google API key
- * @returns A fully configured Gemini adapter instance
- *
- * @example
- * ```typescript
- * const gemini = createGemini("AIza...");
- *
- * const ai = new AI({
- *   adapters: {
- *     gemini,
- *   }
- * });
- * ```
+ * Creates a Gemini text adapter with explicit API key.
+ * Type resolution happens here at the call site.
  */
-export function createGemini(
+export function createGeminiChat<TModel extends (typeof GEMINI_MODELS)[number]>(
+  model: TModel,
   apiKey: string,
-  config?: Omit<GeminiAdapterConfig, 'apiKey'>,
-): GeminiAdapter {
-  return new GeminiAdapter({ apiKey, ...config })
+  config?: Omit<GeminiTextConfig, 'apiKey'>,
+): GeminiTextAdapter<
+  TModel,
+  ResolveProviderOptions<TModel>,
+  ResolveInputModalities<TModel>
+> {
+  return new GeminiTextAdapter({ apiKey, ...config }, model)
 }
 
 /**
- * Create a Gemini adapter with automatic API key detection from environment variables.
- *
- * Looks for `GOOGLE_API_KEY` or `GEMINI_API_KEY` in:
- * - `process.env` (Node.js)
- * - `window.env` (Browser with injected env)
- *
- * @param config - Optional configuration (excluding apiKey which is auto-detected)
- * @returns Configured Gemini adapter instance
- * @throws Error if API key is not found in environment
- *
- * @example
- * ```typescript
- * // Automatically uses GOOGLE_API_KEY or GEMINI_API_KEY from environment
- * const aiInstance = ai(gemini());
- * ```
+ * Creates a Gemini text adapter with automatic API key detection.
+ * Type resolution happens here at the call site.
  */
-export function gemini(
-  config?: Omit<GeminiAdapterConfig, 'apiKey'>,
-): GeminiAdapter {
-  const env =
-    typeof globalThis !== 'undefined' && (globalThis as any).window?.env
-      ? (globalThis as any).window.env
-      : typeof process !== 'undefined'
-        ? process.env
-        : undefined
-  const key = env?.GOOGLE_API_KEY || env?.GEMINI_API_KEY
-
-  if (!key) {
-    throw new Error(
-      'GOOGLE_API_KEY or GEMINI_API_KEY is required. Please set it in your environment variables or use createGemini(apiKey, config) instead.',
-    )
-  }
-
-  return createGemini(key, config)
+export function geminiText<TModel extends (typeof GEMINI_MODELS)[number]>(
+  model: TModel,
+  config?: Omit<GeminiTextConfig, 'apiKey'>,
+): GeminiTextAdapter<
+  TModel,
+  ResolveProviderOptions<TModel>,
+  ResolveInputModalities<TModel>
+> {
+  const apiKey = getGeminiApiKeyFromEnv()
+  return createGeminiChat(model, apiKey, config)
 }