Native Wrapper for Google Gemini (@google/genai) in LangSmith JS SDK

[reedchan7]

Feature request

Hi LangSmith Team,

I'm currently using the @google/genai client for interacting with Google's Gemini models and would like to integrate these calls with LangSmith for tracing.

I've noticed that the langsmith TypeScript SDK provides a convenient wrapOpenAI function for OpenAI clients, which makes integration straightforward by automatically handling much of the tracing setup. This is a fantastic feature!

However, I couldn't find a similar dedicated wrapper for the @google/genai client. While I understand the generic traceable function can be used to wrap any asynchronous function, a dedicated wrapper (e.g., wrapGemini or wrapGoogleGenAI) would offer a more streamlined and consistent developer experience, similar to what's available for OpenAI.

Motivation

With the increasing adoption and capabilities of Google's Gemini models, providing first-class, easy-to-use tracing support within LangSmith is crucial for developers who choose Gemini as their LLM provider. A dedicated wrapper would significantly lower the barrier to entry for integrating Gemini with LangSmith, encouraging broader adoption of LangSmith within the Gemini ecosystem and ensuring developers can easily monitor, debug, and evaluate their Gemini-powered applications. This parity with the OpenAI integration would make LangSmith a more versatile and appealing observability tool for a wider range of LLM users.

[RomainJeff]

We ended up making our own, based on the code of the wrapOpenAI if you still need here it is:

import type { GenerateContentParameters, GenerateContentResponse } from '@google/genai';
import type { RunTreeConfig } from 'langsmith';
import { KVMap } from 'langsmith/schemas';
import { isTraceableFunction, traceable } from 'langsmith/traceable';

// Extra leniency around types in case multiple Google GenAI SDK versions get installed
type GoogleGenAIType = {
  models: {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    generateContent: (...args: any[]) => any;
  };
  chats: {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    create: (...args: any[]) => any;
  };
};

type ExtraRunTreeConfig = Pick<
  Partial<RunTreeConfig>,
  'name' | 'metadata' | 'tags'
>;

type PatchedGoogleGenAIClient<T extends GoogleGenAIType> = T & {
  models: T['models'] & {
    generateContent: {
      (
        request: GenerateContentParameters,
        options?: { langsmithExtra?: ExtraRunTreeConfig }
      ): Promise<GenerateContentResponse>;
    };
  };
};

function _combineGeminiContentChunks(
  chunks: GenerateContentResponse[]
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
): any {
  if (!chunks || chunks.length === 0) {
    return { candidates: [{ content: { parts: [{ text: '' }] } }] };
  }

  const aggregatedOutput = chunks[chunks.length - 1];
  let combinedText = '';
  
  for (const chunk of chunks) {
    if (chunk.candidates?.[0]?.content?.parts) {
      for (const part of chunk.candidates[0].content.parts) {
        if ('text' in part && part.text) {
          combinedText += part.text;
        }
      }
    }
  }

  // Update the final chunk with the combined text
  if (aggregatedOutput.candidates?.[0]?.content?.parts) {
    aggregatedOutput.candidates[0].content.parts = [{ text: combinedText }];
  }

  return aggregatedOutput;
}

const geminiAggregator = (chunks: GenerateContentResponse[]) => {
  return _combineGeminiContentChunks(chunks);
};

function processGeminiCompletion(outputs: Readonly<KVMap>): KVMap {
  const geminiResponse = outputs.outputs as GenerateContentResponse;
  // copy the original object, minus usage metadata processing
  const result = { ...geminiResponse } as KVMap;
  
  // Gemini usage metadata is typically in a different structure
  if (geminiResponse.usageMetadata) {
    result.usage_metadata = {
      input_tokens: geminiResponse.usageMetadata.promptTokenCount ?? 0,
      output_tokens: geminiResponse.usageMetadata.candidatesTokenCount ?? 0,
      total_tokens: geminiResponse.usageMetadata.totalTokenCount ?? 0,
      output_token_details: {
        reasoning: geminiResponse.usageMetadata.thoughtsTokenCount ?? 0,
      },
    };
  }

  return result;
}

/**
 * Wraps a Google GenAI client's completion methods, enabling automatic LangSmith
 * tracing. Method signatures are unchanged, with the exception that you can pass
 * an additional and optional "langsmithExtra" field within the second parameter.
 * @param gemini A Google GenAI client instance.
 * @param options LangSmith options.
 * @example
 * ```ts
 * import { GoogleGenAI } from '@google/genai';
 * import { wrapGemini } from './geminiLangsmithWrapper';
 *
 * const patchedClient = wrapGemini(new GoogleGenAI({ apiKey: 'YOUR_API_KEY' }));
 *
 * const response = await patchedClient.models.generateContent(
 *   {
 *     model: 'gemini-2.0-flash',
 *     contents: 'Why is the sky blue?',
 *   },
 *   {
 *     langsmithExtra: {
 *       metadata: {
 *         additional_data: 'bar',
 *       },
 *     },
 *   },
 * );
 * ```
 */
