跳到主要内容

使用模型记录提示

提示词通常用作 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：使用代码模型自动记录提示词​

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

与 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 魔术命令并在单元格中运行以下代码以生成脚本。

# %%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