mirror of
https://github.com/FunnyWolf/agentic-soc-platform.git
synced 2026-08-22 13:12:56 +02:00
126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
import json
|
||
from typing import Dict, Any, Annotated, List
|
||
|
||
from langchain_core.messages import HumanMessage
|
||
from langgraph.graph import StateGraph, add_messages
|
||
from langgraph.graph.state import CompiledStateGraph
|
||
from pydantic import BaseModel
|
||
|
||
from Lib.baseplaybook import LanggraphPlaybook
|
||
from PLUGINS.LLM.llmapi import LLMAPI
|
||
from PLUGINS.SIRP.nocolyapi import WorksheetRow
|
||
from PLUGINS.SIRP.sirpapi import Alert, Artifact
|
||
from PLUGINS.SIRP.sirpapi import Notice
|
||
from PLUGINS.SIRP.sirpapi import Playbook as SIRPPlaybook
|
||
|
||
|
||
class AgentState(BaseModel):
|
||
messages: Annotated[List[Any], add_messages]
|
||
|
||
alert: Dict
|
||
suggestion: str
|
||
|
||
|
||
class Playbook(LanggraphPlaybook):
|
||
TYPE = "ALERT" # 分类标签
|
||
NAME = "Alert Analysis Agent" # 剧本名称
|
||
|
||
def __init__(self):
|
||
super().__init__() # do not delete this code
|
||
self.init()
|
||
|
||
def init(self):
|
||
def preprocess_node(state: AgentState):
|
||
"""预处理数据"""
|
||
# worksheet = self.param("worksheet")
|
||
rowid = self.param("rowid")
|
||
worksheet = self.param("worksheet")
|
||
alert = WorksheetRow.get(worksheet, rowid, include_system_fields=False)
|
||
artifacts = WorksheetRow.relations(Alert.WORKSHEET_ID, alert.get("rowId"), "artifact", relation_worksheet_id=Artifact.WORKSHEET_ID,
|
||
include_system_fields=False)
|
||
alert["artifact"] = artifacts
|
||
state.alert = alert
|
||
return state
|
||
|
||
# 定义node
|
||
def analyze_node(state: AgentState):
|
||
"""AI分析告警数据"""
|
||
|
||
# 加载system prompt
|
||
system_prompt_template = self.load_system_prompt_template("L3_SOC_Analyst")
|
||
|
||
system_message = system_prompt_template.format()
|
||
|
||
# 构建few-shot示例
|
||
few_shot_examples = [
|
||
# HumanMessage(
|
||
# content=json.dumps({
|
||
# "requirement": ".",
|
||
# })
|
||
# ),
|
||
# AIMessage(
|
||
# content=json.dumps({
|
||
# "function": "the amount of pneumothorax",
|
||
# })
|
||
# ),
|
||
]
|
||
|
||
# 运行
|
||
llm_api = LLMAPI()
|
||
|
||
llm = llm_api.get_model(tag="fast")
|
||
|
||
# 构建消息列表
|
||
messages = [
|
||
system_message,
|
||
*few_shot_examples,
|
||
HumanMessage(content=json.dumps(state.alert))
|
||
]
|
||
response = llm.invoke(messages)
|
||
response = LLMAPI.extract_think(response) # langchain chatollama bug临时方案
|
||
state.suggestion = response.content
|
||
return state
|
||
|
||
def output_node(state: AgentState):
|
||
"""处理分析结果"""
|
||
|
||
suggestion = state.suggestion
|
||
fields = [
|
||
{"id": "suggestion_ai", "value": suggestion},
|
||
]
|
||
rowid = self.param("rowid")
|
||
WorksheetRow.update(Alert.WORKSHEET_ID, rowid, fields)
|
||
|
||
self.agent_state = state
|
||
|
||
Notice.send(self.param("user"), "Alert_Suggestion_Gen_By_LLM output_node Finish", f"rowid：{self.param('rowid')}")
|
||
|
||
SIRPPlaybook.update_status_and_remark(self.param("playbook_rowid"), "Success", "Get suggestion by ai agent completed.") # Success/Failed
|
||
return state
|
||
|
||
# 编译graph
|
||
workflow = StateGraph(AgentState)
|
||
|
||
workflow.add_node("preprocess_node", preprocess_node)
|
||
workflow.add_node("analyze_node", analyze_node)
|
||
workflow.add_node("output_node", output_node)
|
||
|
||
workflow.set_entry_point("preprocess_node")
|
||
workflow.add_edge("preprocess_node", "analyze_node")
|
||
workflow.add_edge("analyze_node", "output_node")
|
||
workflow.set_finish_point("output_node")
|
||
self.agent_state = AgentState(messages=[], alert={}, suggestion="")
|
||
self.graph: CompiledStateGraph = workflow.compile(checkpointer=self.get_checkpointer())
|
||
return True
|
||
|
||
def run(self):
|
||
self.run_graph()
|
||
return
|
||
|
||
|
||
if __name__ == "__main__":
|
||
params_debug = {'rowid': '13782e0a-2423-4fc3-9b16-7f2eb15ae83f', 'worksheet': 'alert'}
|
||
module = Playbook()
|
||
module._params = params_debug
|
||
module.run()
|