export const wrapGemini = <T extends GoogleGenAIType>(
  gemini: T,
  options?: Partial<RunTreeConfig>
): PatchedGoogleGenAIClient<T> => {
  if (
    isTraceableFunction(gemini.models.generateContent)
  ) {
    throw new Error(
      'This instance of Google GenAI client has been already wrapped once.'
    );
  }

  // Some internal Gemini methods call each other, so we need to preserve original
  // Gemini methods.
  const tracedGeminiClient = { ...gemini };

  tracedGeminiClient.models = {
    ...gemini.models,
    generateContent: traceable(
      gemini.models.generateContent.bind(gemini.models),
      {
        name: 'ChatGemini',
        run_type: 'llm',
        aggregator: geminiAggregator,
        argsConfigPath: [1, 'langsmithExtra'],
        getInvocationParams: (payload: unknown) => {
          if (typeof payload !== 'object' || payload == null)
            return undefined;

          const params = payload as GenerateContentParameters;

          // Extract stop sequences if they exist in generation config
          let ls_stop: Array<string> | undefined;
          if (params.config?.stopSequences) {
            ls_stop = params.config.stopSequences;
          }

          return {
            ls_provider: 'google',
            ls_model_type: 'chat',
            ls_model_name: params.model,
            ls_max_tokens: params.config?.maxOutputTokens ?? undefined,
            ls_temperature: params.config?.temperature ?? undefined,
            ls_stop,
          };
        },
        processOutputs: processGeminiCompletion,
        ...options,
      }
    ),
  };

  return tracedGeminiClient as PatchedGoogleGenAIClient<T>;
};

[Eunchan24]

Is developing these features progressing?

[arath36]

would also be helpful to have multimodal support - audio, video and images.

[danias]

We ended up making our own, based on the code of the wrapOpenAI if you still need here it is:

import type { GenerateContentParameters, GenerateContentResponse } from '@google/genai';
import type { RunTreeConfig } from 'langsmith';
import { KVMap } from 'langsmith/schemas';
import { isTraceableFunction, traceable } from 'langsmith/traceable';

// Extra leniency around types in case multiple Google GenAI SDK versions get installed
type GoogleGenAIType = {
models: {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
generateContent: (...args: any[]) => any;
};
chats: {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
create: (...args: any[]) => any;
};
};

type ExtraRunTreeConfig = Pick<
Partial,
'name' | 'metadata' | 'tags'

;

type PatchedGoogleGenAIClient = T & {
models: T['models'] & {
generateContent: {
(
request: GenerateContentParameters,
options?: { langsmithExtra?: ExtraRunTreeConfig }
): Promise;
};
};
};

function _combineGeminiContentChunks(
chunks: GenerateContentResponse[]
// eslint-disable-next-line @typescript-eslint/no-explicit-any
): any {
if (!chunks || chunks.length === 0) {
return { candidates: [{ content: { parts: [{ text: '' }] } }] };
}

const aggregatedOutput = chunks[chunks.length - 1];
let combinedText = '';

for (const chunk of chunks) {
if (chunk.candidates?.[0]?.content?.parts) {
for (const part of chunk.candidates[0].content.parts) {
if ('text' in part && part.text) {
combinedText += part.text;
}
}
}
}

// Update the final chunk with the combined text
if (aggregatedOutput.candidates?.[0]?.content?.parts) {
aggregatedOutput.candidates[0].content.parts = [{ text: combinedText }];
}

return aggregatedOutput;
}

const geminiAggregator = (chunks: GenerateContentResponse[]) => {
return _combineGeminiContentChunks(chunks);
};

function processGeminiCompletion(outputs: Readonly): KVMap {
const geminiResponse = outputs.outputs as GenerateContentResponse;
// copy the original object, minus usage metadata processing
const result = { ...geminiResponse } as KVMap;

// Gemini usage metadata is typically in a different structure
if (geminiResponse.usageMetadata) {
result.usage_metadata = {
input_tokens: geminiResponse.usageMetadata.promptTokenCount ?? 0,
output_tokens: geminiResponse.usageMetadata.candidatesTokenCount ?? 0,
total_tokens: geminiResponse.usageMetadata.totalTokenCount ?? 0,
output_token_details: {
reasoning: geminiResponse.usageMetadata.thoughtsTokenCount ?? 0,
},
};
}

return result;
}

