| 1 | from typing import Dict |
| 2 | from typing import List |
| 3 | from typing import Optional |
| 4 | |
| 5 | from pydantic import BaseModel |
| 6 | from pydantic import Field |
| 7 | from pydantic import model_validator |
| 8 | |
| 9 | |
| 10 | class Fields(BaseModel): |
| 11 | ALERT_ID: Optional[str] = None |
| 12 | ALERT_SOURCE: Optional[str] = None |
| 13 | CUSTOMER_CODE: Optional[str] = None |
| 14 | COPILOT_ALERT_ID: Optional[str] = None |
| 15 | |
| 16 | |
| 17 | class Source(BaseModel): |
| 18 | id: str |
| 19 | event_definition_type: str |
| 20 | event_definition_id: str |
| 21 | origin_context: str |
| 22 | timestamp: str |
| 23 | timestamp_processing: str |
| 24 | timerange_start: Optional[str] = None |
| 25 | timerange_end: Optional[str] = None |
| 26 | streams: List[str] |
| 27 | source_streams: List[str] |
| 28 | message: str |
| 29 | source: str |
| 30 | key_tuple: List[str] |
| 31 | key: str |
| 32 | priority: int |
| 33 | alert: bool |
| 34 | fields: Fields |
| 35 | group_by_fields: Dict = Field(default_factory=dict) |
| 36 | original_alert_id: Optional[str] = Field(None, alias="original_alert_id") |
| 37 | original_alert_index_name: Optional[str] = Field(None, alias="original_alert_index_name") |
| 38 | |
| 39 | @model_validator(mode="before") |
| 40 | @classmethod |
| 41 | def extract_origin_context(cls, data): |
| 42 | if isinstance(data, dict): |
| 43 | origin_context = data.get("origin_context", "") |
| 44 | try: |
| 45 | # Assuming the format is always as given in the example |
| 46 | parts = origin_context.split(":") |
| 47 | if len(parts) == 6: |
| 48 | _, _, _, _, index_name, alert_id = parts |
| 49 | data["original_alert_id"] = alert_id |
| 50 | data["original_alert_index_name"] = index_name |
| 51 | except Exception as e: |
| 52 | # Consider logging the exception to understand what's going wrong |
| 53 | print(f"Error parsing origin_context: {e}") |
| 54 | return data |
| 55 | |
| 56 | |
| 57 | class AlertPayloadItem(BaseModel): |
| 58 | index: str = Field(..., alias="_index") |
| 59 | id: str = Field(..., alias="_id") |
| 60 | score: Optional[float] = Field(None, alias="_score") |
| 61 | source: Source = Field(..., alias="_source") |
| 62 | |
| 63 | |
| 64 | class AlertsPayload(BaseModel): |
| 65 | alerts: List[AlertPayloadItem] |
| 66 | success: Optional[bool] = True |
| 67 | message: Optional[str] = "Success" |