Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Fix API alerts and storing of PII alerts #1075

Merged
merged 1 commit into from
Feb 17, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/codegate/api/v1.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/codegate/api/v1_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
13 changes: 9 additions & 4 deletions src/codegate/db/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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] = {}
Expand Down
2 changes: 1 addition & 1 deletion src/codegate/pipeline/pii/analyzer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
9 changes: 6 additions & 3 deletions src/codegate/pipeline/pii/manager.py
Original file line number Diff line number Diff line change
@@ -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")
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion src/codegate/pipeline/pii/pii.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Empty file added tests/api/__init__.py
Empty file.
161 changes: 142 additions & 19 deletions tests/api/test_v1_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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]

Expand All @@ -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)
Loading
Loading