From 9db2d55633e59b2ab4ae4fd098afac1419c51240 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Mon, 17 Feb 2025 15:52:08 +0200 Subject: [PATCH] Fix API alerts 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. --- src/codegate/api/v1.py | 4 +- src/codegate/api/v1_processing.py | 2 +- src/codegate/db/connection.py | 13 +- src/codegate/pipeline/pii/analyzer.py | 2 +- src/codegate/pipeline/pii/manager.py | 9 +- src/codegate/pipeline/pii/pii.py | 2 +- tests/api/__init__.py | 0 tests/api/test_v1_processing.py | 161 ++++++++++++++++++++++--- tests/pipeline/pii/test_pii_manager.py | 4 +- 9 files changed, 165 insertions(+), 32 deletions(-) create mode 100644 tests/api/__init__.py diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index 3b837a8b..c5ac57d5 100644 --- a/src/codegate/api/v1.py +++ b/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 diff --git a/src/codegate/api/v1_processing.py b/src/codegate/api/v1_processing.py index fc902d59..0dbce577 100644 --- a/src/codegate/api/v1_processing.py +++ b/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 diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index f14c0bc8..9c9abd79 100644 --- a/src/codegate/db/connection.py +++ b/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] = {} diff --git a/src/codegate/pipeline/pii/analyzer.py b/src/codegate/pipeline/pii/analyzer.py index be854aae..a1ed5bed 100644 --- a/src/codegate/pipeline/pii/analyzer.py +++ b/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 = [ diff --git a/src/codegate/pipeline/pii/manager.py b/src/codegate/pipeline/pii/manager.py index e4ac69ab..54112713 100644 --- a/src/codegate/pipeline/pii/manager.py +++ b/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: diff --git a/src/codegate/pipeline/pii/pii.py b/src/codegate/pipeline/pii/pii.py index d4581cd9..36b3e1a1 100644 --- a/src/codegate/pipeline/pii/pii.py +++ b/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) diff --git a/tests/api/__init__.py b/tests/api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/api/test_v1_processing.py b/tests/api/test_v1_processing.py index 28d70fe1..ad8ffcbd 100644 --- a/tests/api/test_v1_processing.py +++ b/tests/api/test_v1_processing.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,7 +282,7 @@ 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", @@ -289,21 +290,21 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row): 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", @@ -311,7 +312,7 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row): type="chat", ), # Another superset - PartialQuestions( + v1_models.PartialQuestions( messages=["apple", "banana", "cherry"], timestamp=datetime.datetime(2023, 5, 2, 10, 0, 0), message_id="pqS2", @@ -319,14 +320,14 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row): 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", @@ -334,7 +335,7 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row): 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", @@ -342,7 +343,7 @@ async def test_get_question_answer(request_msg_list, output_msg_str, row): 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) diff --git a/tests/pipeline/pii/test_pii_manager.py b/tests/pipeline/pii/test_pii_manager.py index 1cd5ab30..229b7314 100644 --- a/tests/pipeline/pii/test_pii_manager.py +++ b/tests/pipeline/pii/test_pii_manager.py @@ -44,7 +44,7 @@ def test_analyze_no_pii(self, manager, mock_analyzer): assert anonymized_text == text assert found_pii == [] assert manager.session_store is session_store - mock_analyzer.analyze.assert_called_once_with(text) + mock_analyzer.analyze.assert_called_once_with(text, context=None) def test_analyze_with_pii(self, manager, mock_analyzer): text = "My email is test@example.com" @@ -71,7 +71,7 @@ def test_analyze_with_pii(self, manager, mock_analyzer): assert found_pii == pii_details assert manager.session_store is session_store assert manager.session_store.mappings[placeholder] == "test@example.com" - mock_analyzer.analyze.assert_called_once_with(text) + mock_analyzer.analyze.assert_called_once_with(text, context=None) def test_restore_pii_no_session(self, manager, mock_analyzer): text = "Anonymized text"