MLflow LlamaIndex Flavor

简介

LlamaIndex 🦙 是一个强大的数据中心化框架，旨在将自定义数据源无缝连接到大语言模型 (LLM) 上。它提供了一套完整的数据结构和工具，简化了为 LLM 使用而摄取、结构化和访问私有或特定领域数据的过程。LlamaIndex 通过提供高效的索引和检索机制，在实现上下文感知 AI 应用方面表现出色，使构建高级问答系统、聊天机器人以及其他需要集成外部知识的 AI 驱动应用变得更加容易。

为什么要将 LlamaIndex 与 MLflow 结合使用？

LlamaIndex 库与 MLflow 的集成，为管理和部署 LlamaIndex 引擎提供了无缝体验。以下是将 LlamaIndex 与 MLflow 结合使用的主要优势：

• MLflow Tracking 允许您在 MLflow 中跟踪索引，并管理 LlamaIndex 项目中众多的组成部分，如提示词 (prompts)、LLM、工作流 (workflows)、工具、全局配置等。
• MLflow Model 将您的 LlamaIndex 索引/引擎/工作流及其所有依赖版本、输入输出接口以及其他重要元数据打包在一起。这使您能够轻松部署 LlamaIndex 模型进行推理，并确保在 AI 开发工作流的不同阶段环境保持一致。
• MLflow Evaluate 在 MLflow 中提供了评估 LLM 应用和 AI 智能体的原生能力。此功能有助于高效评估 LlamaIndex 模型的推理结果，确保稳健的性能分析并促进快速迭代。
• MLflow Tracing 是一款功能强大的可观测性工具，用于监控和调试 LlamaIndex 模型内部发生的情况，帮助您快速识别潜在的瓶颈或问题。凭借其强大的自动记录功能，您无需添加任何额外代码，只需运行一条命令即可检测 LlamaIndex 应用。

入门

在这些入门教程中，您将学习 LlamaIndex 的最基本组件，以及如何利用与 MLflow 的集成，为您的 LlamaIndex 应用带来更好的可维护性和可观测性。

LlamaIndex 工作流与 MLflow

通过构建一个简单的智能体工作流 (agentic Workflow)，开始使用 MLflow 和 LLamaIndex。了解如何记录和加载工作流以进行推理，以及如何启用跟踪 (tracing) 以实现可观测性。

使用 MLflow 构建索引

通过探索 VectorStoreIndex 最简单的索引配置，开始使用 MLflow 和 LlamaIndex。

概念

注意

工作流集成仅适用于 LlamaIndex >= 0.11.0 和 MLflow >= 2.17.0。

Workflow (工作流) 🆕

Workflow 是 LlamaIndex 的事件驱动编排框架。它被设计为一个灵活且可解释的框架，用于构建各种 LLM 应用，例如智能体、RAG 流程、数据提取流水线等。MLflow 支持对 Workflow 对象进行跟踪、评估和溯源，这使得它们更具可观测性和可维护性。

Index (索引)

Index 对象是为快速信息检索而索引的文档集合，为检索增强生成 (RAG) 和智能体等应用提供支持。Index 对象可以直接记录到 MLflow 运行中，并重新加载以用作推理引擎。

Engine (引擎)

Engine 是构建在 Index 对象之上的通用接口，提供了一组与索引交互的 API。LlamaIndex 提供两种类型的引擎：QueryEngine 和 ChatEngine。QueryEngine 只需接收单个查询并根据索引返回响应。ChatEngine 专为会话智能体设计，同时会跟踪会话历史记录。

Settings (设置)

Settings 对象是一个全局服务上下文，捆绑了整个 LlamaIndex 应用中常用的资源。它包括 LLM 模型、嵌入模型、回调等设置。当记录 LlamaIndex 索引/引擎/工作流时，MLflow 会跟踪 Settings 对象的状态，以便您在重新加载模型进行推理时可以轻松重现相同的结果（请注意，API 密钥、不可序列化的对象等不会被跟踪）。

用法

在 MLflow 实验中保存和加载索引

创建索引

index 对象是 LlamaIndex 与 MLflow 集成的核心。使用 LlamaIndex，您可以从文档集合或外部向量存储创建索引。以下代码使用 LlamaIndex 仓库中提供的 Paul Graham 的论文数据创建了一个示例索引。

shell

mkdir -p data
curl -L https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt -o ./data/paul_graham_essay.txt

python

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

documents = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(documents)

将索引记录到 MLflow

