226 lines
7.6 KiB
Python
226 lines
7.6 KiB
Python
import os
|
||
from config import OPENAI_API_KEY,LANGCHAIN_API_KEY
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain.tools import tool
|
||
from langchain_core.messages import HumanMessage,SystemMessage
|
||
from langgraph.graph import StateGraph,START,END
|
||
from langgraph.prebuilt import ToolNode
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
from typing import TypedDict,Annotated
|
||
from langgraph.graph.message import add_messages
|
||
|
||
# LangSmith调试
|
||
os.environ["LANGCHAIN_TRACING_V2"] = "true" # 总开关，决定启用追踪功能
|
||
os.environ["LANGCHAIN_PROJECT"] = "supervisor_agent_ops_system" # 自定义项目名
|
||
os.environ["LANGCHAIN_API_KEY"] = LANGCHAIN_API_KEY
|
||
|
||
llm = ChatOpenAI(
|
||
model="deepseek-chat",
|
||
api_key=OPENAI_API_KEY,
|
||
base_url="https://api.deepseek.com"
|
||
)
|
||
|
||
# === 一、Graph-as-a-Tool ===
|
||
# === 模拟一个不稳定的SSH日志查询过程 ===
|
||
|
||
class SSHState(TypedDict):
|
||
target_ip: str
|
||
attempt: int
|
||
logs: str
|
||
|
||
def connect_ssh(state:SSHState):
|
||
"""模拟SSH连接，第一次连接必定超时"""
|
||
print(f' [子图]正在尝试连接服务器{state["target_ip"]}(第{state["attempt"]}次)')
|
||
if state['attempt'] == 1:
|
||
return {'logs':'ERROR:Connection Timed Out','attempt':state['attempt']+1}
|
||
return {'logs':'CONNECTED','attempt':state['attempt']+1}
|
||
|
||
def grep_system_logs(state:SSHState):
|
||
"""连接成功后读取日志"""
|
||
if state["logs"] == "CONNECTED":
|
||
# 打印查到的结果
|
||
return {'logs':f'SUCCESS: Retrieved logs from {state["target_ip"]}:[ERROR: OutOfMemory at line 4032]'}
|
||
return {'logs':state['logs']} # 保持错误状态
|
||
|
||
def ssh_routing(state:SSHState):
|
||
"""路由逻辑：如果连接失败且尝试次数少于3，重试"""
|
||
if "ERROR" in state['logs'] and state["attempt"] <= 2:
|
||
return "connect"
|
||
return "grep"
|
||
|
||
# 构建子图
|
||
ssh_workflow = StateGraph(SSHState)
|
||
ssh_workflow.add_node("connect",connect_ssh)
|
||
ssh_workflow.add_node("grep",grep_system_logs)
|
||
|
||
ssh_workflow.add_edge(START,"connect")
|
||
ssh_workflow.add_conditional_edges("connect",ssh_routing,{"connect":"connect","grep":"grep"})
|
||
ssh_workflow.add_edge("grep",END)
|
||
|
||
ssh_app = ssh_workflow.compile()
|
||
|
||
# 将子图封装为工具
|
||
@tool
|
||
def analyze_server_logs(ip_address:str):
|
||
"""使用SSH连接服务器并分析最近的错误日志（内含自动重连机制）"""
|
||
result = ssh_app.invoke({"target_ip":ip_address,"attempt":1,"log":""})
|
||
return result['logs']
|
||
|
||
|
||
|
||
# === 二、Human-in-the-Loop ===
|
||
# === 重启服务，高危操作，需要审批 ===
|
||
|
||
@tool
|
||
def restart_service(service_name:str):
|
||
"""重启指定的服务器服务"""
|
||
return f"服务[{service_name}]已成功重启，系统负载已恢复正常"
|
||
|
||
|
||
|
||
# === 三、Multi-Agent 编排 ===
|
||
# === 总控调度 + 专家分工 ===
|
||
|
||
# 1. 共享状态
|
||
class AgentState(TypedDict):
|
||
messages:Annotated[list,add_messages]
|
||
next_speaker:str
|
||
|
||
# 2. 专家节点
|
||
def log_expert(state:AgentState):
|
||
prompt = "你是日志分析专家，使用工具分析服务器日志，找出报错原因。回答需简洁。"
|
||
messages = [SystemMessage(content=prompt)] + state['messages']
|
||
# 绑定子图工具
|
||
tools = [analyze_server_logs]
|
||
response = llm.bind_tools(tools).invoke(messages)
|
||
return {"messages":[response]}
|
||
|
||
def ops_expert(state:AgentState):
|
||
prompt = "你是运维专家。当收到修复指令时，请立即调用 'restart_service' 工具进行修复，不要输出任何额外的解释文本。"
|
||
messages = [SystemMessage(content=prompt)] + state['messages']
|
||
tools = [restart_service]
|
||
# 绑定敏感工具
|
||
response = llm.bind_tools(tools).invoke(messages)
|
||
return {"messages":[response]}
|
||
|
||
# 3. 总控节点（supervisor）
|
||
def supervisor(state:AgentState):
|
||
prompt = """
|
||
你是 IT 运维总指挥。
|
||
专家列表：
|
||
- log_expert
|
||
- ops_expert
|
||
|
||
决策逻辑：
|
||
1. 未知原因 -> log_expert
|
||
2. 已知原因（如OOM、报错） -> ops_expert
|
||
3. 修复完成 -> FINISH
|
||
|
||
【输出约束】
|
||
仅输出下一个专家的名字（如 log_expert），不要包含任何其他字符或标点。
|
||
"""
|
||
messages = [SystemMessage(content=prompt)] + state['messages']
|
||
response = llm.invoke(messages)
|
||
return {"next_speaker":response.content.strip()}
|
||
|
||
# 4. 路由逻辑
|
||
def route_supervisor(state:AgentState):
|
||
if state['next_speaker'] == "FINISH":
|
||
return END
|
||
return state['next_speaker']
|
||
|
||
def should_continue(state:AgentState):
|
||
last_msg = state['messages'][-1]
|
||
if hasattr(last_msg,"tool_calls") and last_msg.tool_calls:
|
||
return "tools"
|
||
return "supervisor"
|
||
|
||
def route_after_tool(state:AgentState):
|
||
return state["next_speaker"]
|
||
|
||
|
||
# === 四、构建主图与集成 ===
|
||
|
||
# 工具集合
|
||
all_tools = [analyze_server_logs,restart_service]
|
||
tool_node = ToolNode(all_tools)
|
||
|
||
# 构建主图
|
||
workflow = StateGraph(AgentState)
|
||
|
||
workflow.add_node("supervisor",supervisor)
|
||
workflow.add_node("log_expert",log_expert)
|
||
workflow.add_node("ops_expert",ops_expert)
|
||
workflow.add_node("tools",tool_node)
|
||
|
||
workflow.add_edge(START,"supervisor")
|
||
workflow.add_conditional_edges("supervisor",route_supervisor)
|
||
for member in ["log_expert","ops_expert"]:
|
||
workflow.add_conditional_edges(
|
||
member,
|
||
should_continue,
|
||
{"tools":"tools","supervisor":"supervisor"}
|
||
)
|
||
workflow.add_conditional_edges("tools",route_after_tool)
|
||
|
||
# 编译图:加入记忆与中断机制(Human-in-the-Loop)
|
||
# 注: 我们在所有工具执行前都暂停，但在运行时进行逻辑判断
|
||
app = workflow.compile(
|
||
checkpointer=MemorySaver(),
|
||
interrupt_before=["tools"]
|
||
)
|
||
|
||
# === 五、运行时逻辑(模拟生产环境的交互) ===
|
||
|
||
if __name__ == '__main__':
|
||
# 模拟一次完整的故障处理流程
|
||
user_input = "服务器 192.168.1.100 报警，响应极慢，请处理。"
|
||
config = {
|
||
"configurable": {"thread_id": "incident_001"}
|
||
}
|
||
|
||
print(f'收到报警 : {user_input}')
|
||
inputs = {"messages":[HumanMessage(content=user_input)]}
|
||
|
||
# 循环执行，直到任务结束
|
||
while 1:
|
||
# 1. 执行图直到中断或结束
|
||
for _ in app.stream(inputs,config,stream_mode="values"):
|
||
pass
|
||
|
||
# 2. 检查当前状态
|
||
snapshot = app.get_state(config)
|
||
next_tasks = snapshot.next
|
||
|
||
# 没有下一步，任务结束
|
||
if not next_tasks:
|
||
print(f" 最终报告:{snapshot.values['messages'][-1].content}")
|
||
break
|
||
|
||
# 3. 处理中断:判断是哪个工具被调用
|
||
if "tools" in next_tasks:
|
||
last_msg = snapshot.values['messages'][-1]
|
||
tool_call = last_msg.tool_calls[0]
|
||
tool_name = tool_call["name"]
|
||
|
||
print(f'\n[系统暂停] 请求调用工具:{tool_name}')
|
||
|
||
# 策略A: 自动放行安全工具(Graph-as-a-Tool)
|
||
if tool_name == "analyze_server_logs":
|
||
print(' -> 这是一个查询类工具，系统自行批准。')
|
||
inputs = None # 继续执行
|
||
continue
|
||
|
||
# 策略B: 拦截高危工具(Human-in-the-Loop)
|
||
elif tool_name == "restart_service":
|
||
print(" -> ⚠️ 警告: 这是一个高危操作!")
|
||
user_approval = input(" -> 请人工审批 (输入 'yes' 允许重启):")
|
||
|
||
if user_approval == "yes":
|
||
print(' -> ✅ 审批通过，正在执行...')
|
||
inputs = None # 继续执行
|
||
else:
|
||
print(' -> ❌️ 审批拒绝，任务终止!')
|
||
# 实际系统中，应通过 ToolMessage 反馈人工拒绝，使 LLM 能继续响应；
|
||
# 当前demo为简化，直接退出
|
||
break |