From 54f3916e5fff8e2d60879c60e53c346396f1583d Mon Sep 17 00:00:00 2001 From: Sainath Singineedi <44405294+sainad2222@users.noreply.github.com> Date: Thu, 19 Dec 2024 16:39:09 +0530 Subject: [PATCH] feat: refactor vectordev provider to support wrapping multiple source types (#2828) --- .../grafana_provider/grafana_provider.py | 16 ++++-- .../prometheus_provider.py | 5 ++ .../vectordev_provider/vectordev_provider.py | 51 +++++++++++++------ scripts/simulate_alerts.py | 2 +- 4 files changed, 53 insertions(+), 21 deletions(-) diff --git a/keep/providers/grafana_provider/grafana_provider.py b/keep/providers/grafana_provider/grafana_provider.py index d26bf88d7..e3f958230 100644 --- a/keep/providers/grafana_provider/grafana_provider.py +++ b/keep/providers/grafana_provider/grafana_provider.py @@ -509,6 +509,8 @@ def simulate_alert(cls, **kwargs) -> dict: if not alert_type: alert_type = random.choice(list(ALERTS.keys())) + to_wrap_with_provider_type = kwargs.get("to_wrap_with_provider_type") + if "payload" in ALERTS[alert_type]: alert_payload = ALERTS[alert_type]["payload"] else: @@ -552,11 +554,15 @@ def simulate_alert(cls, **kwargs) -> dict: fingerprint = hashlib.md5(fingerprint_src.encode()).hexdigest() alert_payload["fingerprint"] = fingerprint - return { - "alerts": [alert_payload], - "severity": alert_payload.get("labels", {}).get("severity"), - "title": alert_type, - } + + final_payload = { + "alerts": [alert_payload], + "severity": alert_payload.get("labels", {}).get("severity"), + "title": alert_type, + } + if to_wrap_with_provider_type: + return {"keep_source_type": "grafana", "event": final_payload} + return final_payload if __name__ == "__main__": diff --git a/keep/providers/prometheus_provider/prometheus_provider.py b/keep/providers/prometheus_provider/prometheus_provider.py index cb3c392a7..b4dd4927e 100644 --- a/keep/providers/prometheus_provider/prometheus_provider.py +++ b/keep/providers/prometheus_provider/prometheus_provider.py @@ -233,6 +233,8 @@ def simulate_alert(cls, **kwargs) -> dict: if not alert_type: alert_type = random.choice(list(ALERTS.keys())) + to_wrap_with_provider_type = kwargs.get("to_wrap_with_provider_type") + alert_payload = ALERTS[alert_type]["payload"] alert_parameters = ALERTS[alert_type].get("parameters", []) # now generate some random data @@ -267,6 +269,9 @@ def simulate_alert(cls, **kwargs) -> dict: fingerprint_src = json.dumps(alert_payload["labels"], sort_keys=True) fingerprint = hashlib.md5(fingerprint_src.encode()).hexdigest() alert_payload["fingerprint"] = fingerprint + if to_wrap_with_provider_type: + return {"keep_source_type": "prometheus", "event": alert_payload} + return alert_payload diff --git a/keep/providers/vectordev_provider/vectordev_provider.py b/keep/providers/vectordev_provider/vectordev_provider.py index c3d6f8a9d..aeaa17aca 100644 --- a/keep/providers/vectordev_provider/vectordev_provider.py +++ b/keep/providers/vectordev_provider/vectordev_provider.py @@ -1,13 +1,20 @@ import dataclasses + +import random import json import pydantic +import logging from keep.api.models.alert import AlertDto from keep.contextmanager.contextmanager import ContextManager from keep.providers.base.base_provider import BaseProvider from keep.providers.models.provider_config import ProviderConfig +from keep.api.models.alert import AlertDto +from keep.providers.providers_factory import ProvidersFactory +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) @pydantic.dataclasses.dataclass class VectordevProviderAuthConfig: @@ -21,6 +28,11 @@ class VectordevProvider(BaseProvider): PROVIDER_CATEGORY = ["Monitoring", "Developer Tools"] PROVIDER_COMING_SOON = True + # Mapping from vector sources to keep providers + SOURCE_TO_PROVIDER_MAP = { + "prometheus": "prometheus", + } + def __init__( self, context_manager: ContextManager, provider_id: str, config: ProviderConfig ): @@ -32,33 +44,42 @@ def validate_config(self): ) def _format_alert( - event: list[dict], provider_instance: "BaseProvider" = None + event: dict, provider_instance: "BaseProvider" = None ) -> AlertDto | list[AlertDto]: events = [] - # event is a list of events - for e in event: - event_json = None - try: - event_json = json.loads(e.get("message")) - except json.JSONDecodeError: - pass - - events.append( + if isinstance(event, list): + events = event + else: + events = [event] + alert_dtos = [] + for e in events: + if "keep_source_type" in e and e["keep_source_type"] in VectordevProvider.SOURCE_TO_PROVIDER_MAP: + provider_class = ProvidersFactory.get_provider_class(VectordevProvider.SOURCE_TO_PROVIDER_MAP[e["keep_source_type"]]) + alert_dtos.extend(provider_class._format_alert(e["message"],provider_instance)) + else: + message_str = json.dumps(e.get("message")) + alert_dtos.append( AlertDto( name="", - host=e.get("host"), - message=e.get("message"), - description=e.get("message"), + message=message_str, + description=message_str, lastReceived=e.get("timestamp"), source_type=e.get("source_type"), source=["vectordev"], - original_event=event_json, + original_event=e.get("message"), ) ) - return events + return alert_dtos def dispose(self): """ No need to dispose of anything, so just do nothing. """ pass + + @classmethod + def simulate_alert(cls, **kwargs) -> dict: + provider = random.choice(list(VectordevProvider.SOURCE_TO_PROVIDER_MAP.values())) + provider_class = ProvidersFactory.get_provider_class(provider) + return provider_class.simulate_alert(to_wrap_with_provider_type=True) + diff --git a/scripts/simulate_alerts.py b/scripts/simulate_alerts.py index 29a38ca4b..45db1310b 100644 --- a/scripts/simulate_alerts.py +++ b/scripts/simulate_alerts.py @@ -36,7 +36,7 @@ async def main(): SLEEP_INTERVAL = float( os.environ.get("SLEEP_INTERVAL", default_sleep_interval) ) - keep_api_key = os.environ.get("KEEP_API_KEY") + keep_api_key = os.environ.get("KEEP_API_KEY") or "keepappkey" keep_api_url = os.environ.get("KEEP_API_URL") or "http://localhost:8080" for i in range(args.workers):