From 0e85ed776a9f09d818c20f22bdb7e594106e5e3e Mon Sep 17 00:00:00 2001 From: Alejandro Ponce de Leon Date: Fri, 7 Feb 2025 11:15:10 +0200 Subject: [PATCH] Fixe on Pydantic model for alerts (#976) Closes: #969 Some small fixes all related with fixing alerts both when adding from secrets and when retrieving them through the endpoint `http://localhost:8989/api/v1/workspaces/openai/messages` --- src/codegate/api/v1.py | 2 +- src/codegate/api/v1_models.py | 12 +++++++++--- src/codegate/db/fim_cache.py | 4 ++-- src/codegate/db/models.py | 19 ++++++++++++------- src/codegate/pipeline/base.py | 11 ++--------- .../codegate_context_retriever/codegate.py | 2 +- src/codegate/pipeline/comment/output.py | 3 ++- src/codegate/pipeline/secrets/secrets.py | 4 ++-- tests/db/test_fim_cache.py | 4 ++-- 9 files changed, 33 insertions(+), 28 deletions(-) diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index edec0e60..3b837a8b 100644 --- a/src/codegate/api/v1.py +++ b/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 diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py index d895edd8..1b883df9 100644 --- a/src/codegate/api/v1_models.py +++ b/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 diff --git a/src/codegate/db/fim_cache.py b/src/codegate/db/fim_cache.py index 482bbe7a..a0b3e9e4 100644 --- a/src/codegate/db/fim_cache.py +++ b/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") diff --git a/src/codegate/db/models.py b/src/codegate/db/models.py index 1841bc5b..2dd7568c 100644 --- a/src/codegate/db/models.py +++ b/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): diff --git a/src/codegate/pipeline/base.py b/src/codegate/pipeline/base.py index 5b8b87eb..0baa322a 100644 --- a/src/codegate/pipeline/base.py +++ b/src/codegate/pipeline/base.py @@ -1,10 +1,8 @@ -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 @@ -12,18 +10,13 @@ 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( diff --git a/src/codegate/pipeline/codegate_context_retriever/codegate.py b/src/codegate/pipeline/codegate_context_retriever/codegate.py index f373e874..bf3e95a8 100644 --- a/src/codegate/pipeline/codegate_context_retriever/codegate.py +++ b/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, diff --git a/src/codegate/pipeline/comment/output.py b/src/codegate/pipeline/comment/output.py index 6210ee13..3a01bf99 100644 --- a/src/codegate/pipeline/comment/output.py +++ b/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 diff --git a/src/codegate/pipeline/secrets/secrets.py b/src/codegate/pipeline/secrets/secrets.py index 63279c41..2fc69d3a 100644 --- a/src/codegate/pipeline/secrets/secrets.py +++ b/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, diff --git a/tests/db/test_fim_cache.py b/tests/db/test_fim_cache.py index 5e2ad547..aa228dd5 100644 --- a/tests/db/test_fim_cache.py +++ b/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