from langchain.chains.router.multi_promptimport MultiPromptChain
from langchain_openaiimport 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?"})

{'input': 'What color are carrots?',
 'text': 'Carrots are most commonly orange, but they can also be found in a variety of other colors including purple, yellow, white, and red. The orange variety is the most popular and widely recognized.'}

pip install-qU langgraph

from operatorimport itemgetter
from typingimport Literal

from langchain_core.output_parsersimport StrOutputParser
from langchain_core.promptsimport ChatPromptTemplate
from langchain_core.runnablesimport RunnableConfig
from langchain_openaiimport ChatOpenAI
from langgraph.graphimport END, START, StateGraph
from typing_extensionsimport TypedDict

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

# Define the prompts we will route to
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}"),
]
)

# Construct the chains we will route to. These format the input query
# into the respective prompt, run it through a chat model, and cast
# the result to a string.
chain_1= prompt_1| llm| StrOutputParser()
chain_2= prompt_2| llm| StrOutputParser()


# Next: define the chain that selects which branch to route to.
# Here we will take advantage of tool-calling features to force
# the output to select one of two desired branches.
route_system="Route the user's query to either the animal or vegetable expert."
route_prompt= ChatPromptTemplate.from_messages(
[
("system", route_system),
("human","{input}"),
]
)


# Define schema for output:
classRouteQuery(TypedDict):
"""Route query to destination expert."""

    destination: Literal["animal","vegetable"]


route_chain= route_prompt| llm.with_structured_output(RouteQuery)


# For LangGraph, we will define the state of the graph to hold the query,
# destination, and final answer.
classState(TypedDict):
    query:str
    destination: RouteQuery
    answer:str


# We define functions for each node, including routing the query:
asyncdefroute_query(state: State, config: RunnableConfig):
    destination=await route_chain.ainvoke(state["query"], config)
return{"destination": destination}


# And one node for each prompt
asyncdefprompt_1(state: State, config: RunnableConfig):
return{"answer":await chain_1.ainvoke(state["query"], config)}


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


# We then define logic that selects the prompt based on the classification
defselect_node(state: State)-> Literal["prompt_1","prompt_2"]:
if state["destination"]=="animal":
return"prompt_1"
else:
return"prompt_2"


# Finally, assemble the multi-prompt chain. This is a sequence of two steps:
# 1) Select "animal" or "vegetable" via the route_chain, and collect the answer
# alongside the input query.
# 2) Route the input query to chain_1 or chain_2, based on the
# selection.
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()

from IPython.displayimport Image

Image(app.get_graph().draw_mermaid_png())

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

{'destination': 'vegetable'}
Carrots are most commonly orange, but they can also come in a variety of other colors, including purple, red, yellow, and white. The different colors often indicate varying flavors and nutritional profiles. For example, purple carrots contain anthocyanins, while orange carrots are rich in beta-carotene, which is converted to vitamin A in the body.