使用LangGraph迁移MultiPromptChain以提升多重提示链效率_await planner.ainvoke({"messages": [("user", state-优快云博客

本文链接：https://blog.youkuaiyun.com/tt_jishu/article/details/148713462

在不断发展的AI技术中，多提示链(MultiPromptChain)通过将输入查询路由到多个LLM链(LLMChain)来处理复杂的问题。然而，MultiPromptChain存在一些局限性，如不支持消息角色和工具调用等通用聊天模型功能。为了解决这些问题，LangGraph提供了一系列优势，包括支持系统角色消息和工具调用，以及流式处理输出步骤。

本篇文章将详细讲解如何使用LangGraph来替代MultiPromptChain，并提供相关代码示例。

技术背景介绍

MultiPromptChain主要用于将输入查询路由到不同的LLM链，从而选择合适的提示模板并生成响应。然而，随着AI交互需求的提升，MultiPromptChain在支持多角色聊天和工具调用方面存在不足。LangGraph通过底层原语实现了一种更灵活的解决方案，为开发者提供了更多功能支持。

核心原理解析

LangGraph允许用户定义复杂的状态图(StateGraph)，并在该图中使用工具调用来选择合适的提示链。通过支持多角色的聊天提示模板和结构化输出解析，LangGraph实现了比MultiPromptChain更为复杂和丰富的功能。

代码实现演示(重点)

安装必要的库

首先，确保安装所需的库：

%pip install -qU langchain-core langchain-openai langgraph

代码实现细节

这里我们将展示如何使用LangGraph来实现多提示链路由。

import os
from getpass import getpass
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph

os.environ["OPENAI_API_KEY"] = getpass()  # 输入你的API密钥

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"])