/**

• Wraps a Google GenAI client's completion methods, enabling automatic LangSmith
• tracing. Method signatures are unchanged, with the exception that you can pass
• an additional and optional "langsmithExtra" field within the second parameter.
• PARAM gemini A Google GenAI client instance.
• PARAM options LangSmith options.
• example
• import { GoogleGenAI } from '@google/genai';
• import { wrapGemini } from './geminiLangsmithWrapper';
• const patchedClient = wrapGemini(new GoogleGenAI({ apiKey: 'YOUR_API_KEY' }));
• const response = await patchedClient.models.generateContent(
• {

• model: 'gemini-2.0-flash',
  

• contents: 'Why is the sky blue?',
  

• },
• {

• langsmithExtra: {
  

•   metadata: {
  

•     additional_data: 'bar',
  

•   },
  

• },
  

• },
• );

*/
export const wrapGemini = (
gemini: T,
options?: Partial
): PatchedGoogleGenAIClient => {
if (
isTraceableFunction(gemini.models.generateContent)
) {
throw new Error(
'This instance of Google GenAI client has been already wrapped once.'
);
}

// Some internal Gemini methods call each other, so we need to preserve original
// Gemini methods.
const tracedGeminiClient = { ...gemini };

tracedGeminiClient.models = {
...gemini.models,
generateContent: traceable(
gemini.models.generateContent.bind(gemini.models),
{
name: 'ChatGemini',
run_type: 'llm',
aggregator: geminiAggregator,
argsConfigPath: [1, 'langsmithExtra'],
getInvocationParams: (payload: unknown) => {
if (typeof payload !== 'object' || payload == null)
return undefined;

      const params = payload as GenerateContentParameters;

      // Extract stop sequences if they exist in generation config
      let ls_stop: Array<string> | undefined;
      if (params.config?.stopSequences) {
        ls_stop = params.config.stopSequences;
      }

      return {
        ls_provider: 'google',
        ls_model_type: 'chat',
        ls_model_name: params.model,
        ls_max_tokens: params.config?.maxOutputTokens ?? undefined,
        ls_temperature: params.config?.temperature ?? undefined,
        ls_stop,
      };
    },
    processOutputs: processGeminiCompletion,
    ...options,
  }
),

};

return tracedGeminiClient as PatchedGoogleGenAIClient;
};

Thanks for sharing!

Here is an updated version that shows input and outputs correctly, including images:

import type { GenerateContentParameters, GenerateContentResponse } from '@google/genai';
import type { RunTreeConfig } from 'langsmith';
import { KVMap } from 'langsmith/schemas';
import { isTraceableFunction, traceable } from 'langsmith/traceable';

// Extra leniency around types in case multiple Google GenAI SDK versions get installed
type GoogleGenAIType = {
  models: {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    generateContent: (...args: any[]) => any;
  };
  chats: {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    create: (...args: any[]) => any;
  };
};

type ExtraRunTreeConfig = Pick<
  Partial<RunTreeConfig>,
  'name' | 'metadata' | 'tags'
>;

type PatchedGoogleGenAIClient<T extends GoogleGenAIType> = T & {
  models: T['models'] & {
    generateContent: {
      (
        request: GenerateContentParameters,
        options?: { langsmithExtra?: ExtraRunTreeConfig }
      ): Promise<GenerateContentResponse>;
    };
  };
};

function _combineGeminiContentChunks(
  chunks: GenerateContentResponse[]
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
): any {
  if (!chunks || chunks.length === 0) {
    return { candidates: [{ content: { parts: [{ text: '' }] } }] };
  }

  const aggregatedOutput = chunks[chunks.length - 1];
  let combinedText = '';
  
  for (const chunk of chunks) {
    if (chunk.candidates?.[0]?.content?.parts) {
      for (const part of chunk.candidates[0].content.parts) {
        if ('text' in part && part.text) {
          combinedText += part.text;
        }
      }
    }
  }

  // Update the final chunk with the combined text
  if (aggregatedOutput.candidates?.[0]?.content?.parts) {
    aggregatedOutput.candidates[0].content.parts = [{ text: combinedText }];
  }

  return aggregatedOutput;
}

const geminiAggregator = (chunks: GenerateContentResponse[]) => {
  return _combineGeminiContentChunks(chunks);
};

