diff --git a/PLAYBOOKS/Alert_Summary_Agent.py b/PLAYBOOKS/Alert_Summary_Agent.py index 5d99f68..d963efb 100644 --- a/PLAYBOOKS/Alert_Summary_Agent.py +++ b/PLAYBOOKS/Alert_Summary_Agent.py @@ -107,12 +107,11 @@ class Playbook(LanggraphPlaybook): indent=2) if state.threat_intel_data else "No threat intelligence data available." logs_summary = [log.model_dump_json() for log in state.logs] if state.logs else [] - human_template = self.load_human_prompt_template("summary_agent_human") - human_message = HumanMessage(content=human_template.format( + human_message = self.load_human_prompt_template("summary_agent_human").format( alert_data=alert.model_dump_json(), threat_intel=ti_summary, siem_logs=logs_summary - )) + ) llm_api = LLMAPI() llm = llm_api.get_model(tag="fast") diff --git a/PLUGINS/Mock/SIEM/main.py b/PLUGINS/Mock/SIEM/main.py index 14d241b..71909ad 100644 --- a/PLUGINS/Mock/SIEM/main.py +++ b/PLUGINS/Mock/SIEM/main.py @@ -19,6 +19,7 @@ class ELKSender: def __init__(self): self.url = f"{CONFIG.ELK_HOST}/_bulk" self.auth = (CONFIG.ELK_USER, CONFIG.ELK_PASS) + self.client = httpx.Client(verify=False, timeout=30.0) def send(self, batch, index_name): payload = "" @@ -26,15 +27,14 @@ class ELKSender: payload += json.dumps({"index": {"_index": index_name}}) + "\n" payload += json.dumps(doc) + "\n" - with httpx.Client(verify=False) as client: - resp = client.post( - self.url, - content=payload, - auth=self.auth, - headers={"Content-Type": "application/x-ndjson"} - ) - if resp.status_code >= 400: - raise Exception(f"ELK Error: {resp.text}") + resp = self.client.post( + self.url, + content=payload, + auth=self.auth, + headers={"Content-Type": "application/x-ndjson"} + ) + if resp.status_code >= 400: + raise Exception(f"ELK Error: {resp.text}") class SplunkSender: