Google Gemini MLflow Observability

MLflow Observability

gemini-cookbookMLflowgemini-apiexamplesgemini

alph-notebooks/gemini-cookbook / MLflow_Observability.ipynb

Export

Run Notebooks

Contents

No cells yet

Add cells to see them here

Copyright 2024 Google LLC.

[ ]

Gemini API: LLM Observability with MLflow

Overview

MLflow is an open-source platform to assist machine learning practitioners and teams in handling the complexities of the machine learning process.

It provides MLflow Tracing that enhances LLM observability in your Generative AI applications by capturing detailed information about the execution of your application’s services. Tracing provides a way to record the inputs, outputs, and metadata associated with each intermediate step of a request, enabling you to easily pinpoint the source of bugs and unexpected behaviors.

MLflow provides a built-in integration with Google Gen AI SDK that enables you to instrument your Gemini calls easily. This cookbook describes the basic usage of the MLflow tracing integration with the google-genai package.

This notebook was contributed by Tomu Hirata.

Have a cool Gemini example? Feel free to share it too!

Installation

Install Google Gen AI SDK (google-genai) and MLflow (mlflow). See troubleshooting for MLFlow installation for other install options.

[2]


[notice] A new release of pip is available: 24.3.1 -> 25.0.1
[notice] To update, run: pip install --upgrade pip
Note: you may need to restart the kernel to use updated packages.

Create Gemini Client with your API key

Let's create an API client and pass your API key. If you do not have API ket yet, visit AI Studio to create one.

[3]

Tracking Server

There are several options to run MLflow tracking server: local tracking server, Databricks Free Trial, and production Databricks. See our documentation for the comparison. In this example, you use Databricks Free Trial which allows easy connection from Colab notebooks and enable you to use managed MLflow for free. Follow the steps below to create an account and generate a Personal Access Token (PAT) to connect to your workspace.

Go to the Databricks Trial Signup Page and create your account
Follow the steps in this guide to create a PAT for your Databricks workspace user
You need following information to connect to your workspace from your notebook
- Databricks Host: Use "https://<your workspace host>.cloud.databricks.com/
- Token: Your personal access token for your Databricks Workspace.

[4]

Then, set the tracking server uri and experiment name.

[ ]

Enable AutoLogging

MLflow Tracing provides automatic tracing capability for Google Gemini. By enabling auto tracing for Gemini by calling the mlflow.gemini.autolog() function, MLflow will capture nested traces and log them to the active MLflow Experiment upon invocation of Gemini Python SDK.

MLflow trace automatically captures the following information about Gemini calls:

Prompts and completion responses
Latencies
Model name
Additional metadata such as temperature, max_tokens, if specified.
Function calling if returned in the response
Any exception if raised

[7]

Call Simple Content Generation

Let's run client.models.generate_content to try a simple text generation use case with MLflow tracing. For Jupyter Notebook users, MLflow provides a convenient way to see the generated traces on your notebook. See this blog for more information. For users who use other platforms, visit "http://localhost:5000" to see MLflow UI.

[ ]

MODEL_ID

[17]

'cold\n'

Multi-Turn Chat Interactions

MLflow tracing captures the structure of your interactions with the Gemini API. Run the following cell to try multi-turn chat and see how MLflow captures the interaction.

[16]

A computer follows instructions, like a recipe, to do things with numbers and pictures.

A computer works by executing a sequence of instructions, called a program, written in a language it understands. These instructions manipulate binary data (0s and 1s) representing information, performing calculations, storing and retrieving data from memory and storage devices, and interacting with input and output devices like the keyboard, screen, and network to complete tasks.

Function Call

If your application uses function calling of the Gemini API, the function definition, function call and function response are captured by MLflow automatically.

[18]

[19]

[24]

'2508\n'

Conclusion

That's all for this cookbook. MLflow tracing provides many features that are not included in this notebook and actively releases new features. Visit MLflow Gemini tracing integration for more configurations and MLflow Tracing Overview for general offerings.