function processGeminiCompletion(outputs: Readonly<KVMap>): KVMap {
  const geminiResponse = outputs.outputs as GenerateContentResponse;
  
  // Extract the text content from the response
  const textContent = geminiResponse.candidates?.[0]?.content?.parts?.[0]?.text || '';
  
  // Return in chat-like format that LangSmith expects
  const result: KVMap = {
    content: textContent,
    role: 'assistant'
  };
  
  // Add usage metadata if available
  if (geminiResponse.usageMetadata) {
    result.usage_metadata = {
      input_tokens: geminiResponse.usageMetadata.promptTokenCount ?? 0,
      output_tokens: geminiResponse.usageMetadata.candidatesTokenCount ?? 0,
      total_tokens: geminiResponse.usageMetadata.totalTokenCount ?? 0,
      output_token_details: {
        reasoning: geminiResponse.usageMetadata.thoughtsTokenCount ?? 0,
      },
    };
  }

  return result;
}

function processGeminiInputs(inputs: Readonly<KVMap>): KVMap {
  // Extract the contents from the inputs
  const contents = inputs.contents as Array<{
    role: string;
    parts: Array<{ text?: string; inlineData?: { data: string; mimeType: string } }>;
  }>;

  if (!contents || !Array.isArray(contents)) {
    return inputs;
  }

  // Convert to chat-like format that LangSmith expects
  const messages = contents.map((content) => {
    const parts = content.parts || [];
    const textParts = parts.filter(part => 'text' in part && part.text);
    const imageParts = parts.filter(part => 'inlineData' in part);

    let messageContent = '';
    if (textParts.length > 0) {
      messageContent = textParts.map(part => (part as any).text).join('\n');
    }

    // Convert Gemini image format to LangSmith-compatible format
    const convertedParts = [];
    
    // Add text parts
    if (textParts.length > 0) {
      convertedParts.push({
        type: 'text',
        text: messageContent
      });
    }
    
    // Add image parts in LangSmith-compatible format
    imageParts.forEach(part => {
      const inlineData = (part as any).inlineData;
      convertedParts.push({
        type: 'image_url',
        image_url: {
          url: `data:${inlineData.mimeType};base64,${inlineData.data}`,
          detail: 'high'
        }
      });
    });

    return {
      role: content.role,
      content: convertedParts.length === 1 && convertedParts[0].type === 'text' 
        ? convertedParts[0].text 
        : convertedParts
    };
  });

  return {
    messages,
    model: inputs.model,
    ...(inputs.config && { config: inputs.config })
  };
}

/**
 * Wraps a Google GenAI client's completion methods, enabling automatic LangSmith
 * tracing. Method signatures are unchanged, with the exception that you can pass
 * an additional and optional "langsmithExtra" field within the second parameter.
 * @param gemini A Google GenAI client instance.
 * @param options LangSmith options.
 * @example
 * ```ts
 * import { GoogleGenAI } from '@google/genai';
 * import { wrapGemini } from './geminiLangsmithWrapper';
 *
 * const patchedClient = wrapGemini(new GoogleGenAI({ apiKey: 'YOUR_API_KEY' }));
 *
 * const response = await patchedClient.models.generateContent(
 *   {
 *     model: 'gemini-2.0-flash',
 *     contents: 'Why is the sky blue?',
 *   },
 *   {
 *     langsmithExtra: {
 *       metadata: {
 *         additional_data: 'bar',
 *       },
 *     },
 *   },
 * );
 * ```
 */
export const wrapGemini = <T extends GoogleGenAIType>(
  gemini: T,
  options?: Partial<RunTreeConfig>
): PatchedGoogleGenAIClient<T> => {
  if (
    isTraceableFunction(gemini.models.generateContent)
  ) {
    throw new Error(
      'This instance of Google GenAI client has been already wrapped once.'
    );
  }

  // Some internal Gemini methods call each other, so we need to preserve original
  // Gemini methods.
  const tracedGeminiClient = { ...gemini };

  tracedGeminiClient.models = {
    ...gemini.models,
    generateContent: traceable(
      gemini.models.generateContent.bind(gemini.models),
      {
        name: 'ChatGemini',
        run_type: 'llm',
        aggregator: geminiAggregator,
        argsConfigPath: [1, 'langsmithExtra'],
        getInvocationParams: (payload: unknown) => {
          if (typeof payload !== 'object' || payload == null)
            return undefined;

          const params = payload as GenerateContentParameters;

          // Extract stop sequences if they exist in generation config
          let ls_stop: Array<string> | undefined;
          if (params.config?.stopSequences) {
            ls_stop = params.config.stopSequences;
          }

          return {
            ls_provider: 'google',
            ls_model_type: 'chat',
            ls_model_name: params.model,
            ls_max_tokens: params.config?.maxOutputTokens ?? undefined,
            ls_temperature: params.config?.temperature ?? undefined,
            ls_stop,
          };
        },
        processInputs: processGeminiInputs,
        processOutputs: processGeminiCompletion,
        ...options,
      }
    ),
  };

  return tracedGeminiClient as PatchedGoogleGenAIClient<T>;
};