MLflow Langchain 自动记录

MLflow LangChain 风格支持自动记录,这是一个强大的功能,允许您在不使用显式日志语句的情况下记录关于 LangChain 模型和执行的关键细节。MLflow LangChain 自动记录涵盖了模型的各个方面,包括轨迹、模型、签名等。

注意

MLflow 的 LangChain Autologging 功能在 MLflow 2.14.0 版本中进行了全面改进。如果您使用的是较早版本的 MLflow,请参阅 此处 的旧版文档以获取适用的 autologging 文档。

备注

MLflow LangChain Autologging 已验证与 LangChain 版本 0.1.0 到 0.2.3 兼容。超出此范围,该功能可能无法按预期工作。要安装兼容版本的 LangChain,请运行以下命令:

pip install mlflow[langchain] --upgrade

快速入门

要为 LangChain 模型启用自动日志记录,请在脚本或笔记本的开头调用 mlflow.langchain.autolog()。这将默认自动记录跟踪,以及在您明确启用它们时记录其他工件,如模型、输入示例和模型签名。有关配置的更多信息,请参阅 配置自动日志记录 部分。

import mlflow

mlflow.langchain.autolog()

# Enable other optional logging
# mlflow.langchain.autolog(log_models=True, log_input_examples=True)

# Your LangChain model code here
...

一旦你调用了链,你可以在 MLflow UI 中查看记录的轨迹和工件。

通过 autolog 进行 LangChain 跟踪

配置自动日志记录

MLflow LangChain 自动记录可以记录有关模型及其推理的各种信息。默认情况下,仅启用跟踪日志记录,但您可以通过在调用 mlflow.langchain.autolog() 时设置相应的参数来启用其他信息的自动记录。有关其他配置,请参阅 API 文档。

目标

默认

参数

描述

跟踪

true

log_traces

是否为模型生成并记录跟踪。有关跟踪功能的更多详细信息,请参阅 MLflow 跟踪

模型工件

false

log_models

如果设置为 True,LangChain 模型在调用时将被记录。支持的模型包括 ChainAgentExecutorBaseRetrieverSimpleChatModelChatPromptTemplate 以及 Runnable 类型的子集。请参阅 MLflow 仓库 获取支持模型的完整列表。

模型签名

false

log_model_signatures

如果设置为 True,在推理过程中,描述模型输入和输出的 模型签名 将与 Langchain 模型工件一起收集和记录。此选项仅在 log_models 启用时可用。

输入示例

false

log_input_examples

如果设置为 True ,推理数据中的输入示例将与 LangChain 模型工件一起在推理期间收集并记录。此选项仅在 log_models 启用时可用。

输入与输出 (已弃用)

false

log_inputs_outputs

如果设置为 True,模型调用时将记录输入和输出。此功能已弃用,并将在未来移除。请改用 log_traces

例如,要禁用跟踪日志记录,并改为启用模型日志记录,请运行以下代码:

import mlflow

mlflow.langchain.autolog(
    log_traces=False,
    log_models=True,
)

备注

MLflow 不支持包含检索器的链的自动模型记录。保存检索器需要额外的 loader_fnpersist_dir 信息来加载模型。如果你想记录带有检索器的模型,请按照 retriever_chain 示例中的方法手动记录模型。

LangChain Autologging 的示例代码

import os
from operator import itemgetter

from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnableLambda

import mlflow

# Uncomment the following to use the full abilities of langchain autologgin
# %pip install `langchain_community>=0.0.16`
# These two libraries enable autologging to log text analysis related artifacts
# %pip install textstat spacy

assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."

# Enable mlflow langchain autologging
# Note: We only support auto-logging models that do not contain retrievers
mlflow.langchain.autolog(
    log_input_examples=True,
    log_model_signatures=True,
    log_models=True,
    log_inputs_outputs=True,
    registered_model_name="lc_model",
)

prompt_with_history_str = """
Here is a history between you and a human: {chat_history}

Now, please answer this question: {question}
"""
prompt_with_history = PromptTemplate(
    input_variables=["chat_history", "question"], template=prompt_with_history_str
)


def extract_question(input):
    return input[-1]["content"]


def extract_history(input):
    return input[:-1]


llm = OpenAI(temperature=0.9)

# Build a chain with LCEL
chain_with_history = (
    {
        "question": itemgetter("messages") | RunnableLambda(extract_question),
        "chat_history": itemgetter("messages") | RunnableLambda(extract_history),
    }
    | prompt_with_history
    | llm
    | StrOutputParser()
)

