MLflow

MLflow 是一个多功能的开源平台，用于管理机器学习生命周期中的工作流程和工件。它内置了与许多流行的 ML 库的集成，但可以与任何库、算法或部署工具一起使用。它被设计为可扩展的，因此您可以编写插件来支持新的工作流程、库和工具。

在 LangChain 集成的上下文中，MLflow 提供了以下功能

• 实验跟踪：MLflow 跟踪并存储 LangChain 实验中的工件，包括模型、代码、提示、指标等。
• 依赖管理：MLflow 自动记录模型依赖项，确保开发环境和生产环境之间的一致性。
• 模型评估：MLflow 提供评估 LangChain 应用程序的本机功能。
• 跟踪：MLflow 允许您可视化跟踪数据流经 LangChain 链、代理、检索器或其他组件的过程。

注意：跟踪功能仅在 MLflow 2.14.0 及更高版本中可用。

此笔记本演示如何使用 MLflow 跟踪您的 LangChain 实验。有关此功能的更多信息以及使用 LangChain 和 MLflow 的教程和示例，请参阅MLflow LangChain 集成文档。

设置

安装 MLflow Python 包

%pip install google-search-results num

%pip install mlflow -qU

此示例使用 OpenAI LLM。 如果需要，您可以跳过以下命令并使用其他 LLM。

%pip install langchain-openai -qU

import os

# Set MLflow tracking URI if you have MLflow Tracking Server running
os.environ["MLFLOW_TRACKING_URI"] = ""
os.environ["OPENAI_API_KEY"] = ""

首先，让我们创建一个专用的 MLflow 实验，以便跟踪我们的模型和工件。 虽然您可以选择跳过此步骤并使用默认实验，但我们强烈建议将您的运行和工件组织到单独的实验中，以避免混乱并保持清晰、结构化的工作流程。

import mlflow

mlflow.set_experiment("LangChain MLflow Integration")

概述

使用以下方法之一将 MLflow 与您的 LangChain 应用程序集成

1. 自动日志记录：使用 mlflow.langchain.autolog() 命令启用无缝跟踪，这是我们推荐的首选方案，用于利用 LangChain MLflow 集成。
2. 手动日志记录：使用 MLflow API 记录 LangChain 链和代理，从而可以精细控制要在实验中跟踪的内容。
3. 自定义回调：在调用链时手动传递 MLflow 回调，从而可以对工作负载进行半自动自定义，例如跟踪特定调用。

场景 1：MLFlow 自动日志记录

要开始使用自动日志记录，只需调用 mlflow.langchain.autolog()。在此示例中，我们将 log_models 参数设置为 True，这将允许将链定义及其依赖库记录为 MLflow 模型，从而提供全面的跟踪体验。

import mlflow

mlflow.langchain.autolog(
    # These are optional configurations to control what information should be logged automatically (default: False)
    # For the full list of the arguments, refer to https://mlflow.org/docs/latest/llms/langchain/autologging.html#id1
    log_models=True,
    log_input_examples=True,
    log_model_signatures=True,
)

定义链

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant that translates {input_language} to {output_language}.",
        ),
        ("human", "{input}"),
    ]
)
parser = StrOutputParser()

chain = prompt | llm | parser

API 参考：StrOutputParser | ChatPromptTemplate | ChatOpenAI

调用链

请注意，此步骤可能比平时多花费几秒钟，因为 MLflow 在后台运行多个后台任务，将模型、跟踪和工件记录到跟踪服务器。

test_input = {
    "input_language": "English",
    "output_language": "German",
    "input": "I love programming.",
}

chain.invoke(test_input)

'Ich liebe das Programmieren.'

花点时间浏览 MLflow 跟踪 UI，您可以在其中更深入地了解正在记录哪些信息。

• 跟踪 - 导航到实验中的“跟踪”选项卡，然后单击第一行的请求 ID 链接。显示的跟踪树可视化了链调用的调用堆栈，使您可以深入了解每个组件如何在链中执行。
• MLflow 模型 - 由于我们设置了 log_model=True，MLflow 会自动创建一个 MLflow 运行来跟踪您的链定义。导航到最新的“运行”页面并打开“工件”选项卡，该选项卡列出了作为 MLflow 模型记录的文件工件，包括依赖项、输入示例、模型签名等。

