57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import operator
|
|
|
|
from typing import TypedDict, List, Annotated
|
|
from langgraph.graph import END, StateGraph
|
|
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
|
|
|
|
from src.multi_agent.config import LLM
|
|
from src.multi_agent.util.weather_search import search_weather
|
|
|
|
llm = LLM
|
|
|
|
|
|
class WeatherState(TypedDict):
|
|
user_request: str
|
|
city: str
|
|
result: str
|
|
|
|
|
|
def requirement_analysis(state: WeatherState):
|
|
user_request = state.get('user_request')
|
|
analysis_prompt = ChatPromptTemplate([
|
|
('system', '你是一个天气查询agent，你负责从用户的需求中提取出需要查询天气的城市或地区，你只需要输出这个地名，不需要输出其他任何内容'),
|
|
('user', f'当前的用户需求是:'),
|
|
MessagesPlaceholder(variable_name="user_request"),
|
|
])
|
|
analysis_chain = analysis_prompt | llm
|
|
llm_resp = analysis_chain.invoke({'user_request': [("user", user_request)]})
|
|
return {'city': llm_resp.content}
|
|
|
|
|
|
def get_weather_data(state: WeatherState):
|
|
city = state.get('city')
|
|
weather_data = search_weather(city)
|
|
return {'result': weather_data}
|
|
|
|
|
|
def enter_chain(input_data: dict):
|
|
return {'user_request': input_data['next_plan']}
|
|
|
|
|
|
workflow = StateGraph(WeatherState)
|
|
# 其实完全没必要使用子agent，但演示一下子agent如何并网
|
|
|
|
workflow.add_node('需求分析', requirement_analysis)
|
|
workflow.add_node('获取天气', get_weather_data)
|
|
|
|
workflow.set_entry_point('需求分析')
|
|
|
|
workflow.add_edge('需求分析', '获取天气')
|
|
workflow.add_edge('获取天气', END)
|
|
|
|
weather_graph = workflow.compile()
|
|
weather_chain = enter_chain | weather_graph
|
|
|
|
|
|
if __name__ == '__main__':
|
|
requirement_analysis({}) |