Rebase and formatting changes

Author: aponcedeleonch

src/codegate/db/connection.py

@@ -198,7 +198,6 @@ async def record_output_non_stream(
     async def record_alerts(self, alerts: List[Alert]) -> None:
         if not alerts:
             return
-        logger.info(f"Recording {len(alerts)} alerts.", alerts=alerts)
         sql = text(
             """
                 INSERT INTO alerts (

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -6,12 +6,13 @@
 from codegate.config import Config
 from codegate.inference.inference_engine import LlamaCppInferenceEngine
 from codegate.pipeline.base import (
+    AlertSeverity,
     PipelineContext,
     PipelineResult,
     PipelineStep,
 )
-from src.codegate.storage.storage_engine import StorageEngine
-from src.codegate.utils.utils import generate_vector_string
+from codegate.storage.storage_engine import StorageEngine
+from codegate.utils.utils import generate_vector_string
 
 logger = structlog.get_logger("codegate")
 
@@ -54,7 +55,7 @@ def generate_context_str(self, objects: list[object]) -> str:
         return context_str
 
     async def __lookup_packages(self, user_query: str):
-        ## Check which packages are referenced in the user query
+        # Check which packages are referenced in the user query
         request = {
             "messages": [
                 {"role": "system", "content": Config.get_config().prompts.lookup_packages},
@@ -125,16 +126,23 @@ async def process(
             # is list
             message = new_request["messages"][last_user_idx]
             if isinstance(message["content"], str):
-                message["content"] = f'Context: {context_str} \n\n Query: {message["content"]}'
+                context_msg = f'Context: {context_str} \n\n Query: {message["content"]}'
+                context.add_alert(
+                    self.name, trigger_string=context_msg, severity_category=AlertSeverity.CRITICAL
+                )
+                message["content"] = context_msg
             elif isinstance(message["content"], (list, tuple)):
                 for item in message["content"]:
                     if isinstance(item, dict) and item.get("type") == "text":
-                        item["text"] = f'Context: {context_str} \n\n Query: {item["text"]}'
-
-                    return PipelineResult(
-                        request=new_request,
-                        context=context,
-                    )
+                        context_msg = f'Context: {context_str} \n\n Query: {item["text"]}'
+                        context.add_alert(
+                            self.name,
+                            trigger_string=context_msg,
+                            severity_category=AlertSeverity.CRITICAL,
+                        )
+                        item["text"] = context_msg
+
+                    return PipelineResult(request=new_request, context=context)
 
         # Fall through
         return PipelineResult(request=request, context=context)

src/codegate/pipeline/extract_snippets/output.py

@@ -98,11 +98,6 @@ async def process_chunk(
 
             # Add an alert to the context
             input_context.add_alert(self.name, trigger_string=complete_comment)
-            logger.info(
-                f"Added alert {self.name} for new code snippet: {complete_comment}",
-                alerts=input_context.alerts_raised,
-                num=len(input_context.alerts_raised),
-            )
 
             return chunks
 

src/codegate/pipeline/secrets/secrets.py

@@ -157,7 +157,7 @@ def _redeact_text(
                 logger.info(f"Original: {secret['original']}")
                 logger.info(f"Encrypted: REDACTED<${secret['encrypted']}>")
 
-            print(f"\nProtected text:\n{protected_string}")
+            logger.info(f"\nProtected text:\n{protected_string}")
             return "".join(protected_text)
 
     async def process(

src/codegate/pipeline/system_prompt/codegate.py

@@ -1,4 +1,4 @@
-from typing import Optional
+import json
 
 from litellm import ChatCompletionRequest, ChatCompletionSystemMessage
 
@@ -16,9 +16,7 @@ class SystemPrompt(PipelineStep):
     """
 
     def __init__(self, system_prompt: str):
-        self._system_message = ChatCompletionSystemMessage(
-            content=system_prompt, role="system"
-        )
+        self._system_message = ChatCompletionSystemMessage(content=system_prompt, role="system")
 
     @property
     def name(self) -> str:
@@ -46,11 +44,17 @@ async def process(
 
         if request_system_message is None:
             # Add system message
+            context.add_alert(self.name, trigger_string=json.dumps(self._system_message))
             new_request["messages"].insert(0, self._system_message)
         elif "codegate" not in request_system_message["content"].lower():
             # Prepend to the system message
-            request_system_message["content"] = self._system_message["content"] + \
-                "\n Here are additional instructions. \n " + request_system_message["content"]
+            prepended_message = (
+                self._system_message["content"]
+                + "\n Here are additional instructions. \n "
+                + request_system_message["content"]
+            )
+            context.add_alert(self.name, trigger_string=prepended_message)
+            request_system_message["content"] = prepended_message
 
         return PipelineResult(
             request=new_request,

src/codegate/server.py

@@ -7,13 +7,13 @@
 from codegate.dashboard.dashboard import dashboard_router
 from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
 from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
-from codegate.pipeline.system_prompt.codegate import SystemPrompt
 from codegate.pipeline.extract_snippets.extract_snippets import CodeSnippetExtractor
 from codegate.pipeline.extract_snippets.output import CodeCommentStep
 from codegate.pipeline.output import OutputPipelineProcessor, OutputPipelineStep
 from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.pipeline.secrets.secrets import CodegateSecrets, SecretUnredactionStep
 from codegate.pipeline.secrets.signatures import CodegateSignatures
+from codegate.pipeline.system_prompt.codegate import SystemPrompt
 from codegate.pipeline.version.version import CodegateVersion
 from codegate.providers.anthropic.provider import AnthropicProvider
 from codegate.providers.llamacpp.provider import LlamaCppProvider

src/codegate/storage/storage_engine.py

@@ -1,9 +1,9 @@
 import structlog
 import weaviate
+import weaviate.classes as wvc
 from weaviate.classes.config import DataType
 from weaviate.classes.query import MetadataQuery
 from weaviate.embedded import EmbeddedOptions
-import weaviate.classes as wvc
 
 from codegate.config import Config
 from codegate.inference.inference_engine import LlamaCppInferenceEngine