inputs = {"messages": [{"role": "user", "content": "Who owns MLflow?"}]}

print(chain_with_history.invoke(inputs))
# sample output:
# "1. Databricks\n2. Microsoft\n3. Google\n4. Amazon\n\nEnter your answer: 1\n\n
# Correct! MLflow is an open source project developed by Databricks. ...

# We automatically log the model and trace related artifacts
# A model with name `lc_model` is registered, we can load it back as a PyFunc model
model_name = "lc_model"
model_version = 1
loaded_model = mlflow.pyfunc.load_model(f"models:/{model_name}/{model_version}")
print(loaded_model.predict(inputs))

工作原理

MLflow LangChain Autologging 使用两种方式记录跟踪和其他工件。通过 LangChain 的 Callbacks 框架实现跟踪。其他工件通过修补支持模型的 调用函数 来记录。在典型场景中,您不需要关心内部实现细节,但本节简要概述了其工作原理。

MLflow 跟踪回调

MlflowLangchainTracer 是一个回调处理器,它被注入到 langchain 模型推理过程中以自动记录跟踪。它在链的一系列动作如 on_chain_starton_llm_start 开始时启动一个新的跨度,并在动作完成时结束它。各种元数据如跨度类型、动作名称、输入、输出、延迟等会自动记录到跨度中。

自定义回调

有时你可能希望自定义在跟踪中记录的信息。你可以通过创建一个继承自 MlflowLangchainTracer 的自定义回调处理程序来实现这一点。以下示例展示了如何在聊天模型开始运行时记录一个额外的属性到跨度中。

from mlflow.langchain.langchain_tracer import MlflowLangchainTracer

class CustomLangchainTracer(MlflowLangchainTracer):

    # Override the handler functions to customize the behavior. The method signature is defined by LangChain Callbacks.
    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        tags: Optional[List[str]] = None,
        parent_run_id: Optional[UUID] = None,
        metadata: Optional[Dict[str, Any]] = None,
        name: Optional[str] = None,
        **kwargs: Any,
    ):
        """Run when a chat model starts running."""
        attributes = {
            **kwargs,
            **metadata,
            # Add additional attribute to the span
            "version": "1.0.0",
        }

        # Call the _start_span method at the end of the handler function to start a new span.
        self._start_span(
            span_name=name or self._assign_span_name(serialized, "chat model"),
            parent_run_id=parent_run_id,
            span_type=SpanType.CHAT_MODEL,
            run_id=run_id,
            inputs=messages,
            attributes=kwargs,
        )

日志工件的补丁函数

其他工件,如模型,通过修补受支持模型的调用函数来记录,以插入日志调用。MLflow 修补了以下函数:

  • invoke

  • batch

  • stream

  • get_relevant_documents (用于检索器)

  • __call__ (用于链和代理执行器)

  • ainvoke

  • abatch

  • astream

警告

MLflow 支持异步函数(例如 ainvokeabatchastream)的自动日志记录,然而,日志记录操作本身不是异步的,可能会阻塞主线程。调用函数本身仍然是非阻塞的,并返回一个协程对象,但日志记录的开销可能会减慢模型推理过程。在使用带有自动日志记录的异步函数时,请注意这一副作用。

故障排除

如果你在使用 MLflow LangChain 风格时遇到任何问题,请参考 FAQ <../index.html#faq>。如果你仍有疑问,请随时在 MLflow Github 仓库 中提出问题。

如何在自动记录期间抑制警告消息?

MLflow Langchain Autologging 在幕后调用各种日志记录函数和 LangChain 工具。其中一些可能会生成对 autologging 过程不重要的警告消息。如果你想抑制这些警告消息,请将 silent=True 传递给 mlflow.langchain.autolog() 函数。

import mlflow

mlflow.langchain.autolog(silent=True)

# No warning messages will be emitted from autologging

我无法加载由 mlflow langchain autologging 记录的模型

MLflow LangChain 自动记录不支持原生保存或加载的几种模型类型。

  • 模型包含 langchain 检索器

    LangChain 检索器不支持 MLflow 自动日志记录。如果你的模型包含一个检索器,你需要使用 mlflow.langchain.log_model API 手动记录模型。由于加载这些模型需要指定 loader_fnpersist_dir 参数,请查看 retriever_chain 中的示例。

  • 无法序列化某些对象

    对于LangChain不支持原生存储或加载的某些模型,我们将在保存对象时使用pickle。由于此功能,您的cloudpickle版本在保存和加载环境中必须一致,以确保对象引用正确解析。为了进一步保证对象表示的正确性,您应确保您的环境中安装了至少版本为2的`pydantic`。

