mirror of
https://github.com/FunnyWolf/agentic-soc-platform.git
synced 2026-08-22 13:12:56 +02:00
121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
import json
|
|
|
|
from langchain_core.messages import (
|
|
BaseMessage,
|
|
SystemMessage,
|
|
HumanMessage,
|
|
AIMessage,
|
|
ToolMessage
|
|
)
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langgraph.graph.state import CompiledStateGraph
|
|
from pydantic import BaseModel
|
|
|
|
from Lib.baseapi import BaseAPI
|
|
from Lib.llmapi import BaseAgentState
|
|
from Lib.log import logger
|
|
from PLUGINS.SIRP.sirpapi import Message
|
|
from PLUGINS.SIRP.sirpapi import Playbook, Notice
|
|
from PLUGINS.SIRP.sirpmodel import PlaybookModel, PlaybookJobStatus, MessageModel, MessageType
|
|
|
|
|
|
class BasePlaybook(BaseAPI):
|
|
RUN_AS_JOB = True # 是否作为后台任务运行
|
|
NAME = None
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.logger = logger
|
|
# noinspection PyTypeChecker
|
|
self._playbook_model: PlaybookModel = None
|
|
|
|
@property
|
|
def param_source_rowid(self):
|
|
return self._playbook_model.source_rowid
|
|
|
|
@property
|
|
def param_user_input(self):
|
|
return self._playbook_model.user_input
|
|
|
|
def update_playbook_status(self, job_status: PlaybookJobStatus, remark: str):
|
|
model_tmp = PlaybookModel()
|
|
model_tmp.rowid = self._playbook_model.rowid
|
|
model_tmp.job_status = job_status
|
|
model_tmp.remark = remark
|
|
rowid = Playbook.update(model_tmp)
|
|
return rowid
|
|
|
|
def send_notice(self, title: str, body: str) -> bool:
|
|
result = Notice.send(self._playbook_model.user, title, body)
|
|
return result
|
|
|
|
def execute(self):
|
|
try:
|
|
self.run()
|
|
except Exception as e:
|
|
self.logger.exception(e)
|
|
self.update_playbook_status(PlaybookJobStatus.FAILED, str(e))
|
|
|
|
|
|
class LanggraphPlaybook(BasePlaybook):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.graph: CompiledStateGraph = None
|
|
self.agent_state = None
|
|
|
|
@staticmethod
|
|
def get_checkpointer():
|
|
checkpointer = MemorySaver()
|
|
return checkpointer
|
|
|
|
def add_message_to_playbook(self, message: BaseMessage | BaseModel, node=None):
|
|
|
|
message_model = MessageModel()
|
|
message_model.playbook = [self._playbook_model.rowid]
|
|
message_model.node = node
|
|
|
|
# handle content
|
|
if isinstance(message, BaseMessage):
|
|
message_model.content = message.content
|
|
|
|
if isinstance(message, SystemMessage):
|
|
message_model.type = MessageType.SYSTEM
|
|
elif isinstance(message, HumanMessage):
|
|
message_model.type = MessageType.HUMAN
|
|
elif isinstance(message, AIMessage):
|
|
if hasattr(message, 'tool_calls') and message.tool_calls:
|
|
message_model.type = MessageType.AI
|
|
message_model.data = json.dumps(message.tool_calls)
|
|
else:
|
|
message_model.type = MessageType.AI
|
|
elif isinstance(message, ToolMessage):
|
|
try:
|
|
json_data = {"name": message.name, "tool_call_id": message.tool_call_id, "result": json.loads(message.content)}
|
|
except json.decoder.JSONDecodeError:
|
|
json_data = {"name": message.name, "tool_call_id": message.tool_call_id, "result": message.content}
|
|
message_model.type = MessageType.TOOL
|
|
message_model.data = json.dumps(json_data)
|
|
elif isinstance(message, BaseModel):
|
|
message_model.type = MessageType.AI
|
|
message_model.data = message.model_dump_json(exclude_none=True, exclude_defaults=True, exclude_unset=True)
|
|
else:
|
|
logger.warning(f"Unknown message type: {message.type}.")
|
|
|
|
row_id = Message.create(message_model)
|
|
return row_id
|
|
|
|
# langgraph interface
|
|
def run_graph(self):
|
|
self.graph.checkpointer.delete_thread(self.module_name)
|
|
config = RunnableConfig()
|
|
config["configurable"] = {"thread_id": self.module_name}
|
|
if self.agent_state is None:
|
|
self.agent_state = BaseAgentState()
|
|
for event in self.graph.stream(self.agent_state, config, stream_mode="values"):
|
|
self.logger.debug(event)
|
|
|
|
def run(self):
|
|
self.run_graph()
|
|
return self.agent_state
|