调用已记录的链

接下来，让我们加载模型并验证我们是否可以重现相同的预测，从而确保一致性和可靠性。

有两种加载模型的方法

1. mlflow.langchain.load_model(MODEL_URI) - 这会将模型加载为原始 LangChain 对象。
2. mlflow.pyfunc.load_model(MODEL_URI) - 这会将模型加载到 PythonModel 包装器中，并使用 predict() API 封装预测逻辑，其中包含其他逻辑，例如模式强制。

# Replace YOUR_RUN_ID with the Run ID displayed on the MLflow UI
loaded_model = mlflow.langchain.load_model("runs:/{YOUR_RUN_ID}/model")
loaded_model.invoke(test_input)

'Ich liebe Programmieren.'

pyfunc_model = mlflow.pyfunc.load_model("runs:/{YOUR_RUN_ID}/model")
pyfunc_model.predict(test_input)

['Ich liebe das Programmieren.']

配置自动日志记录

mlflow.langchain.autolog() 函数提供了几个参数，允许对记录到 MLflow 的工件进行细粒度控制。有关可用配置的完整列表，请参阅最新的MLflow LangChain 自动日志记录文档。

场景 2：从代码手动记录 Agent

前提条件

此示例使用 SerpAPI，这是一个搜索引擎 API，作为 Agent 检索 Google 搜索结果的工具。LangChain 原生集成了 SerpAPI，只需一行代码即可为您的 Agent 配置该工具。

开始使用

• 通过 pip 安装所需的 Python 包：pip install google-search-results numexpr。
• 在 SerpAPI 官方网站 创建一个帐户并检索 API 密钥。
• 在环境变量中设置 API 密钥：os.environ["SERPAPI_API_KEY"] = "YOUR_API_KEY"

定义 Agent

在此示例中，我们将把 Agent 定义作为代码记录，而不是直接提供 Python 对象并将其保存为序列化格式。这种方法有几个优点

1. 无需序列化：通过将模型保存为代码，我们避免了序列化的需要，当处理本身不支持序列化的组件时，序列化可能会有问题。这种方法也消除了在不同环境中反序列化模型时出现不兼容问题的风险。
2. 更好的透明度：通过检查保存的代码文件，您可以深入了解模型的工作原理。这与诸如 pickle 之类的序列化格式形成对比，在 pickle 中，模型的行为保持不透明，直到它被加载回来，可能会暴露诸如远程代码执行之类的安全风险。

首先，创建一个单独的 .py 文件来定义 Agent 实例。

为了节省时间，您可以运行以下单元格来生成一个 Python 文件 agent.py，其中包含 Agent 定义代码。在实际的开发场景中，您会在另一个笔记本或手工编写的 Python 脚本中定义它。

script_content = """
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain_openai import ChatOpenAI
import mlflow

llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
tools = load_tools(["serpapi", "llm-math"], llm=llm)
agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)

# IMPORTANT: call set_model() to register the instance to be logged.
mlflow.models.set_model(agent)
"""

with open("agent.py", "w") as f:
    f.write(script_content)

API 参考：AgentType | initialize_agent | load_tools | ChatOpenAI

记录 Agent

返回原始笔记本并运行以下单元格以记录您在 agent.py 文件中定义的 Agent。

question = "How long would it take to drive to the Moon with F1 racing cars?"

with mlflow.start_run(run_name="search-math-agent") as run:
    info = mlflow.langchain.log_model(
        lc_model="agent.py",  # Specify the relative code path to the agent definition
        artifact_path="model",
        input_example=question,
    )

print("The agent is successfully logged to MLflow!")

The agent is successfully logged to MLflow!

现在，打开 MLflow UI 并导航到“运行详细信息”页面中的“工件”选项卡。您应该看到 agent.py 文件已成功记录，以及其他模型工件，例如依赖项、输入示例等。

调用已记录的 Agent

现在加载 Agent 并调用它。有两种加载模型的方法

# Let's turn on the autologging with default configuration, so we can see the trace for the agent invocation.
mlflow.langchain.autolog()

