Skip to content

Commit

Permalink
Fix alert type code-comment
Browse files Browse the repository at this point in the history
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`
  • Loading branch information
aponcedeleonch committed Jan 24, 2025
1 parent e05f49a commit c2240f8
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions src/codegate/pipeline/extract_snippets/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from litellm import ModelResponse
from litellm.types.utils import Delta, StreamingChoices

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

# Add an alert to the context
context.add_alert(
self.name, trigger_string=comment, severity_category=AlertSeverity.CRITICAL
)

return comment

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

# Add an alert to the context
input_context.add_alert(self.name, trigger_string=complete_comment)

return chunks

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

0 comments on commit c2240f8

Please sign in to comment.