Merge pull request #227 from stacklok/add-notification-sse

feat: add alerts sse notification endpoint

Author: peppescg

src/codegate/dashboard/dashboard.py

@@ -1,15 +1,16 @@
 import asyncio
-from typing import List
+from typing import List, AsyncGenerator
 
 import structlog
 from fastapi import APIRouter
+from fastapi.responses import StreamingResponse
 
 from codegate.dashboard.post_processing import (
     parse_get_alert_conversation,
     parse_messages_in_conversations,
 )
 from codegate.dashboard.request_models import AlertConversation, Conversation
-from codegate.db.connection import DbReader
+from codegate.db.connection import DbReader, alert_queue
 
 logger = structlog.get_logger("codegate")
 
@@ -34,3 +35,19 @@ def get_alerts() -> List[AlertConversation]:
     """
     alerts_prompt_output = asyncio.run(db_reader.get_alerts_with_prompt_and_output())
     return asyncio.run(parse_get_alert_conversation(alerts_prompt_output))
+
+
+async def generate_sse_events() -> AsyncGenerator[str, None]:
+    """
+    SSE generator from queue
+    """
+    while True:
+        message = await alert_queue.get()
+        yield f"data: {message}\n\n"
+
+@dashboard_router.get("/dashboard/alerts_notification")
+async def stream_sse():
+    """
+    Send alerts event
+    """
+    return StreamingResponse(generate_sse_events(), media_type="text/event-stream")

src/codegate/db/connection.py

@@ -5,7 +5,6 @@
 import uuid
 from pathlib import Path
 from typing import AsyncGenerator, AsyncIterator, List, Optional
-
 import structlog
 from litellm import ChatCompletionRequest, ModelResponse
 from pydantic import BaseModel
@@ -20,7 +19,7 @@
 )
 
 logger = structlog.get_logger("codegate")
-
+alert_queue = asyncio.Queue()
 
 class DbCodeGate:
 
@@ -213,7 +212,9 @@ async def record_alerts(self, alerts: List[Alert]) -> None:
         async with asyncio.TaskGroup() as tg:
             for alert in alerts:
                 try:
-                    tg.create_task(self._insert_pydantic_model(alert, sql))
+                    result = tg.create_task(self._insert_pydantic_model(alert, sql))
+                    if result and alert.trigger_category == "critical":
+                        await alert_queue.put(f"New alert detected: {alert.timestamp}")
                 except Exception as e:
                     logger.error(f"Failed to record alert: {alert}.", error=str(e))
         return None