跳到主要内容

使用模型记录提示

提示通常用作 GenAI 应用程序的一部分。管理提示与模型之间的关联对于跟踪模型的演变和确保不同环境之间的一致性至关重要。MLflow 提示注册表与 MLflow 的模型跟踪功能集成，允许您跟踪模型和应用程序使用了哪些提示（以及版本）。

基本用法​

要使用关联的提示记录模型，请在 log_model 方法中使用 prompts 参数。prompts 参数接受与模型关联的提示 URL 列表或提示对象。关联的提示会显示在模型的 MLflow UI 运行中。

import mlflow

with mlflow.start_run():
    mlflow.<flavor>.log_model(
        model,
        ...
        # Specify a list of prompt URLs or prompt objects.
        prompts=["prompts:/summarization-prompt/2"]
    )
警告

用于将提示与模型关联的 prompts 参数仅支持 GenAI 风格，例如 OpenAI、LangChain、LlamaIndex、DSPy 等。有关支持的风格的完整列表，请参阅 GenAI 风格。

示例 1：使用 LangChain 记录提示​

1. 创建提示​

如果您尚未创建提示，请按照此页面中的说明创建新提示。

2. 使用注册的提示定义链​

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# Load registered prompt
prompt = mlflow.genai.load_prompt("prompts:/summarization-prompt/2")

# Create LangChain prompt object
langchain_prompt = ChatPromptTemplate.from_messages(
    [
        (
            # IMPORTANT: Convert prompt template from double to single curly braces format
            "system",
            prompt.to_single_brace_format(),
        ),
        ("placeholder", "{messages}"),
    ]
)

# Define the LangChain chain
llm = ChatOpenAI()
chain = langchain_prompt | llm

# Invoke the chain
response = chain.invoke({"num_sentences": 1, "sentences": "This is a test sentence."})
print(response)

3. 将链记录到 MLflow​

然后将链记录到 MLflow，并在 prompts 参数中指定提示 URL

with mlflow.start_run(run_name="summarizer-model"):
    mlflow.langchain.log_model(
        chain, name="model", prompts=["prompts:/summarization-prompt/2"]
    )

现在您可以在 MLflow UI 中查看与模型关联的提示

Associated Prompts

此外，您可以在提示详细信息页面中查看使用特定提示的模型（运行）列表

Associated Prompts

示例 2：使用 Models-from-Code 自动记录提示​

Models-from-Code 是一项允许您在代码中定义和记录模型的功能。使用代码记录模型带来了诸多好处，例如可移植性、可读性、避免序列化等等。

结合 MLflow 提示注册表，该功能解锁了管理提示版本的更多灵活性。值得注意的是，如果您的模型代码使用 MLflow 提示注册表中的提示，MLflow 会自动将其与模型一起记录。

在以下示例中，我们使用 LangGraph 定义了一个使用注册提示的非常简单的聊天机器人。

1. 创建提示​

import mlflow

# Register a new prompt
prompt = mlflow.genai.register_prompt(
    name="chat-prompt",
    template="You are an expert in programming. Please answer the user's question about programming.",
)

2. 使用注册提示定义图​

创建包含以下内容的 Python 脚本 chatbot.py。

提示

如果您正在使用 Jupyter notebook，您可以取消注释 %writefile magic 命令并在单元格中运行以下代码以生成脚本。

# %%writefile chatbot.py

import mlflow
from typing import Annotated
from typing_extensions import TypedDict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages


class State(TypedDict):
    messages: list


llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.1)
system_prompt = mlflow.genai.load_prompt("prompts:/chat-prompt/1")


def add_system_message(state: State):
    return {
        "messages": [
            {
                "role": "system",
                "content": system_prompt.to_single_brace_format(),
            },
            *state["messages"],
        ]
    }


def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}


graph_builder = StateGraph(State)
graph_builder.add_node("add_system_message", add_system_message)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "add_system_message")
graph_builder.add_edge("add_system_message", "chatbot")
graph_builder.add_edge("chatbot", END)

graph = graph_builder.compile()

mlflow.models.set_model(graph)

3. 将图记录到 MLflow​

在 model 参数中指定脚本的文件路径

with mlflow.start_run():
    model_info = mlflow.langchain.log_model(
        lc_model="./chatbot.py",
        name="graph",
    )

我们这次没有指定 prompts 参数，但 MLflow 会自动将脚本中加载的提示记录到已记录的模型中。现在您可以在 MLflow UI 中查看关联的提示

Associated Prompts

4. 重新加载图并调用​

最后，让我们重新加载图并调用它，看看聊天机器人的实际运行情况。

# Enable MLflow tracing for LangChain to view the prompt passed to LLM.
mlflow.langchain.autolog()

# Load the graph
graph = mlflow.langchain.load_model(model_info.model_uri)

graph.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is the difference between multi-threading and multi-processing?",
            }
        ]
    }
)

Chatbot