MLflow Callback Handler in LangChain - Track Your LLM Experiments

The MLflow callback handler in LangChain allows you to track and monitor your LLM experiments by logging metrics, artifacts and traces to MLflow's tracking server. In this guide, we'll explore how to use this powerful callback to get insights into your LangChain applications.

What is MLflowCallbackHandler?

MLflowCallbackHandler is a callback handler in LangChain that integrates with MLflow to track experiments and log information about LLM runs. It captures important metrics, artifacts and execution traces when running chains, agents, chat models and other LangChain components.

The handler can log:

• Metrics like token usage and number of LLM calls
• Chain/Agent execution traces as JSON
• Input/output data as Pandas DataFrames
• Chain/Agent definitions as MLflow models
• Text complexity metrics (when spaCy is installed)

Reference

Here are the key methods and attributes of MLflowCallbackHandler:

How to Use MLflowCallbackHandler

Basic Setup

First, import and initialize the callback handler:

from langchain_community.callbacks import MlflowCallbackHandler

mlflow_callback = MlflowCallbackHandler(
    experiment="langchain",  # Name of the experiment
    name="my-run",  # Name of the run 
    tracking_uri="http://localhost:5000"  # MLflow tracking server URI
)

Using with Chains

You can use the callback handler when invoking chains to track their execution:

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# Create a simple chain
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant."),
    ("human", "{input}")
])
chain = prompt | ChatOpenAI() 

# Invoke chain with MLflow tracking
chain.invoke(
    {"input": "What is MLflow?"}, 
    config={"callbacks": [mlflow_callback]}
)

# Important: Flush tracked data at the end
mlflow_callback.flush_tracker()

Using with Agents

The handler can also track agent executions:

from langchain.agents import initialize_agent, load_tools
from langchain_openai import ChatOpenAI

# Create an agent
llm = ChatOpenAI()
tools = load_tools(["serpapi", "llm-math"], llm=llm)
agent = initialize_agent(tools, llm, agent="zero-shot-react-description")

# Run agent with MLflow tracking
agent.invoke(
    "What was the temperature in SF yesterday?",
    config={"callbacks": [mlflow_callback]}
)

mlflow_callback.flush_tracker()

Configuring What to Track

You can configure which components to track using the ignore flags:

mlflow_callback = MlflowCallbackHandler(
    ignore_llm=False,       # Track LLM calls
    ignore_chain=False,     # Track chains
    ignore_agent=False,     # Track agents  
    ignore_chat_model=False # Track chat models
)

Accessing the MLflow UI

After running your code with the callback handler, you can view the tracked data in the MLflow UI:

1. Start the MLflow UI server:

mlflow ui

2. Open http://localhost:5000 in your browser to see:

• Metrics like token usage
• Execution traces
• Input/output artifacts
• Model information

Remember to always call flush_tracker() at the end of your runs to ensure all data is written to MLflow.

This callback handler provides valuable insights into your LangChain applications and helps with experiment tracking, debugging and optimization.