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
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
| """RAG Agent 服务 - 基于 LangGraph 的智能代理
使用 langchain_qwq 的 ChatQwen 原生集成，
支持真正的流式输出和更好的模型适配。
"""
from typing import Annotated, Any, AsyncGenerator, Dict, Sequence
from langchain.agents import create_agent
from langchain_core.messages import (
BaseMessage,
HumanMessage,
RemoveMessage,
SystemMessage,
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
from loguru import logger
from typing_extensions import TypedDict
from langchain_qwq import ChatQwen
from app.config import config
from app.tools import DEFAULT_LOCAL_AGENT_TOOLS
from app.agent.mcp_client import (
get_mcp_client_with_retry,
load_mcp_tools_safe,
format_exception_chain,
suggest_mcp_transport,
)
# 阿里千问大模型和langchain集成参考： https://docs.langchain.com/oss/python/integrations/chat/qwen
# 注意：需要配置环境变量 DASHSCOPE_API_BASE=https://dashscope.aliyuncs.com/compatible-mode/v1 否则默认访问的是新加坡站点
# 同时也需要配置环境变量 DASHSCOPE_API_KEY=your_api_key
class AgentState(TypedDict):
"""Agent 状态"""
messages: Annotated[Sequence[BaseMessage], add_messages]
def trim_messages_middleware(state: AgentState) -> dict[str, Any] | None:
"""
修剪消息历史，只保留最近的几条消息以适应上下文窗口
策略：
- 保留第一条系统消息（System Message）
- 保留最近的 6 条消息（3 轮对话）
- 当消息少于等于 7 条时，不做修剪
Args:
state: Agent 状态
Returns:
包含修剪后消息的字典，如果无需修剪则返回 None
"""
messages = state["messages"]
# 如果消息数量较少，无需修剪
if len(messages) <= 7:
return None
# 提取第一条系统消息
first_msg = messages[0]
# 保留最近的 6 条消息（确保包含完整的对话轮次）
recent_messages = messages[-6:] if len(messages) % 2 == 0 else messages[-7:]
# 构建新的消息列表
new_messages = [first_msg] + list(recent_messages)
logger.debug(f"修剪消息历史: {len(messages)} -> {len(new_messages)} 条")
return {
"messages": [
RemoveMessage(id=REMOVE_ALL_MESSAGES),
*new_messages
]
}
class RagAgentService:
"""RAG Agent 服务 - 使用 LangGraph + ChatQwen 原生集成"""
def __init__(self, streaming: bool = True):
"""初始化 RAG Agent 服务
Args:
streaming: 是否启用流式输出，默认为 True
"""
self.model_name = config.rag_model
self.streaming = streaming
self.system_prompt = self._build_system_prompt()
self.model = ChatQwen(
model=self.model_name,
api_key=config.dashscope_api_key,
temperature=0.7,
streaming=streaming,
)
# 定义基础工具（与 AIOps Planner/Executor 使用同一套默认本地工具）
self.tools = list(DEFAULT_LOCAL_AGENT_TOOLS)
# MCP 客户端（延迟初始化，使用全局管理）
self.mcp_tools: list = []
# 创建内存检查点（用于会话管理）
self.checkpointer = MemorySaver()
# Agent 初始化（会在异步方法中完成）
self.agent = None
self._agent_initialized = False
logger.info(f"RAG Agent 服务初始化完成 (ChatQwen), model={self.model_name}, streaming={streaming}")
async def _initialize_agent(self):
"""异步初始化 Agent（包括 MCP 工具）"""
if self._agent_initialized:
return
for name, server in config.mcp_servers.items():
hint = suggest_mcp_transport(
str(server.get("url", "")),
str(server.get("transport", "")),
)
if hint:
logger.warning(f"MCP 配置 [{name}]: {hint}")
mcp_client = await get_mcp_client_with_retry()
mcp_tools, mcp_err = await load_mcp_tools_safe(mcp_client)
if mcp_err:
logger.warning(
f"MCP 工具加载失败，将仅使用本地工具继续运行:\n{mcp_err}"
)
self.mcp_tools = []
else:
self.mcp_tools = mcp_tools
logger.info(f"成功加载 {len(mcp_tools)} 个 MCP 工具")
all_tools = self.tools + self.mcp_tools
self.agent = create_agent(
self.model,
tools=all_tools,
checkpointer=self.checkpointer,
)
self._agent_initialized = True
if all_tools:
tool_names = [tool.name if hasattr(tool, "name") else str(tool) for tool in all_tools]
logger.info(f"可用工具列表: {', '.join(tool_names)}")
def _build_system_prompt(self) -> str:
"""
构建系统提示词
注意：LangChain 框架会自动将工具信息传递给 LLM，
因此系统提示词中无需列举具体的工具列表。
Returns:
str: 系统提示词
"""
from textwrap import dedent
return dedent("""
你是一个专业的AI助手，能够使用多种工具来帮助用户解决问题。
工作原则:
1. 理解用户需求，选择合适的工具来完成任务
2. 当需要获取实时信息或专业知识时，主动使用相关工具
3. 基于工具返回的结果提供准确、专业的回答
4. 如果工具无法提供足够信息，请诚实地告知用户
回答要求:
- 保持友好、专业的语气
- 回答简洁明了，重点突出
- 基于事实，不编造信息
- 如有不确定的地方，明确说明
请根据用户的问题，灵活使用可用工具，提供高质量的帮助。
""").strip()
async def query(
self,
question: str,
session_id: str,
) -> str:
"""
非流式处理用户问题（一次性返回完整答案）
Args:
question: 用户问题
session_id: 会话ID（作为 thread_id）
Returns:
str: 完整答案
"""
try:
await self._initialize_agent()
logger.info(f"[会话 {session_id}] RAG Agent 收到查询（非流式）: {question}")
# 构建消息列表（系统提示 + 用户问题）
messages = [
SystemMessage(content=self.system_prompt),
HumanMessage(content=question)
]
# 构建 Agent 输入
agent_input = {"messages": messages}
# 配置 thread_id（用于会话持久化）
config_dict = {
"configurable": {
"thread_id": session_id
}
}
result = await self.agent.ainvoke(
input=agent_input,
config=config_dict,
)
# 提取最终答案
messages_result = result.get("messages", [])
if messages_result:
last_message = messages_result[-1]
answer = last_message.content if hasattr(last_message, 'content') else str(last_message)
# 记录工具调用
if hasattr(last_message, "tool_calls") and last_message.tool_calls:
tool_names = [tc.get("name", "unknown") for tc in last_message.tool_calls]
logger.info(f"[会话 {session_id}] Agent 调用了工具: {tool_names}")
logger.info(f"[会话 {session_id}] RAG Agent 查询完成（非流式）")
return answer
logger.warning(f"[会话 {session_id}] Agent 返回结果为空")
return ""
except Exception as e:
logger.error(
f"[会话 {session_id}] RAG Agent 查询失败（非流式）: "
f"{format_exception_chain(e)}"
)
raise
async def query_stream(
self,
question: str,
session_id: str,
) -> AsyncGenerator[Dict[str, Any], None]:
"""
流式处理用户问题（逐步返回答案片段）
Args:
question: 用户问题
session_id: 会话ID（作为 thread_id）
Yields:
Dict[str, Any]: 包含流式数据的字典
- type: "content" | "tool_call" | "complete" | "error"
- data: 具体内容
"""
try:
await self._initialize_agent()
logger.info(f"[会话 {session_id}] RAG Agent 收到查询（流式）: {question}")
# 构建消息列表（系统提示 + 用户问题）
messages = [
SystemMessage(content=self.system_prompt),
HumanMessage(content=question)
]
# 构建 Agent 输入
agent_input = {"messages": messages}
# 配置 thread_id（用于会话持久化）
config_dict = {
"configurable": {
"thread_id": session_id
}
}
async for token, metadata in self.agent.astream(
input=agent_input,
config=config_dict,
stream_mode="messages",
):
node_name = metadata.get('langgraph_node', 'unknown') if isinstance(metadata, dict) else 'unknown'
message_type = type(token).__name__
if message_type in ("AIMessage", "AIMessageChunk"):
content_blocks = getattr(token, 'content_blocks', None)
if content_blocks and isinstance(content_blocks, list):
for block in content_blocks:
if isinstance(block, dict) and block.get('type') == 'text':
text_content = block.get('text', '')
if text_content:
yield {
"type": "content",
"data": text_content,
"node": node_name
}
logger.info(f"[会话 {session_id}] RAG Agent 查询完成（流式）")
yield {"type": "complete"}
except Exception as e:
detail = format_exception_chain(e)
logger.error(
f"[会话 {session_id}] RAG Agent 查询失败（流式）: {detail}"
)
yield {"type": "error", "data": detail}
def get_session_history(self, session_id: str) -> list:
"""
获取会话历史（从 MemorySaver checkpointer 中读取）
Args:
session_id: 会话ID（即 thread_id）
Returns:
list: 消息历史列表 [{"role": "user|assistant", "content": "...", "timestamp": "..."}]
"""
try:
# 使用 checkpointer 的 get 方法获取最新的检查点
config = {"configurable": {"thread_id": session_id}}
# 获取该 thread 的最新检查点
checkpoint_tuple = self.checkpointer.get(config)
if not checkpoint_tuple:
logger.info(f"获取会话历史: {session_id}, 消息数量: 0")
return []
# checkpoint_tuple 可能是命名元组或普通元组，安全地提取 checkpoint
# 通常第一个元素是 checkpoint 数据
if hasattr(checkpoint_tuple, 'checkpoint'):
checkpoint_data = checkpoint_tuple.checkpoint # type: ignore
else:
# 如果是普通元组，第一个元素是 checkpoint
checkpoint_data = checkpoint_tuple[0] if checkpoint_tuple else {}
# 从检查点中提取消息
messages = checkpoint_data.get("channel_values", {}).get("messages", [])
# 转换为前端需要的格式
history = []
for msg in messages:
# 跳过系统消息
if isinstance(msg, SystemMessage):
continue
role = "user" if isinstance(msg, HumanMessage) else "assistant"
content = msg.content if hasattr(msg, 'content') else str(msg)
# 提取时间戳（如果有的话）
timestamp = getattr(msg, 'timestamp', None)
if timestamp:
history.append({
"role": role,
"content": content,
"timestamp": timestamp
})
else:
from datetime import datetime
history.append({
"role": role,
"content": content,
"timestamp": datetime.now().isoformat()
})
logger.info(f"获取会话历史: {session_id}, 消息数量: {len(history)}")
return history
except Exception as e:
logger.error(f"获取会话历史失败: {session_id}, 错误: {e}")
return []
def clear_session(self, session_id: str) -> bool:
"""
清空会话历史（从 MemorySaver checkpointer 中删除）
Args:
session_id: 会话ID（即 thread_id）
Returns:
bool: 是否成功
"""
try:
# 使用 checkpointer 的 delete_thread 方法删除该 thread 的所有检查点
self.checkpointer.delete_thread(session_id)
logger.info(f"已清除会话历史: {session_id}")
return True
except Exception as e:
logger.error(f"清空会话历史失败: {session_id}, 错误: {e}")
return False
async def cleanup(self):
"""清理资源"""
try:
logger.info("清理 RAG Agent 服务资源...")
# MCP 客户端由全局管理器统一管理，无需手动清理
logger.info("RAG Agent 服务资源已清理")
except Exception as e:
logger.error(f"清理资源失败: {e}")
# 全局单例 - 启用流式输出
rag_agent_service = RagAgentService(streaming=True)
|
