LangGraph实战案例之MessagesState类继承和能力扩展

import os
from copyreg import constructor

from langchain_core.prompt_values import StringPromptValue
from langgraph.graph import StateGraph, MessagesState
from langgraph.graph.state import CompiledStateGraph, StateGraph  # 重复了
from langgraph.constants import START,END
from langchain_core.messages import AIMessage
from app.bailian.common import llm

key_keywords="key_keywords"
key_search="key_search"
key_answer="key_answer"

class BaiduSearchSate(MessagesState):
    search_question:str
    search_keywords:str
    search_results:str

def node_key_keywords(state:BaiduSearchSate):
    last_message = state['messages'][-1]
    content = last_message.content
    state['search_question'] = content
    prompt = StringPromptValue(text=f"请从下面的信息中提取需要在百度中搜索的关键词，并且直接返回结果:{content}")
    message = llm.invoke(input=prompt)
    state['search_keywords'] = message
    return state

def node_key_search(state:BaiduSearchSate):
    state['search_results'] ="百度查询的天气结果是: 今天0度 有点冷"
    return state

def node_key_answer(state:BaiduSearchSate):
    content = state['search_results']
    question =   state['search_question']
    message = llm.invoke(input=f"""
    # 要求
    请结合百度提供的结果 回答用户的问题   {question}
    
    {content}
    """)

    state['messages'].append(message)
    return state

def output_graph_image(graph,file_name):
    png_data = graph.get_graph().draw_png()
    output_file_dir = os.path.dirname(__file__)
    output_file_path = os.path.join(output_file_dir,file_name+'.png')
    with open(output_file_path,'wb') as f:
        f.write(png_data)


state_graph = StateGraph(MessagesState)
state_graph.add_node(key_keywords,node_key_keywords)
state_graph.add_node(key_search,node_key_search)
state_graph.add_node(key_answer,node_key_answer)

state_graph.add_edge(START,key_keywords)
state_graph.add_edge(key_keywords,key_search)
state_graph.add_edge(key_search,key_answer)
state_graph.add_edge(key_answer,END)

compiled_graph = state_graph.compile()

#output_graph_image(compiled_graph,'output_graph_image')

results = compiled_graph.stream({
    "messages":[("user","请问北京今天的天气如何")]
})

for result in results:
    print(result)