您可以使用 mlflow.llama_index.log_model() 函数将 index 对象记录到 MLflow 实验中。

这里的一个关键步骤是指定 engine_type 参数。引擎类型的选择不会影响索引本身，但它决定了您在重新加载模型进行推理时查询索引的接口方式。

1. QueryEngine (engine_type="query") 专为简单的查询-响应系统设计，接收单个查询字符串并返回响应。
2. ChatEngine (engine_type="chat") 专为会话智能体设计，它会跟踪会话历史记录并回复用户消息。
3. Retriever (engine_type="retriever") 是一个更底层的组件，它返回匹配查询的前 k 个相关文档。

以下代码是使用 chat 引擎类型将索引记录到 MLflow 的示例。

python

import mlflow

mlflow.set_experiment("llama-index-demo")

with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        index,
        name="index",
        engine_type="chat",
        input_example="What did the author do growing up?",
    )

注意

上述代码片段直接将索引对象传递给 log_model 函数。此方法仅适用于默认的 SimpleVectorStore 向量存储，它仅将嵌入的文档保留在内存中。如果您的索引使用外部向量存储（如 QdrantVectorStore 或 DatabricksVectorSearch），您可以使用“代码模型”(Model-from-Code) 记录方法。有关详细信息，请参阅如何使用外部向量存储记录和加载索引。

提示

在底层，MLflow 会在索引对象上调用 as_query_engine() / as_chat_engine() / as_retriever() 方法，将其转换为相应的引擎实例。

重新加载索引以进行推理

保存的索引可以使用 mlflow.pyfunc.load_model() 函数重新加载以进行推理。此函数提供了一个由 LlamaIndex 引擎支持的 MLflow Python 模型，并使用了记录时指定的引擎类型。

python

import mlflow

model = mlflow.pyfunc.load_model(model_info.model_uri)

response = model.predict("What was the first program the author wrote?")
print(response)
# >> The first program the author wrote was on the IBM 1401 ...

# The chat engine keeps track of the conversation history
response = model.predict("How did the author feel about it?")
print(response)
# >> The author felt puzzled by the first program ...

提示

若要重新加载索引本身而不是引擎，请使用 mlflow.llama_index.load_model() 函数。

python

index = mlflow.llama_index.load_model("runs:/<run_id>/index")

启用跟踪 (Tracing)

您可以通过调用 mlflow.llama_index.autolog() 函数为您的 LlamaIndex 代码启用跟踪。MLflow 会自动将 LlamaIndex 执行的输入和输出记录到活动的 MLflow 实验中，为您提供模型行为的详细视图。

python

import mlflow

mlflow.llama_index.autolog()

chat_engine = index.as_chat_engine()
response = chat_engine.chat("What was the first program the author wrote?")

然后，您可以导航到 MLflow UI，选择实验，并打开“Traces”选项卡，查看引擎所做预测的记录轨迹。看到聊天引擎如何协调并执行多项任务来回答您的问题，是非常令人印象深刻的！

您可以通过运行相同的函数并将 disable 参数设置为 True 来禁用跟踪。

python

mlflow.llama_index.autolog(disable=True)

注意

跟踪功能支持异步预测和流式响应，但不支持异步与流式的结合，例如 astream_chat 方法。

常见问题 (FAQ)

如何使用外部向量存储记录和加载索引？

如果您的索引使用默认的 SimpleVectorStore，您可以使用 mlflow.llama_index.log_model() 函数将索引直接记录到 MLflow。MLflow 会将内存中的索引数据（嵌入文档）持久化到 MLflow 工件存储中，这允许在不重新索引文档的情况下使用相同数据重新加载索引。

然而，当索引使用 DatabricksVectorSearch 和 QdrantVectorStore 等外部向量存储时，索引数据存储在远程，它们不支持本地序列化。因此，您不能直接记录带有这些存储的索引。对于这种情况，您可以使用 Model-from-Code（代码模型）记录，它提供了对索引保存过程的更多控制，并允许您使用外部向量存储。

要使用 Model-from-Code 记录，首先需要创建一个单独的 Python 文件来定义索引。如果您在 Jupyter Notebook 中，可以使用 %%writefile 魔术命令将单元格代码保存到 Python 文件中。

python

# %%writefile index.py

# Create Qdrant client with your own settings.
client = qdrant_client.QdrantClient(
    host="localhost",
    port=6333,
)

