使用代码模型构建LangGraph
在本博客中，我们将指导您使用MLflow创建一个LangGraph聊天机器人。通过将MLflow与LangGraph创建和管理循环图的能力相结合，您可以以可扩展的方式创建强大的有状态、多智能体应用程序。
在这篇文章中，我们将展示如何利用MLflow的功能创建一个可序列化和可服务的MLflow模型，该模型可以轻松地在各种服务器上进行跟踪、版本控制和部署。我们将使用langchain flavor结合MLflow的model from code功能。
什么是LangGraph？
LangGraph 是一个用于构建基于LLM的有状态多参与者应用的库，用于创建智能体和多智能体工作流。与其他LLM框架相比，它具有以下核心优势：
- 循环与分支: 在你的应用中实现循环和条件判断。
- 持久性: 在图的每个步骤后自动保存状态。随时暂停和恢复图执行，以支持错误恢复、人在回路工作流、时间旅行等功能。
- 人工介入循环: 中断图执行以批准或编辑智能体计划的下一个操作。
- 流式支持: 当每个节点产生输出时进行流式传输（包括令牌流式传输）。
- 与LangChain的集成: LangGraph与LangChain无缝集成。
LangGraph 允许您定义包含循环的流程，这对于大多数智能体架构至关重要，使其区别于基于DAG的解决方案。作为一个非常底层的框架，它提供了对应用程序流程和状态的细粒度控制，这对于创建可靠的智能体至关重要。此外，LangGraph 包含内置持久化功能，支持高级人机交互和记忆特性。
LangGraph 灵感来源于 Pregel 和 Apache Beam。其公共接口设计借鉴了 NetworkX。LangGraph 由 LangChain Inc 开发，即 LangChain 的创建者，但也可独立于 LangChain 使用。
如需完整演练，请查看LangGraph 快速入门，若想深入了解 LangGraph 的设计基础，请参阅概念指南。
1 - 设置
首先，我们需要安装所需的依赖项。在这个示例中，我们将使用OpenAI作为我们的LLM，但使用LangChain与LangGraph可以轻松替换任何其他受支持的LLM或LLM提供商。
%%capture
%pip install langchain_openai==0.2.0 langchain==0.3.0 langgraph==0.2.27
%pip install -U mlflow
接下来，让我们获取相关的密钥。如LangGraph 快速入门中所示，getpass 是在交互式 jupyter 环境中插入密钥的好方法。
import os
# Set required environment variables for authenticating to OpenAI
# Check additional MLflow tutorials for examples of authentication if needed
# https://mlflow.org/docs/latest/llms/openai/guide/index.html#direct-openai-service-usage
assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."
2 - 自定义工具
虽然这是一个演示，但将可重用的实用工具分离到单独的文件/目录中是良好的实践。下面我们创建三个通用实用工具，理论上在构建额外的 MLflow + LangGraph 实现时会很有价值。
请注意，我们使用魔法命令 %%writefile 在 jupyter notebook 环境中创建一个新文件。如果您在交互式笔记本之外运行此操作，只需创建以下文件，省略 %%writefile {FILE_NAME}.py 行。
%%writefile langgraph_utils.py
# omit this line if directly creating this file; this command is purely for running within Jupyter
import os
from typing import Union
from langgraph.pregel.io import AddableValuesDict
def _langgraph_message_to_mlflow_message(
langgraph_message: AddableValuesDict,
) -> dict:
langgraph_type_to_mlflow_role = {
"human": "user",
"ai": "assistant",
"system": "system",
}
if type_clean := langgraph_type_to_mlflow_role.get(langgraph_message.type):
return {"role": type_clean, "content": langgraph_message.content}
else:
raise ValueError(f"Incorrect role specified: {langgraph_message.type}")
def get_most_recent_message(response: AddableValuesDict) -> dict:
most_recent_message = response.get("messages")[-1]
return _langgraph_message_to_mlflow_message(most_recent_message)["content"]
def increment_message_history(
response: AddableValuesDict, new_message: Union[dict, AddableValuesDict]
) -> list[dict]:
if isinstance(new_message, AddableValuesDict):
new_message = _langgraph_message_to_mlflow_message(new_message)
message_history = [
_langgraph_message_to_mlflow_message(message)
for message in response.get("messages")
]
return message_history + [new_message]
在此步骤结束时，您应该会在当前目录中看到一个名为 langgraph_utils.py 的新文件。
请注意，最佳实践是添加单元测试并将项目合理组织到逻辑结构化的目录中。
3 - 记录LangGraph模型
太棒了！现在我们已经有了一些可重用的工具函数位于./langgraph_utils.py中，我们准备好使用MLflow官方的LangGraph风味来记录模型了。
3.1 - 创建我们的模型代码文件
快速了解一些背景。MLflow 旨在将模型工件序列化到 MLflow 跟踪服务器。许多流行的机器学习包没有强大的序列化和反序列化支持，因此 MLflow 通过 models from code 功能来增强这一功能。借助代码生成模型，我们能够利用 Python 作为序列化格式，而不是像 JSON 或 pkl 这样的流行替代方案。这带来了极大的灵活性和稳定性。
要创建一个包含来自代码的模型的Python文件，我们必须执行以下步骤：
- 创建一个新的python文件。我们称之为
graph.py。 - 定义我们的langgraph图。
- 利用 mlflow.models.set_model 向 MLflow 指明 Python 脚本中哪个对象是我们感兴趣的模型。
就这样！
%%writefile graph.py
# omit this line if directly creating this file; this command is purely for running within Jupyter
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.graph.state import CompiledStateGraph
import mlflow
import os
from typing import TypedDict, Annotated
def load_graph() -> CompiledStateGraph:
"""Create example chatbot from LangGraph Quickstart."""
assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."
class State(TypedDict):
messages: Annotated[list, add_messages]
graph_builder = StateGraph(State)
llm = ChatOpenAI()
def chatbot(state: State):
return {"messages": [llm.invoke(state["messages"])]}
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()
return graph
# Set are model to be leveraged via model from code
mlflow.models.set_model(load_graph())
3.2 - 使用"来自代码的模型"进行记录
创建此实现后，我们可以利用标准 MLflow API 来记录模型。
import mlflow
with mlflow.start_run() as run_id:
model_info = mlflow.langchain.log_model(
lc_model="graph.py", # Path to our model Python file
artifact_path="langgraph",
)
model_uri = model_info.model_uri
4 - 使用已记录的模型
既然我们已经成功记录了一个模型，就可以加载它并利用它进行推理。
在下面的代码中，我们演示了我们的链具有聊天机器人功能！
import mlflow
# Custom utilities for handling chat history
from langgraph_utils import (
increment_message_history,
get_most_recent_message,
)
# Enable tracing
mlflow.set_experiment("Tracing example") # In Databricks, use an absolute path. Visit Databricks docs for more.
mlflow.langchain.autolog()
# Load the model
loaded_model = mlflow.langchain.load_model(model_uri)
# Show inference and message history functionality
print("-------- Message 1 -----------")
message = "What's my name?"
payload = {"messages": [{"role": "user", "content": message}]}
response = loaded_model.invoke(payload)
print(f"User: {message}")
print(f"Agent: {get_most_recent_message(response)}")
print("\n-------- Message 2 -----------")
message = "My name is Morpheus."
new_messages = increment_message_history(response, {"role": "user", "content": message})
payload = {"messages": new_messages}
response = loaded_model.invoke(payload)
print(f"User: {message}")
print(f"Agent: {get_most_recent_message(response)}")
print("\n-------- Message 3 -----------")
message = "What is my name?"
new_messages = increment_message_history(response, {"role": "user", "content": message})
payload = {"messages": new_messages}
response = loaded_model.invoke(payload)
print(f"User: {message}")
print(f"Agent: {get_most_recent_message(response)}")
输出：
-------- Message 1 -----------
User: What's my name?
Agent: I'm sorry, I cannot guess your name as I do not have access to that information. If you would like to share your name with me, feel free to do so.
-------- Message 2 -----------
User: My name is Morpheus.
Agent: Nice to meet you, Morpheus! How can I assist you today?
-------- Message 3 -----------
User: What is my name?
Agent: Your name is Morpheus.
4.1 - MLflow Tracing
在结束之前，让我们演示一下MLflow tracing。
MLflow Tracing 是一项功能，通过捕获应用程序服务执行的详细信息，增强您生成式人工智能（GenAI）应用中的LLM可观测性。追踪提供了一种记录请求每个中间步骤的输入、输出和元数据的方法，使您能够轻松定位错误和意外行为的来源。
按照跟踪服务器文档中所述启动MLflow服务器。进入MLflow UI后，我们可以看到实验和相应的追踪记录。
如您所见，我们已经记录了追踪信息，只需点击感兴趣的实验，然后选择"追踪"标签页即可轻松查看。
点击其中一个追踪后，我们现在可以看到单个查询的运行执行情况。请注意，我们会记录输入、输出以及大量优质元数据，例如使用情况和调用参数。随着我们的应用从使用量和复杂性两个维度扩展，这个线程安全且高性能的追踪系统将确保对应用的稳健监控。
5 - 总结
本教程有许多逻辑上的扩展，然而MLflow组件基本可以保持不变。一些例子包括将聊天记录持久化到数据库、实现更复杂的langgraph对象、将该解决方案投入生产等等！
总而言之，本教程涵盖的内容如下：
- 创建一个简单的LangGraph链。
- 利用 MLflow model from code 功能记录我们的图。
- 通过标准 MLflow API 加载模型。
- 利用 MLflow tracing 查看图执行。
编程愉快！