Skip to content

Commit 17e2c33

Browse files
Fix alert type code-comment (#759)
Right now it was being recorded every time that the step was executed. With this PR now it's only recorded when a malicious package is found and we elevate the severity type to `critical`
1 parent 81061c0 commit 17e2c33

File tree

1 file changed

+6
-4
lines changed
  • src/codegate/pipeline/extract_snippets

1 file changed

+6
-4
lines changed

src/codegate/pipeline/extract_snippets/output.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from litellm import ModelResponse
66
from litellm.types.utils import Delta, StreamingChoices
77

8-
from codegate.pipeline.base import CodeSnippet, PipelineContext
8+
from codegate.pipeline.base import AlertSeverity, CodeSnippet, PipelineContext
99
from codegate.pipeline.extract_snippets.extract_snippets import extract_snippets
1010
from codegate.pipeline.output import OutputPipelineContext, OutputPipelineStep
1111
from codegate.storage import StorageEngine
@@ -85,6 +85,11 @@ async def _snippet_comment(self, snippet: CodeSnippet, context: PipelineContext)
8585
archived packages: {libobjects_text}\n"
8686
comment += "\n### 🚨 Warnings\n" + "\n".join(warnings) + "\n"
8787

88+
# Add an alert to the context
89+
context.add_alert(
90+
self.name, trigger_string=comment, severity_category=AlertSeverity.CRITICAL
91+
)
92+
8893
return comment
8994

9095
def _split_chunk_at_code_end(self, content: str) -> tuple[str, str]:
@@ -147,9 +152,6 @@ async def process_chunk(
147152
chunks.append(self._create_chunk(chunk, after))
148153
complete_comment += after
149154

150-
# Add an alert to the context
151-
input_context.add_alert(self.name, trigger_string=complete_comment)
152-
153155
return chunks
154156

155157
# Pass through all other content that does not create a new snippet

0 commit comments

Comments
 (0)