# Here we simply load vector store from the existing collection to avoid
# re-indexing documents, because this Python file is executed every time
# when the model is loaded. If you don't have an existing collection, create
# a new one by following the official tutorial:
# https://docs.llamaindex.org.cn/en/stable/examples/vector_stores/QdrantIndexDemo/
vector_store = QdrantVectorStore(client=client, collection_name="my_collection")
index = VectorStoreIndex.from_vector_store(vector_store=vector_store)

# IMPORTANT: call set_model() method to tell MLflow to log this index
mlflow.models.set_model(index)

然后，通过将 Python 文件路径传递给 mlflow.llama_index.log_model() 函数来记录索引。全局 Settings 对象会作为模型的一部分正常保存。

python

import mlflow

with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        "index.py",
        name="index",
        engine_type="query",
    )

记录的索引可以使用 mlflow.llama_index.load_model() 或 mlflow.pyfunc.load_model() 函数重新加载，方式与本地索引相同。

python

index = mlflow.llama_index.load_model(model_info.model_uri)
index.as_query_engine().query("What is MLflow?")

注意

传递给 set_model() 方法的对象必须是与记录期间指定的引擎类型兼容的 LlamaIndex 索引。未来版本将添加对更多对象的支持。

如何记录和加载 LlamaIndex 工作流？

MLflow 通过 Model-from-Code 功能支持记录和加载 LlamaIndex 工作流。有关记录和加载 LlamaIndex 工作流的详细示例，请参阅 LlamaIndex Workflows with MLflow 笔记本。

python

import mlflow

with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        "/path/to/workflow.py",
        name="model",
        input_example={"input": "What is MLflow?"},
    )

记录的工作流可以使用 mlflow.llama_index.load_model() 或 mlflow.pyfunc.load_model() 函数重新加载。

python

# Use mlflow.llama_index.load_model to load the workflow object as is
workflow = mlflow.llama_index.load_model(model_info.model_uri)
await workflow.run(input="What is MLflow?")

# Use mlflow.pyfunc.load_model to load the workflow as a MLflow Pyfunc Model
# with standard inference APIs for deployment and evaluation.
pyfunc = mlflow.pyfunc.load_model(model_info.model_uri)
pyfunc.predict({"input": "What is MLflow?"})

警告

MLflow PyFunc 模型不支持异步推理。当您使用 mlflow.pyfunc.load_model() 加载工作流时，predict 方法变为同步，并将阻塞直到工作流执行完成。当使用 MLflow Deployment 或 Databricks Model Serving 将记录的 LlamaIndex 工作流部署到生产端点时，这也适用。

我有一个以 query 引擎类型记录的索引。我可以将其作为 chat 引擎加载回来吗？

虽然无法在原地更新记录模型的引擎类型，但您可以随时加载索引并以所需的引擎类型重新记录它。此过程不需要重新创建索引，因此是在不同引擎类型之间切换的有效方法。

python

import mlflow

# Log the index with the query engine type first
with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        index,
        name="index-query",
        engine_type="query",
    )

# Load the index back and re-log it with the chat engine type
index = mlflow.llama_index.load_model(model_info.model_uri)
with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(
        index,
        name="index-chat",
        # Specify the chat engine type this time
        engine_type="chat",
    )

或者，您可以利用它们在加载的 LlamaIndex 原生索引对象上的标准推理 API，具体如下：

• index.as_chat_engine().chat("hi")
• index.as_query_engine().query("hi")
• index.as_retriever().retrieve("hi")

如何使用不同的 LLM 与加载的引擎进行推理？

将索引保存到 MLflow 时，它会持久化全局 Settings 对象作为模型的一部分。该对象包含引擎要使用的 LLM 和嵌入模型等设置。

python

import mlflow
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI

Settings.llm = OpenAI("gpt-4o-mini")

# MLflow saves GPT-4o-Mini as the LLM to use for inference
with mlflow.start_run():
    model_info = mlflow.llama_index.log_model(index, name="index", engine_type="chat")

稍后当您重新加载索引时，持久化的设置也会在全局范围内应用。这意味着加载的引擎将使用与记录时相同的 LLM。

但是，有时您可能希望使用不同的 LLM 进行推理。在这种情况下，您可以在加载索引后直接更新全局 Settings 对象。

python

import mlflow

# Load the index back
loaded_index = mlflow.llama_index.load_model(model_info.model_uri)

assert Settings.llm.model == "gpt-4o-mini"


# Update the settings to use GPT-4 instead
Settings.llm = OpenAI("gpt-4")
query_engine = loaded_index.as_query_engine()
response = query_engine.query("What is the capital of France?")