MLflow Typescript SDK - Gemini

Seamlessly integrate MLflow Tracing with Gemini to automatically trace your Claude API calls.

| Package | NPM | Description | | -------------------- | --------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------- | | @mlflow/gemini | | Auto-instrumentation integration for Gemini. |

Installation

npm install @mlflow/gemini

The package includes the @mlflow/core package and @google/genai package as peer dependencies. Depending on your package manager, you may need to install these two packages separately.

Quickstart

Start MLflow Tracking Server if you don't have one already:

pip install mlflow
mlflow server --backend-store-uri sqlite:///mlruns.db --port 5000

Self-hosting MLflow server requires Python 3.10 or higher. If you don't have one, you can also use managed MLflow service for free to get started quickly.

Instantiate MLflow SDK in your application:

import * as mlflow from '@mlflow/core';

mlflow.init({
  trackingUri: 'http://localhost:5000',
  experimentId: '<experiment-id>',
});

Create a trace for Gemini:

import { tracedGemini } from '@mlflow/gemini';
import { GoogleGenAI } from '@google/genai';

const gemini = new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY });
const client = tracedGemini(gemini);

const response = await client.models.generateContent({
  model: 'gemini-2.0-flash-001',
  contents: 'Hello Gemini',
});

View traces in MLflow UI:

Documentation 📘

Official documentation for MLflow Typescript SDK can be found here.

License

This project is licensed under the Apache License 2.0.