@cryptotaxi247 / CoPilot / commits / 70083b9a

550 alert engine tuning (#551)

* Enhance auto alert creation with batch processing and improved response model * precommit fixes

taylor_socfortress committed Dec 4, 2025 at 08:47 UTC 70083b9acd08c2b35ac5aaa0bae2310bc0cc01ef
3 files changed +218 -35
backend/app/incidents/routes/incident_alert.py
+119 -27
@@ -1,4 +1,5 @@
1 import os
2 +from typing import Optional
3
4 from fastapi import APIRouter
5 from fastapi import Depends
@@ -162,47 +163,138 @@ async def create_alert_manual_route(
163 return CreateAlertResponse(success=True, message="Alert created in CoPilot", alert_id=await create_alert(create_alert_request, session))
164
165
166 +# @incidents_alerts_router.post(
167 +# "/create/auto",
168 +# response_model=CreateAlertResponse,
169 +# description="Is invoked by the scheduler to create an incident alert in CoPilot",
170 +# )
171 +# async def create_alert_auto_route(
172 +# session: AsyncSession = Depends(get_db),
173 +# ) -> AutoCreateAlertResponse:
174 +# """
175 +# Create an incident alert in CoPilot. Automatically create an incident alert within CoPilot.
176 +# This queries the `gl-events-*` indices for alerts that have not been created in CoPilot.
177 +# It is important to note that Graylog must be configured for the alerts.
178 +
179 +# Args:
180 +# create_alert_request (CreateAlertRequest): The request object containing the details of the alert to be created.
181 +# session (AsyncSession, optional): The database session. Defaults to Depends(get_session).
182 +
183 +# Returns:
184 +# CreateAlertResponse: The response object containing the result of the alert creation.
185 +# """
186 +# alerts = await get_alerts_not_created_in_copilot()
187 +# logger.info(f"Alerts to create in CoPilot: {alerts}")
188 +# if len(alerts.alerts) == 0:
189 +# return AutoCreateAlertResponse(success=False, message="No alerts to create in CoPilot")
190 +
191 +# created_alerts_count = 0
192 +
193 +# for alert in alerts.alerts:
194 +# try:
195 +# logger.info(f"Creating alert {alert} in CoPilot")
196 +# create_alert_request = CreateAlertRequest(
197 +# index_name=await get_original_alert_index_name(origin_context=alert.source.origin_context),
198 +# alert_id=await get_original_alert_id(alert.source.origin_context),
199 +# )
200 +# logger.info(f"Creating alert {create_alert_request.alert_id} in CoPilot")
201 +# alert_id = await create_alert(create_alert_request, session)
202 +# # ! ADD THE COPILOT ALERT ID TO GRAYLOG EVENT INDEX # !
203 +# await add_copilot_alert_id(index_data=CreateAlertRequest(index_name=alert.index, alert_id=alert.id), alert_id=alert_id)
204 +# created_alerts_count += 1
205 +# except Exception as e:
206 +# logger.error(f"Failed to create alert {alert} in CoPilot: {e}")
207 +
208 +
209 @incidents_alerts_router.post(
210 "/create/auto",
167 - response_model=CreateAlertResponse,
211 + response_model=AutoCreateAlertResponse,
212 description="Is invoked by the scheduler to create an incident alert in CoPilot",
213 )
214 async def create_alert_auto_route(
215 + batch_size: Optional[int] = Query(100, ge=10, le=500, description="Number of alerts to process per batch"),
216 + max_batches: Optional[int] = Query(10, ge=1, le=50, description="Maximum number of batches to process in one run"),
217 session: AsyncSession = Depends(get_db),
218 ) -> AutoCreateAlertResponse:
219 """
174 - Create an incident alert in CoPilot. Automatically create an incident alert within CoPilot.
220 + Create incident alerts in CoPilot in batches. Automatically create incident alerts within CoPilot.
221 This queries the `gl-events-*` indices for alerts that have not been created in CoPilot.
176 - It is important to note that Graylog must be configured for the alerts.
222 +
223 + Processing is done in batches to prevent memory issues with large numbers of alerts.
224 + The scheduler will call this endpoint multiple times until all alerts are processed.
225
226 Args:
179 - create_alert_request (CreateAlertRequest): The request object containing the details of the alert to be created.
180 - session (AsyncSession, optional): The database session. Defaults to Depends(get_session).
227 + batch_size: Number of alerts to process per batch (default 100, max 500)
228 + max_batches: Maximum number of batches to process in one scheduler run (default 10, max 50)
229 + session (AsyncSession): The database session.
230
231 Returns:
183 - CreateAlertResponse: The response object containing the result of the alert creation.
232 + AutoCreateAlertResponse: The response object containing the result of the alert creation.
233 """
185 - alerts = await get_alerts_not_created_in_copilot()
186 - logger.info(f"Alerts to create in CoPilot: {alerts}")
187 - if len(alerts.alerts) == 0:
188 - return AutoCreateAlertResponse(success=False, message="No alerts to create in CoPilot")
189 -
190 - created_alerts_count = 0
191 -
192 - for alert in alerts.alerts:
193 - try:
194 - logger.info(f"Creating alert {alert} in CoPilot")
195 - create_alert_request = CreateAlertRequest(
196 - index_name=await get_original_alert_index_name(origin_context=alert.source.origin_context),
197 - alert_id=await get_original_alert_id(alert.source.origin_context),
198 - )
199 - logger.info(f"Creating alert {create_alert_request.alert_id} in CoPilot")
200 - alert_id = await create_alert(create_alert_request, session)
201 - # ! ADD THE COPILOT ALERT ID TO GRAYLOG EVENT INDEX # !
202 - await add_copilot_alert_id(index_data=CreateAlertRequest(index_name=alert.index, alert_id=alert.id), alert_id=alert_id)
203 - created_alerts_count += 1
204 - except Exception as e:
205 - logger.error(f"Failed to create alert {alert} in CoPilot: {e}")
234 + total_created = 0
235 + total_failed = 0
236 + batches_processed = 0
237 +
238 + logger.info(f"Starting auto alert creation with batch_size={batch_size}, max_batches={max_batches}")
239 +
240 + for batch_num in range(max_batches):
241 + # Fetch the next batch
242 + alerts_payload, total_remaining = await get_alerts_not_created_in_copilot(batch_size=batch_size)
243 +
244 + if len(alerts_payload.alerts) == 0:
245 + logger.info(f"No more alerts to process after {batches_processed} batches")
246 + break
247 +
248 + logger.info(f"Processing batch {batch_num + 1}/{max_batches}: {len(alerts_payload.alerts)} alerts.)")
249 + logger.info(f"Total remaining alerts after this batch: {total_remaining}")
250 +
251 + # Process this batch
252 + batch_created = 0
253 + batch_failed = 0
254 +
255 + for alert in alerts_payload.alerts:
256 + try:
257 + create_alert_request = CreateAlertRequest(
258 + index_name=await get_original_alert_index_name(origin_context=alert.source.origin_context),
259 + alert_id=await get_original_alert_id(alert.source.origin_context),
260 + )
261 +
262 + alert_id = await create_alert(create_alert_request, session)
263 +
264 + # Add the CoPilot alert ID to Graylog event index
265 + await add_copilot_alert_id(index_data=CreateAlertRequest(index_name=alert.index, alert_id=alert.id), alert_id=alert_id)
266 +
267 + batch_created += 1
268 + total_created += 1
269 +
270 + except Exception as e:
271 + logger.error(f"Failed to create alert {alert.id} from index {alert.index}: {e}")
272 + batch_failed += 1
273 + total_failed += 1
274 +
275 + batches_processed += 1
276 + logger.info(f"Batch {batch_num + 1} complete: {batch_created} created, {batch_failed} failed")
277 +
278 + # If we processed fewer alerts than the batch size, we're done
279 + if len(alerts_payload.alerts) < batch_size:
280 + logger.info("Processed final batch (fewer alerts than batch size)")
281 + break
282 +
283 + message = f"Processed {batches_processed} batches: {total_created} alerts created, {total_failed} failed"
284 +
285 + if total_remaining > 0:
286 + message += f". {total_remaining} alerts remaining for next run"
287 +
288 + logger.info(message)
289 +
290 + return AutoCreateAlertResponse(
291 + success=True,
292 + message=message,
293 + alerts_created=total_created,
294 + alerts_failed=total_failed,
295 + batches_processed=batches_processed,
296 + alerts_remaining=max(0, total_remaining - len(alerts_payload.alerts)) if batches_processed < max_batches else total_remaining,
297 + )
298
299
300 @incidents_alerts_router.post(
backend/app/incidents/schema/incident_alert.py
+21
@@ -44,9 +44,30 @@ class CreateAlertResponse(BaseModel):
44 alert_id: int = Field(..., description="The alert id as created in CoPilot.")
45
46
47 +# class AutoCreateAlertResponse(BaseModel):
48 +# success: bool
49 +# message: str
50 +
51 +
52 class AutoCreateAlertResponse(BaseModel):
53 success: bool
54 message: str
55 + alerts_created: int = 0
56 + alerts_failed: int = 0
57 + batches_processed: int = 0
58 + alerts_remaining: int = 0
59 +
60 + class Config:
61 + json_schema_extra = {
62 + "example": {
63 + "success": True,
64 + "message": "Processed 5 batches: 487 alerts created, 13 failed. 2000 alerts remaining for next run",
65 + "alerts_created": 487,
66 + "alerts_failed": 13,
67 + "batches_processed": 5,
68 + "alerts_remaining": 2000,
69 + },
70 + }
71
72
73 class IndexNamesResponse(BaseModel):
backend/app/incidents/services/alert_collection.py
+78 -8
@@ -1,3 +1,6 @@
1 +from typing import List
2 +from typing import Tuple
3 +
4 from loguru import logger
5
6 from app.connectors.wazuh_indexer.utils.universal import (
@@ -57,22 +60,89 @@ async def fetch_alerts_for_index(es_client, index, query):
60 return [AlertPayloadItem(**hit) for hit in hits]
61
62
60 -async def get_alerts_not_created_in_copilot() -> AlertsPayload:
63 +async def fetch_alerts_batch(es_client, index: str, query: dict, batch_size: int = 100) -> Tuple[List[AlertPayloadItem], int]:
64 + """
65 + Fetches a single batch of alerts for a given index that match the query.
66 +
67 + Args:
68 + es_client: Wazuh Indexer client
69 + index: Index name to query
70 + query: Query to execute
71 + batch_size: Number of alerts to fetch (default 100)
72 +
73 + Returns:
74 + Tuple of (list of alerts, total count of matching documents)
75 + """
76 + try:
77 + response = await es_client.search(
78 + index=index,
79 + body=query,
80 + size=batch_size,
81 + sort=[{"@timestamp": {"order": "asc"}}], # Process oldest first
82 + )
83 +
84 + hits = response["hits"]["hits"]
85 + total = response["hits"]["total"]["value"] if isinstance(response["hits"]["total"], dict) else response["hits"]["total"]
86 +
87 + logger.info(f"Fetched {len(hits)} alerts from index {index}. Total available: {total}")
88 +
89 + return [AlertPayloadItem(**hit) for hit in hits], total
90 + except Exception as e:
91 + logger.error(f"Error fetching alerts from index {index}: {e}")
92 + return [], 0
93 +
94 +
95 +# async def get_alerts_not_created_in_copilot() -> AlertsPayload:
96 +# """
97 +# Get the Graylog event indices. Then get all the results from the list of indices, where `copilot_alert_id` does not exist.
98 +# """
99 +# indices = await return_graylog_events_index_names()
100 +# logger.info(f"Indices: {indices}")
101 +# es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
102 +# query = await construct_query()
103 +
104 +# alerts_not_created = []
105 +# for index in indices:
106 +# alerts = await fetch_alerts_for_index(es_client, index, query)
107 +# alerts_not_created.extend(alerts)
108 +
109 +# logger.info(f"Alerts not created: {len(alerts_not_created)} alerts found")
110 +# return AlertsPayload(alerts=alerts_not_created)
111 +
112 +
113 +async def get_alerts_not_created_in_copilot(batch_size: int = 100) -> Tuple[AlertsPayload, int]:
114 """
62 - Get the Graylog event indices. Then get all the results from the list of indices, where `copilot_alert_id` does not exist.
115 + Get a batch of alerts that have not been created in CoPilot yet.
116 +
117 + Args:
118 + batch_size: Maximum number of alerts to return (default 100)
119 +
120 + Returns:
121 + Tuple of (AlertsPayload with alerts, total count remaining)
122 """
123 indices = await return_graylog_events_index_names()
65 - logger.info(f"Indices: {indices}")
124 + logger.info(f"Checking indices: {indices}")
125 +
126 es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
127 query = await construct_query()
128
69 - alerts_not_created = []
129 + alerts_to_process = []
130 + total_remaining = 0
131 +
132 + # Fetch from each index until we have enough alerts or run out
133 for index in indices:
71 - alerts = await fetch_alerts_for_index(es_client, index, query)
72 - alerts_not_created.extend(alerts)
134 + if len(alerts_to_process) >= batch_size:
135 + break
136 +
137 + remaining_to_fetch = batch_size - len(alerts_to_process)
138 + alerts, index_total = await fetch_alerts_batch(es_client, index, query, remaining_to_fetch)
139 +
140 + alerts_to_process.extend(alerts)
141 + total_remaining += index_total
142 +
143 + logger.info(f"Returning {len(alerts_to_process)} alerts. Total remaining across all indices: {total_remaining}")
144
74 - logger.info(f"Alerts not created: {len(alerts_not_created)} alerts found")
75 - return AlertsPayload(alerts=alerts_not_created)
145 + return AlertsPayload(alerts=alerts_to_process), total_remaining
146
147
148 async def get_original_alert_id(origin_context: str):