旧版本文档

MLflow LangChain Autologging 功能在 MLflow 2.14.0 中得到了大幅更新。如果你使用的是 MLflow 的早期版本,请参考以下文档。

备注

要使用 MLflow LangChain 自动日志记录,请将 langchain 升级到 0.1.0 版本 或更高版本。根据您现有的环境,您可能需要手动安装 langchain_community>=0.0.16 以启用工件和指标的自动日志记录。(此行为将在未来修改为可选导入)如果自动日志记录未按预期记录工件,请检查 stdout 日志中的警告消息。对于 langchain_community==0.0.16,您需要手动安装 textstatspacy 库,并重新启动任何活动的交互式环境(例如,笔记本环境)。在 Databricks 上,您可以通过执行 dbutils.library.restartPython() 来强制 Python REPL 重新启动,从而使新安装的库可用。

MLflow langchain autologging 将 MlflowCallbackHandler 注入到 langchain 模型推理过程中,以自动记录指标和工件。只有在调用 mlflow.langchain.autolog() 时将 log_models 设置为 True,并且被调用的对象属于以下支持的模型类型时,才会记录模型:ChainAgentExecutorBaseRetrieverRunnableSequenceRunnableParallelRunnableBranchSimpleChatModelChatPromptTemplateRunnableLambdaRunnablePassthrough。未来将支持更多模型类型。

备注

我们对所有支持的 langchain 模型的 invoke 函数、Chains 和 AgentExecutors 模型的 __call__ 函数,以及 BaseRetrievers 的 get_relevant_documents 函数进行了修补,以便仅在这些函数被调用时,MLflow 才会自动记录指标和工件。如果模型包含检索器,我们不支持自动记录该模型,因为它需要保存 loader_fnpersist_dir 以便加载模型。如果您想记录带有检索器的模型,请手动记录模型。

以下指标和工件默认会被记录(取决于所涉及的模型):

工件:

工件名称

解释

table_action_records.html

每个动作的详细信息,包括链、工具、llms、代理、检索器。

table_session_analysis.html

每个提示步骤的提示和输出详情;令牌使用情况;文本分析指标

chat_html.html

LLM 输入和输出细节

llm_start_x_prompt_y.json

包含在llm generate 调用期间传递的提示和kwargs

llm_end_x_generation_y.json

包含 LLM 结果的 llm_output

ent-<生成文本的哈希字符串>.html

使用 spacy 的 “en_core_web_sm” 模型以 ent 样式生成文本的可视化(如果已安装 spacy 并且已下载模型)

dep-<生成文本的哈希字符串>.html

使用 spacy 的 “en_core_web_sm” 模型以 dep 样式生成文本的可视化(如果已安装 spacy 并且已下载模型)

llm_new_tokens_x.json

记录在推理过程中添加到LLM的新令牌

chain_start_x.json

在推理过程中记录输入和链相关信息

chain_end_x.json

记录链输出

tool_start_x.json

记录工具的名称,推理过程中的描述信息

tool_end_x.json

记录工具的观察结果

retriever_start_x.json

记录推理过程中检索器的信息

retriever_end_x.json

记录检索器的结果文档

agent_finish_x.json

记录 ActionAgent 的最终返回值,包括输出和日志

agent_action_x.json

记录 ActionAgent 的操作细节

on_text_x.json

在推理过程中记录文本

inference_inputs_outputs.json

每次推理调用的输入和输出细节(默认记录,可以通过在启用autolog时设置`log_inputs_outputs=False`来关闭)

指标:

度量类型

详情

基本指标

步骤, 开始, 结束, 错误, 文本计数器, 链开始, 链结束, llm开始, llm结束, llm流, 工具开始, 工具结束, 代理结束, 检索器结束, 检索器开始 (它们是每个组件调用的计数)

文本分析指标

flesch_reading_ease, flesch_kincaid_grade, smog_index, coleman_liau_index, automated_readability_index, dale_chall_readability_score, difficult_words, linsear_write_formula, gunning_fog, fernandez_huerta, szigriszt_pazos, gutierrez_polini, crawford, gulpease_index, osman (如果安装了 textstat 库,它们是生成文本的文本分析指标)

备注

每次推理调用都会将这些工件记录到一个名为 artifacts-<session_id>-<idx> 的单独目录中,其中 session_id 是随机生成的 uuid,idx 是推理调用的索引。session_id 也保存在 inference_inputs_outputs.json 文件中,因此您可以轻松找到每次推理调用的相应工件。