# Load the model back
agent = mlflow.pyfunc.load_model(info.model_uri)

# Invoke
agent.predict(question)

Downloading artifacts: 100%|██████████| 10/10 [00:00<00:00, 331.57it/s]

['It would take approximately 1194.5 hours to drive to the Moon with an F1 racing car.']

导航到实验中的“跟踪”选项卡，然后单击第一行的请求 ID 链接。该跟踪可视化了 Agent 在单个预测调用中如何执行多个任务

1. 确定回答问题所需的子任务。
2. 搜索 F1 赛车的速度。
3. 搜索地球到月球的距离。
4. 使用 LLM 计算除法。

场景 3. 使用 MLflow 回调

MLflow 回调提供了一种半自动化的方法来在 MLflow 中跟踪您的 LangChain 应用程序。有两个主要的回调可用

1. MlflowLangchainTracer: 主要用于生成跟踪，可在 mlflow >= 2.14.0 中使用。
2. MLflowCallbackHandler: 将指标和工件记录到 MLflow 跟踪服务器。

MlflowLangchainTracer

当使用 MlflowLangchainTracer 回调调用链或 Agent 时，MLflow 将自动生成调用堆栈的跟踪并将其记录到 MLflow 跟踪服务器。结果与 mlflow.langchain.autolog() 完全相同，但当您只想跟踪特定调用时，这尤其有用。另一方面，自动记录适用于同一笔记本/脚本中的所有调用。

from mlflow.langchain.langchain_tracer import MlflowLangchainTracer

mlflow_tracer = MlflowLangchainTracer()

# This call generates a trace
chain.invoke(test_input, config={"callbacks": [mlflow_tracer]})

# This call does not generate a trace
chain.invoke(test_input)

在哪里传递回调

LangChain 支持两种传递回调实例的方式：（1）请求时间回调 - 将其传递给 invoke 方法或与 with_config() 绑定；（2）构造函数回调 - 在链构造函数中设置它们。当使用 MlflowLangchainTracer 作为回调时，您必须使用请求时间回调。将其设置在构造函数中只会将回调应用于顶级对象，从而阻止其传播到子组件，从而导致不完整的跟踪。有关此行为的更多信息，请参阅 回调文档 了解更多详细信息。

# OK
chain.invoke(test_input, config={"callbacks": [mlflow_tracer]})
chain.with_config(callbacks=[mlflow_tracer])
# NG
chain = TheNameOfSomeChain(callbacks=[mlflow_tracer])

支持的方法

MlflowLangchainTracer 支持来自 Runnable 接口 的以下调用方法。

• 标准接口：invoke、stream、batch
• 异步接口：astream、ainvoke、abatch、astream_log、astream_events

不保证其他方法完全兼容。

MlflowCallbackHandler

MlflowCallbackHandler 是驻留在 LangChain Community 代码库中的回调处理程序。

可以为链/Agent 调用传递此回调，但必须通过调用 flush_tracker() 方法显式完成。

当使用回调调用链时，它会执行以下操作

1. 创建一个新的 MLflow 运行，或者如果活动的 MLflow 实验中存在运行，则检索活动的运行。
2. 记录诸如 LLM 调用次数、令牌使用情况和其他相关指标的度量。如果链/代理包含 LLM 调用且您已安装 spacy 库，它还会记录文本复杂性指标，例如 flesch_kincaid_grade。
3. 将内部步骤记录为 JSON 文件（这是跟踪的旧版本）。
4. 将链的输入和输出记录为 Pandas Dataframe。
5. 使用链/代理实例调用 flush_tracker() 方法，将链/代理记录为 MLflow 模型。

from langchain_community.callbacks import MlflowCallbackHandler

mlflow_callback = MlflowCallbackHandler()

chain.invoke("What is LangChain callback?", config={"callbacks": [mlflow_callback]})

mlflow_callback.flush_tracker()

API 参考：MlflowCallbackHandler

参考资料

要了解有关该功能的更多信息并访问使用 LangChain 和 MLflow 的教程和示例，请参阅 MLflow 关于 LangChain 集成的文档。

MLflow 还提供了几个关于 LangChain 集成的 教程 和 示例

• 快速入门
• RAG 教程
• 代理示例