-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdemo.py
More file actions
166 lines (145 loc) · 6.23 KB
/
Copy pathdemo.py
File metadata and controls
166 lines (145 loc) · 6.23 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
from langgraph.graph import StateGraph, END
from typing import TypedDict, Optional
import requests
# ===================== 1. 定义状态（必须返回完整字典） =====================
class WeatherState(TypedDict):
user_input: str # 用户输入
weather_result: Optional[str] # 天气结果
route_decision: Optional[str] # 新增：存储大模型的路由判断结果
# ===================== 2. 大模型配置（适配llama3.2:1b） =====================
OLLAMA_BASE_URL = "http://localhost:11434/api/generate"
LLM_MODEL = "llama3.2:1b" # 更换为llama3.2:1b模型
def llm_classify_intent(user_input: str) -> str:
"""调用llama3.2:1b判断意图，仅返回 weather / other"""
# 适配llama3.2的提示词风格（更简洁，避免复杂指令）
prompt = f"""
Task: Classify user input as either 'weather' or 'other'.
Rules:
1. 'weather': Any input about weather (temperature, rain, sun, cold/hot, wind, etc.), e.g., "Is Shenzhen hot today?", "Will Shanghai rain?".
2. 'other': All non-weather inputs, e.g., "Hello", "What is Python?".
Only return 'weather' or 'other' (lowercase, no extra words).
User input: {user_input}
Answer:
"""
try:
response = requests.post(
OLLAMA_BASE_URL,
json={
"model": LLM_MODEL,
"prompt": prompt.strip(),
"stream": False,
"temperature": 0.0 # llama3.2用0温度更稳定
},
timeout=20 # 延长超时（1b模型加载稍慢）
)
response.raise_for_status()
result = response.json()["response"].strip().lower()
# 严格校验llama3.2的输出（避免模型返回多余内容）
if "weather" in result:
return "weather"
else:
return "other"
except Exception as e:
print(f"大模型分类出错：{e}")
return "other"
# ===================== 3. 天气工具（不变） =====================
class WeatherTool:
def get_weather(self, city: str) -> str:
try:
url = f"http://wttr.in/{city}?format=3"
response = requests.get(url, timeout=10)
response.raise_for_status()
print(response.text)
prompt = f"将结果整理一下，用中文回答, {response.text}"
response = requests.post(
OLLAMA_BASE_URL,
json={
"model": LLM_MODEL,
"prompt": prompt.strip(),
"stream": False,
"temperature": 0.0 # llama3.2用0温度更稳定
},
timeout=20 # 延长超时（1b模型加载稍慢）
)
response.raise_for_status()
return f"{response.text}（数据来源：wttr.in）"
except Exception as e:
return f"查询{city}天气失败：{str(e)[:50]}"
weather_tool = WeatherTool()
# ===================== 4. 节点定义（核心修复：返回完整state字典） =====================
def route_node(state: WeatherState) -> WeatherState:
"""路由节点：大模型判断意图，写入state，返回完整字典"""
# 1. 调用大模型判断意图
intent = llm_classify_intent(state["user_input"])
# 2. 将判断结果写入state（关键：不再直接返回字符串）
if intent == "weather":
state["route_decision"] = "weather_node"
else:
state["route_decision"] = "default_node"
# 3. 必须返回完整的state字典（核心修复点）
return state
def weather_node(state: WeatherState) -> WeatherState:
"""天气节点：返回完整state字典"""
user_input = state["user_input"]
# 提取城市名
cities = ["北京", "上海", "广州", "深圳", "杭州", "成都", "重庆", "南京", "武汉"]
city = ""
for c in cities:
if c in user_input:
city = c
break
if not city:
state["weather_result"] = "请告诉我具体城市，比如：北京天气、深圳今天热吗？"
else:
state["weather_result"] = weather_tool.get_weather(city)
return state # 返回完整字典
def default_node(state: WeatherState) -> WeatherState:
"""默认节点：返回完整state字典"""
state["weather_result"] = "我现在只支持天气查询哦～比如：北京天气、深圳今天热吗？"
return state # 返回完整字典
# ===================== 5. 构建图（修复条件边路由逻辑） =====================
def build_weather_graph():
graph = StateGraph(WeatherState)
# 添加节点
graph.add_node("route", route_node)
graph.add_node("weather_node", weather_node)
graph.add_node("default_node", default_node)
# 设置入口点
graph.set_entry_point("route")
# 核心修复：条件边从state中读取路由值（而非直接用节点返回值）
graph.add_conditional_edges(
"route", # 从路由节点出发
lambda state: state["route_decision"], # 从state中提取路由决策
{
"weather_node": "weather_node", # 路由值→目标节点
"default_node": "default_node"
}
)
# 连接结束节点
graph.add_edge("weather_node", END)
graph.add_edge("default_node", END)
return graph.compile()
# ===================== 6. 运行示例 =====================
if __name__ == "__main__":
# 检查Ollama和模型
try:
requests.get("http://localhost:11434", timeout=5)
except:
print("⚠️ 未检测到Ollama服务，请先启动：ollama serve")
print("⚠️ 并下载模型：ollama pull llama3.2:1b")
exit(1)
# 构建图
weather_agent = build_weather_graph()
print("=== LangGraph+llama3.2:1b 天气查询助手 ===")
print("输入示例：北京天气、深圳今天热吗？\n输入 exit 退出\n")
while True:
user_input = input("你：").strip()
if user_input.lower() == "exit":
break
# 调用图（传入完整state字典）
result = weather_agent.invoke({
"user_input": user_input,
"weather_result": None,
"route_decision": None # 新增的路由字段必须初始化
})
print(f"AI：{result['weather_result']}\n")