使用 LangGraph 迁移和优化 MultiPromptChain

老铁们，今天我们来聊聊如何从传统的 MultiPromptChain 迁移到更强大的 LangGraph 实现。如果你之前用过 MultiPromptChain，你可能已经知道它的基本工作流程：给定一个输入查询，它会使用一个 LLM 来从一系列提示中选择一个，格式化查询并生成响应。不过，这个老方法不支持一些现代聊天模型的功能，比如消息角色和工具调用。

说到这里，LangGraph 的优势就来了！它不仅支持聊天提示模板（包括带有系统和其他角色的消息），还可以在路由步骤中调用工具，并支持流式处理单个步骤和输出 token。那我们就来看看这俩家伙的对比吧。

快速准备

先确保环境准备好，这里需要安装最新版的 langchain-openai：

%pip install -qU langchain-core langchain-openai

import os
from getpass import getpass

os.environ["OPENAI_API_KEY"] = getpass()

传统的 MultiPromptChain

老铁们，我先前踩过这个坑，传统的 MultiPromptChain 是这样使用的：

from langchain.chains.router.multi_prompt import MultiPromptChain
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini")

prompt_1_template = """
You are an expert on animals. Please answer the below query:

{input}
"""

prompt_2_template = """
You are an expert on vegetables. Please answer the below query:

{input}
"""

prompt_infos = [
    {"name": "animals", "description": "prompt for an animal expert", "prompt_template": prompt_1_template,},
    {"name": "vegetables", "description": "prompt for a vegetable expert", "prompt_template": prompt_2_template,},
]

chain = MultiPromptChain.from_prompts(llm, prompt_infos)

调用链以获取结果：

chain.invoke({"input": "What color are carrots?"})

这波操作在 LangSmith 中你可以看到两个步骤：路由查询和选择的最终提示。

LangGraph 方式

接下来看看新方法：LangGraph。

pip install -qU langgraph

from operator import itemgetter
from typing import Literal

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict

llm = ChatOpenAI(model="gpt-4o-mini")

prompt_1 = ChatPromptTemplate.from_messages([("system", "You are an expert on animals."), ("human", "{input}")])
prompt_2 = ChatPromptTemplate.from_messages([("system", "You are an expert on vegetables."), ("human", "{input}")])

chain_1 = prompt_1 | llm | StrOutputParser()
chain_2 = prompt_2 | llm | StrOutputParser()

route_system = "Route the user's query to either the animal or vegetable expert."
route_prompt = ChatPromptTemplate.from_messages([("system", route_system), ("human", "{input}")])

class RouteQuery(TypedDict):
    destination: Literal["animal", "vegetable"]

route_chain = route_prompt | llm.with_structured_output(RouteQuery)

class State(TypedDict):
    query: str
    destination: RouteQuery
    answer: str

async def route_query(state: State, config: RunnableConfig):
    destination = await route_chain.ainvoke(state["query"], config)
    return {"destination": destination}

async def prompt_1(state: State, config: RunnableConfig):
    return {"answer": await chain_1.ainvoke(state["query"], config)}

async def prompt_2(state: State, config: RunnableConfig):
    return {"answer": await chain_2.ainvoke(state["query"], config)}

def select_node(state: State) -> Literal["prompt_1", "prompt_2"]:
    return "prompt_1" if state["destination"] == "animal" else "prompt_2"

graph = StateGraph(State)
graph.add_node("route_query", route_query)
graph.add_node("prompt_1", prompt_1)
graph.add_node("prompt_2", prompt_2)

graph.add_edge(START, "route_query")
graph.add_conditional_edges("route_query", select_node)
graph.add_edge("prompt_1", END)
graph.add_edge("prompt_2", END)
app = graph.compile()

调用新的链：

state = await app.ainvoke({"query": "what color are carrots"})
print(state["destination"])
print(state["answer"])

通过 LangSmith 的追踪可以看到工具调用是如何路由查询的，而选择的提示又是如何生成答案的。

总结

说白了，新方法就是通过工具调用和低级别的图实现大大增强了灵活性和扩展性。LangGraph 真正让链式调用的过程变得更丝滑。开发过程中遇到问题也可以在评论区交流~

今天的技术分享就到这里，希望对大家有帮助。开发过程中遇到问题也可以在评论区交流~

—END—