Fixe on Pydantic model for alerts

src/codegate/api/v1.py

@@ -12,7 +12,7 @@
 from codegate import __version__
 from codegate.api import v1_models, v1_processing
 from codegate.db.connection import AlreadyExistsError, DbReader
-from codegate.pipeline.base import AlertSeverity
+from codegate.db.models import AlertSeverity
 from codegate.providers import crud as provendcrud
 from codegate.workspaces import crud
 

src/codegate/api/v1_models.py

@@ -1,4 +1,5 @@
 import datetime
+import json
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
@@ -154,11 +155,16 @@ class Alert(pydantic.BaseModel):
 
     @staticmethod
     def from_db_model(db_model: db_models.Alert) -> "Alert":
+        try:
+            trigger_string = json.loads(db_model.trigger_string)
+        except Exception:
+            trigger_string = db_model.trigger_string
+        snippet = json.loads(db_model.code_snippet) if db_model.code_snippet else None
         return Alert(
             id=db_model.id,
             prompt_id=db_model.prompt_id,
-            code_snippet=db_model.code_snippet,
-            trigger_string=db_model.trigger_string,
+            code_snippet=snippet,
+            trigger_string=trigger_string,
             trigger_type=db_model.trigger_type,
             trigger_category=db_model.trigger_category,
             timestamp=db_model.timestamp,
@@ -169,7 +175,7 @@ def from_db_model(db_model: db_models.Alert) -> "Alert":
     code_snippet: Optional[CodeSnippet]
     trigger_string: Optional[Union[str, dict]]
     trigger_type: str
-    trigger_category: Optional[str]
+    trigger_category: db_models.AlertSeverity
     timestamp: datetime.datetime
 
 

src/codegate/db/fim_cache.py

@@ -8,8 +8,8 @@
 from pydantic import BaseModel
 
 from codegate.config import Config
-from codegate.db.models import Alert
-from codegate.pipeline.base import AlertSeverity, PipelineContext
+from codegate.db.models import Alert, AlertSeverity
+from codegate.pipeline.base import PipelineContext
 
 logger = structlog.get_logger("codegate")
 

src/codegate/db/models.py

@@ -5,14 +5,19 @@
 from pydantic import BaseModel, StringConstraints
 
 
+class AlertSeverity(str, Enum):
+    INFO = "info"
+    CRITICAL = "critical"
+
+
 class Alert(BaseModel):
-    id: Any
-    prompt_id: Any
-    code_snippet: Optional[Any]
-    trigger_string: Optional[Any]
-    trigger_type: Any
-    trigger_category: Optional[Any]
-    timestamp: Any
+    id: str
+    prompt_id: str
+    code_snippet: Optional[str]
+    trigger_string: Optional[str]
+    trigger_type: str
+    trigger_category: AlertSeverity
+    timestamp: datetime.datetime
 
 
 class Output(BaseModel):

src/codegate/pipeline/base.py

@@ -1,29 +1,22 @@
-import dataclasses
 import datetime
 import json
 import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Any, Dict, List, Optional
 
 import structlog
 from litellm import ChatCompletionRequest, ModelResponse
 from pydantic import BaseModel
 
 from codegate.clients.clients import ClientType
-from codegate.db.models import Alert, Output, Prompt
+from codegate.db.models import Alert, AlertSeverity, Output, Prompt
 from codegate.extract_snippets.message_extractor import CodeSnippet
 from codegate.pipeline.secrets.manager import SecretsManager
 
 logger = structlog.get_logger("codegate")
 
 
-class AlertSeverity(Enum):
-    INFO = "info"
-    CRITICAL = "critical"
-
-
 @dataclass
 class PipelineSensitiveData:
     manager: SecretsManager
@@ -80,7 +73,7 @@ def add_alert(
             logger.warning("No code snippet or trigger string provided for alert. Will not create")
             return
 
-        code_snippet_str = json.dumps(dataclasses.asdict(code_snippet)) if code_snippet else None
+        code_snippet_str = code_snippet.model_dump_json() if code_snippet else None
 
         self.alerts_raised.append(
             Alert(

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -5,9 +5,9 @@
 from litellm import ChatCompletionRequest
 
 from codegate.clients.clients import ClientType
+from codegate.db.models import AlertSeverity
 from codegate.extract_snippets.factory import MessageCodeExtractorFactory
 from codegate.pipeline.base import (
-    AlertSeverity,
     PipelineContext,
     PipelineResult,
     PipelineStep,

src/codegate/pipeline/comment/output.py

@@ -5,11 +5,12 @@
 from litellm import ModelResponse
 from litellm.types.utils import Delta, StreamingChoices
 
+from codegate.db.models import AlertSeverity
 from codegate.extract_snippets.message_extractor import (
     CodeSnippet,
     DefaultCodeSnippetExtractor,
 )
-from codegate.pipeline.base import AlertSeverity, PipelineContext
+from codegate.pipeline.base import PipelineContext
 from codegate.pipeline.output import OutputPipelineContext, OutputPipelineStep
 from codegate.storage import StorageEngine
 from codegate.utils.package_extractor import PackageExtractor

src/codegate/pipeline/secrets/secrets.py

@@ -2,14 +2,14 @@
 from abc import abstractmethod
 from typing import List, Optional, Tuple
 
-from codegate.extract_snippets.factory import MessageCodeExtractorFactory
 import structlog
 from litellm import ChatCompletionRequest, ChatCompletionSystemMessage, ModelResponse
 from litellm.types.utils import Delta, StreamingChoices
 
 from codegate.config import Config
+from codegate.db.models import AlertSeverity
+from codegate.extract_snippets.factory import MessageCodeExtractorFactory
 from codegate.pipeline.base import (
-    AlertSeverity,
     CodeSnippet,
     PipelineContext,
     PipelineResult,

tests/db/test_fim_cache.py

@@ -6,8 +6,8 @@
 import pytest
 
 from codegate.db.fim_cache import CachedFim, FimCache
-from codegate.db.models import Alert
-from codegate.pipeline.base import AlertSeverity, PipelineContext
+from codegate.db.models import Alert, AlertSeverity
+from codegate.pipeline.base import PipelineContext
 
 fim_python_message = """
 # Path: folder/testing_file.py