Fix API alerts (#1075)

We had some small bugs related with alerts:
1. The function for deduping alerts was only returning alerts related to secrets.
All the other alerts were being dropped when calling this function.
2. We were returning non-critical alerts in the `/messages/` endpoint
3. PII alerts were never recorded because the context was never passed to the function.
Discovered by @kd

This PR fixes the 3 of them and introduces some unit tests.

Author: aponcedeleonch

src/codegate/api/v1.py

@@ -414,7 +414,9 @@ async def get_workspace_messages(workspace_name: str) -> List[v1_models.Conversa
 
     try:
         prompts_with_output_alerts_usage = (
-            await dbreader.get_prompts_with_output_alerts_usage_by_workspace_id(ws.id)
+            await dbreader.get_prompts_with_output_alerts_usage_by_workspace_id(
+                ws.id, AlertSeverity.CRITICAL.value
+            )
         )
         conversations, _ = await v1_processing.parse_messages_in_conversations(
             prompts_with_output_alerts_usage

src/codegate/api/v1_processing.py

@@ -534,4 +534,4 @@ async def remove_duplicate_alerts(alerts: List[v1_models.Alert]) -> List[v1_mode
         seen[key] = alert
         unique_alerts.append(alert)
 
-    return list(seen.values())
+    return unique_alerts

src/codegate/db/connection.py

@@ -586,7 +586,7 @@ async def get_prompts_with_output(self, workpace_id: str) -> List[GetPromptWithO
         return prompts
 
     async def get_prompts_with_output_alerts_usage_by_workspace_id(
-        self, workspace_id: str
+        self, workspace_id: str, trigger_category: Optional[str] = None
     ) -> List[GetPromptWithOutputsRow]:
         """
         Get all prompts with their outputs, alerts and token usage by workspace_id.
@@ -602,12 +602,17 @@ async def get_prompts_with_output_alerts_usage_by_workspace_id(
             LEFT JOIN outputs o ON p.id = o.prompt_id
             LEFT JOIN alerts a ON p.id = a.prompt_id
             WHERE p.workspace_id = :workspace_id
+            AND a.trigger_category LIKE :trigger_category
             ORDER BY o.timestamp DESC, a.timestamp DESC
             """  # noqa: E501
         )
-        conditions = {"workspace_id": workspace_id}
-        rows = await self._exec_select_conditions_to_pydantic(
-            IntermediatePromptWithOutputUsageAlerts, sql, conditions, should_raise=True
+        # If trigger category is None we want to get all alerts
+        trigger_category = trigger_category if trigger_category else "%"
+        conditions = {"workspace_id": workspace_id, "trigger_category": trigger_category}
+        rows: List[IntermediatePromptWithOutputUsageAlerts] = (
+            await self._exec_select_conditions_to_pydantic(
+                IntermediatePromptWithOutputUsageAlerts, sql, conditions, should_raise=True
+            )
         )
 
         prompts_dict: Dict[str, GetPromptWithOutputsRow] = {}

src/codegate/pipeline/pii/analyzer.py

@@ -100,7 +100,7 @@ def __init__(self):
         PiiAnalyzer._instance = self
 
     def analyze(
-        self, text: str, context: Optional["PipelineContext"] = None
+        self, text: str, context: Optional[PipelineContext] = None
     ) -> Tuple[str, List[Dict[str, Any]], PiiSessionStore]:
         # Prioritize credit card detection first
         entities = [

src/codegate/pipeline/pii/manager.py

@@ -1,7 +1,8 @@
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import structlog
 
+from codegate.pipeline.base import PipelineContext
 from codegate.pipeline.pii.analyzer import PiiAnalyzer, PiiSessionStore
 
 logger = structlog.get_logger("codegate")
@@ -52,9 +53,11 @@ def session_store(self) -> PiiSessionStore:
         # Always return the analyzer's current session store
         return self.analyzer.session_store
 
-    def analyze(self, text: str) -> Tuple[str, List[Dict[str, Any]]]:
+    def analyze(
+        self, text: str, context: Optional[PipelineContext] = None
+    ) -> Tuple[str, List[Dict[str, Any]]]:
         # Call analyzer and get results
-        anonymized_text, found_pii, _ = self.analyzer.analyze(text)
+        anonymized_text, found_pii, _ = self.analyzer.analyze(text, context=context)
 
         # Log found PII details (without modifying the found_pii list)
         if found_pii:

src/codegate/pipeline/pii/pii.py

@@ -80,7 +80,7 @@ async def process(
             if "content" in message and message["content"]:
                 # This is where analyze and anonymize the text
                 original_text = str(message["content"])
-                anonymized_text, pii_details = self.pii_manager.analyze(original_text)
+                anonymized_text, pii_details = self.pii_manager.analyze(original_text, context)
 
                 if pii_details:
                     total_pii_found += len(pii_details)

tests/api/__init__.py

@@ -4,13 +4,14 @@
 
 import pytest
 
-from codegate.api.v1_models import PartialQuestions
+from codegate.api import v1_models
 from codegate.api.v1_processing import (
     _get_partial_question_answer,
     _group_partial_messages,
     _is_system_prompt,
     parse_output,
     parse_request,
+    remove_duplicate_alerts,
 )
 from codegate.db.models import GetPromptWithOutputsRow
 
@@ -193,14 +194,14 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row):
         # 1) No subsets: all items stand alone
         (
             [
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["A"],
                     timestamp=datetime.datetime(2023, 1, 1, 0, 0, 0),
                     message_id="pq1",
                     provider="providerA",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["B"],
                     timestamp=datetime.datetime(2023, 1, 1, 0, 0, 1),
                     message_id="pq2",
@@ -214,14 +215,14 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row):
         #    - "Hello" is a subset of "Hello, how are you?"
         (
             [
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["Hello"],
                     timestamp=datetime.datetime(2022, 1, 1, 0, 0, 0),
                     message_id="pq1",
                     provider="providerA",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["Hello", "How are you?"],
                     timestamp=datetime.datetime(2022, 1, 1, 0, 0, 10),
                     message_id="pq2",
@@ -238,28 +239,28 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row):
         #    superset.
         (
             [
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["Hello"],
                     timestamp=datetime.datetime(2023, 1, 1, 10, 0, 0),
                     message_id="pq1",
                     provider="providerA",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["Hello"],
                     timestamp=datetime.datetime(2023, 1, 1, 11, 0, 0),
                     message_id="pq2",
                     provider="providerA",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["Hello"],
                     timestamp=datetime.datetime(2023, 1, 1, 12, 0, 0),
                     message_id="pq3",
                     provider="providerA",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["Hello", "Bye"],
                     timestamp=datetime.datetime(2023, 1, 1, 11, 0, 5),
                     message_id="pq4",
@@ -281,68 +282,68 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row):
         (
             [
                 # Superset
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["hi", "welcome", "bye"],
                     timestamp=datetime.datetime(2023, 5, 1, 9, 0, 0),
                     message_id="pqS1",
                     provider="providerB",
                     type="chat",
                 ),
                 # Subsets for pqS1
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["hi", "welcome"],
                     timestamp=datetime.datetime(2023, 5, 1, 9, 0, 5),
                     message_id="pqA1",
                     provider="providerB",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["hi", "bye"],
                     timestamp=datetime.datetime(2023, 5, 1, 9, 0, 10),
                     message_id="pqA2",
                     provider="providerB",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["hi", "bye"],
                     timestamp=datetime.datetime(2023, 5, 1, 9, 0, 12),
                     message_id="pqA3",
                     provider="providerB",
                     type="chat",
                 ),
                 # Another superset
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["apple", "banana", "cherry"],
                     timestamp=datetime.datetime(2023, 5, 2, 10, 0, 0),
                     message_id="pqS2",
                     provider="providerB",
                     type="chat",
                 ),
                 # Subsets for pqS2
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["banana"],
                     timestamp=datetime.datetime(2023, 5, 2, 10, 0, 1),
                     message_id="pqB1",
                     provider="providerB",
                     type="chat",
                 ),
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["apple", "banana"],
                     timestamp=datetime.datetime(2023, 5, 2, 10, 0, 3),
                     message_id="pqB2",
                     provider="providerB",
                     type="chat",
                 ),
                 # Another item alone, not a subset nor superset
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["xyz"],
                     timestamp=datetime.datetime(2023, 5, 3, 8, 0, 0),
                     message_id="pqC1",
                     provider="providerB",
                     type="chat",
                 ),
                 # Different provider => should remain separate
-                PartialQuestions(
+                v1_models.PartialQuestions(
                     messages=["hi", "welcome"],
                     timestamp=datetime.datetime(2023, 5, 1, 9, 0, 10),
                     message_id="pqProvDiff",
@@ -394,7 +395,7 @@ def test_group_partial_messages(pq_list, expected_group_ids):
     # Execute
     grouped = _group_partial_messages(pq_list)
 
-    # Convert from list[list[PartialQuestions]] -> list[list[str]]
+    # Convert from list[list[v1_models.PartialQuestions]] -> list[list[str]]
     # so we can compare with expected_group_ids easily.
     grouped_ids = [[pq.message_id for pq in group] for group in grouped]
 
@@ -406,3 +407,125 @@ def test_group_partial_messages(pq_list, expected_group_ids):
                 is_matched = True
                 break
         assert is_matched
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "alerts,expected_count,expected_ids",
+    [
+        # Test Case 1: Non-secret alerts pass through unchanged
+        (
+            [
+                v1_models.Alert(
+                    id="1",
+                    prompt_id="p1",
+                    code_snippet=None,
+                    trigger_string="test1",
+                    trigger_type="other-alert",
+                    trigger_category="info",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 0),
+                ),
+                v1_models.Alert(
+                    id="2",
+                    prompt_id="p2",
+                    code_snippet=None,
+                    trigger_string="test2",
+                    trigger_type="other-alert",
+                    trigger_category="info",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 1),
+                ),
+            ],
+            2,  # Expected count
+            ["1", "2"],  # Expected IDs preserved
+        ),
+        # Test Case 2: Duplicate secrets within 5 seconds - keep newer only
+        (
+            [
+                v1_models.Alert(
+                    id="1",
+                    prompt_id="p1",
+                    code_snippet=None,
+                    trigger_string="secret1 Context xyz",
+                    trigger_type="codegate-secrets",
+                    trigger_category="critical",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 0),
+                ),
+                v1_models.Alert(
+                    id="2",
+                    prompt_id="p2",
+                    code_snippet=None,
+                    trigger_string="secret1 Context abc",
+                    trigger_type="codegate-secrets",
+                    trigger_category="critical",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 3),
+                ),
+            ],
+            1,  # Expected count
+            ["2"],  # Only newer alert ID
+        ),
+        # Test Case 3: Similar secrets beyond 5 seconds - keep both
+        (
+            [
+                v1_models.Alert(
+                    id="1",
+                    prompt_id="p1",
+                    code_snippet=None,
+                    trigger_string="secret1 Context xyz",
+                    trigger_type="codegate-secrets",
+                    trigger_category="critical",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 0),
+                ),
+                v1_models.Alert(
+                    id="2",
+                    prompt_id="p2",
+                    code_snippet=None,
+                    trigger_string="secret1 Context abc",
+                    trigger_type="codegate-secrets",
+                    trigger_category="critical",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 6),
+                ),
+            ],
+            2,  # Expected count
+            ["1", "2"],  # Both alerts preserved
+        ),
+        # Test Case 4: Mix of secret and non-secret alerts
+        (
+            [
+                v1_models.Alert(
+                    id="1",
+                    prompt_id="p1",
+                    code_snippet=None,
+                    trigger_string="secret1 Context xyz",
+                    trigger_type="codegate-secrets",
+                    trigger_category="critical",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 0),
+                ),
+                v1_models.Alert(
+                    id="2",
+                    prompt_id="p2",
+                    code_snippet=None,
+                    trigger_string="non-secret alert",
+                    trigger_type="other-alert",
+                    trigger_category="info",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 1),
+                ),
+                v1_models.Alert(
+                    id="3",
+                    prompt_id="p3",
+                    code_snippet=None,
+                    trigger_string="secret1 Context abc",
+                    trigger_type="codegate-secrets",
+                    trigger_category="critical",
+                    timestamp=datetime.datetime(2023, 1, 1, 12, 0, 3),
+                ),
+            ],
+            2,  # Expected count
+            ["2", "3"],  # Non-secret alert and newest secret alert
+        ),
+    ],
+)
+async def test_remove_duplicate_alerts(alerts, expected_count, expected_ids):
+    result = await remove_duplicate_alerts(alerts)
+    assert len(result) == expected_count
+    result_ids = [alert.id for alert in result]
+    assert sorted(result_ids) == sorted(expected_ids)