Merge pull request #197 from stacklok/setup-alerts

Record alerts as part of the pipeline

Author: aponcedeleonch

scripts/import_packages.py

@@ -1,5 +1,7 @@
 import asyncio
+from datetime import date
 import json
+import os
 
 import weaviate
 from weaviate.classes.config import DataType, Property
@@ -14,7 +16,8 @@ class PackageImporter:
     def __init__(self):
         self.client = weaviate.WeaviateClient(
             embedded_options=EmbeddedOptions(
-                persistence_data_path="./weaviate_data", grpc_port=50052
+                persistence_data_path="./weaviate_data", grpc_port=50052,
+                additional_env_vars={"ENABLE_MODULES": "backup-filesystem", "BACKUP_FILESYSTEM_PATH": os.getenv("BACKUP_FILESYSTEM_PATH", "/tmp")}
             )
         )
         self.json_files = [
@@ -86,6 +89,9 @@ async def run_import(self):
         self.setup_schema()
         await self.add_data()
 
+        #  take a backup of the data
+        await self.client.backup.create(backup_id="backup-"+date.today().strftime("%Y-%m-%d"), backend="filesystem", wait_for_completion=True)
+
 
 if __name__ == "__main__":
     importer = PackageImporter()

sql/queries/queries.sql

@@ -1,99 +1,24 @@
--- name: CreatePrompt :one
-INSERT INTO prompts (
-    id,
-    timestamp,
-    provider,
-    request,
-    type
-) VALUES (?, ?, ?, ?, ?) RETURNING *;
-
--- name: GetPrompt :one
-SELECT * FROM prompts WHERE id = ?;
-
--- name: ListPrompts :many
-SELECT * FROM prompts 
-ORDER BY timestamp DESC 
-LIMIT ? OFFSET ?;
-
--- name: CreateOutput :one
-INSERT INTO outputs (
-    id,
-    prompt_id,
-    timestamp,
-    output
-) VALUES (?, ?, ?, ?) RETURNING *;
-
--- name: GetOutput :one
-SELECT * FROM outputs WHERE id = ?;
-
--- name: GetOutputsByPromptId :many
-SELECT * FROM outputs 
-WHERE prompt_id = ? 
-ORDER BY timestamp DESC;
-
--- name: CreateAlert :one
-INSERT INTO alerts (
-    id,
-    prompt_id,
-    output_id,
-    code_snippet,
-    trigger_string,
-    trigger_type,
-    trigger_category,
-    timestamp
-) VALUES (?, ?, ?, ?, ?, ?, ?, ?) RETURNING *;
-
--- name: GetAlert :one
-SELECT * FROM alerts WHERE id = ?;
-
--- name: ListAlertsByPrompt :many
-SELECT * FROM alerts 
-WHERE prompt_id = ? 
-ORDER BY timestamp DESC;
-
--- name: GetSettings :one
-SELECT * FROM settings ORDER BY id LIMIT 1;
-
--- name: UpsertSettings :one
-INSERT INTO settings (
-    id,
-    ip,
-    port,
-    llm_model,
-    system_prompt,
-    other_settings
-) VALUES (?, ?, ?, ?, ?, ?)
-ON CONFLICT(id) DO UPDATE SET
-    ip = excluded.ip,
-    port = excluded.port,
-    llm_model = excluded.llm_model,
-    system_prompt = excluded.system_prompt,
-    other_settings = excluded.other_settings
-RETURNING *;
-
--- name: GetPromptWithOutputsAndAlerts :many
+-- name: GetPromptWithOutputs :many
 SELECT 
     p.*,
     o.id as output_id,
     o.output,
-    a.id as alert_id,
-    a.code_snippet,
-    a.trigger_string,
-    a.trigger_type,
-    a.trigger_category
+    o.timestamp as output_timestamp
 FROM prompts p
 LEFT JOIN outputs o ON p.id = o.prompt_id
-LEFT JOIN alerts a ON p.id = a.prompt_id
-WHERE p.id = ?
-ORDER BY o.timestamp DESC, a.timestamp DESC;
-
+ORDER BY o.timestamp DESC;
 
--- name: GetPromptWithOutputs :many
+-- name: GetAlertsWithPromptAndOutput :many
 SELECT 
-    p.*,
+    a.*,
+    p.timestamp as prompt_timestamp,
+    p.provider,
+    p.request,
+    p.type,
     o.id as output_id,
     o.output,
     o.timestamp as output_timestamp
-FROM prompts p
+FROM alerts a
+LEFT JOIN prompts p ON p.id = a.prompt_id
 LEFT JOIN outputs o ON p.id = o.prompt_id
-ORDER BY o.timestamp DESC;
+ORDER BY a.timestamp DESC;

sql/schema/schema.sql

@@ -22,14 +22,12 @@ CREATE TABLE outputs (
 CREATE TABLE alerts (
     id TEXT PRIMARY KEY,  -- UUID stored as TEXT
     prompt_id TEXT NOT NULL,
-    output_id TEXT NOT NULL,
-    code_snippet TEXT NOT NULL,  -- VARCHAR(255)
-    trigger_string TEXT NOT NULL, -- VARCHAR(255)
+    code_snippet TEXT,  -- We check in code that not both code_snippet and trigger_string are NULL
+    trigger_string TEXT, -- VARCHAR(255)
     trigger_type TEXT NOT NULL,   -- VARCHAR(50)
     trigger_category TEXT,
     timestamp DATETIME NOT NULL,
-    FOREIGN KEY (prompt_id) REFERENCES prompts(id),
-    FOREIGN KEY (output_id) REFERENCES outputs(id)
+    FOREIGN KEY (prompt_id) REFERENCES prompts(id)
 );
 
 -- Settings table
@@ -45,7 +43,6 @@ CREATE TABLE settings (
 -- Create indexes for foreign keys and frequently queried columns
 CREATE INDEX idx_outputs_prompt_id ON outputs(prompt_id);
 CREATE INDEX idx_alerts_prompt_id ON alerts(prompt_id);
-CREATE INDEX idx_alerts_output_id ON alerts(output_id);
 CREATE INDEX idx_prompts_timestamp ON prompts(timestamp);
 CREATE INDEX idx_outputs_timestamp ON outputs(timestamp);
 CREATE INDEX idx_alerts_timestamp ON alerts(timestamp);

src/codegate/dashboard/dashboard.py

@@ -4,8 +4,12 @@
 import structlog
 from fastapi import APIRouter
 
-from codegate.dashboard.post_processing import match_conversations, parse_get_prompt_with_output
-from codegate.dashboard.request_models import Conversation
+from codegate.dashboard.post_processing import (
+    match_conversations,
+    parse_get_alert_conversation,
+    parse_get_prompt_with_output,
+)
+from codegate.dashboard.request_models import AlertConversation, Conversation
 from codegate.db.connection import DbReader
 
 logger = structlog.get_logger("codegate")
@@ -28,3 +32,18 @@ async def get_messages() -> List[Conversation]:
 
     conversations = await match_conversations(partial_conversations)
     return conversations
+
+
+@dashboard_router.get("/dashboard/alerts")
+async def get_alerts() -> List[AlertConversation]:
+    """
+    Get all the messages from the database and return them as a list of conversations.
+    """
+    alerts_prompt_output = await db_reader.get_alerts_with_prompt_and_output()
+
+    # Parse the prompts and outputs in parallel
+    async with asyncio.TaskGroup() as tg:
+        tasks = [tg.create_task(parse_get_alert_conversation(row)) for row in alerts_prompt_output]
+    alert_conversations = [task.result() for task in tasks if task.result() is not None]
+
+    return alert_conversations

src/codegate/dashboard/post_processing.py

@@ -1,16 +1,17 @@
 import asyncio
 import json
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 import structlog
 
 from codegate.dashboard.request_models import (
+    AlertConversation,
     ChatMessage,
     Conversation,
     PartialConversation,
     QuestionAnswer,
 )
-from codegate.db.queries import GetPromptWithOutputsRow
+from codegate.db.queries import GetAlertsWithPromptAndOutputRow, GetPromptWithOutputsRow
 
 logger = structlog.get_logger("codegate")
 
@@ -40,7 +41,7 @@ async def parse_request(request_str: str) -> Optional[str]:
     try:
         request = json.loads(request_str)
     except Exception as e:
-        logger.exception(f"Error parsing request: {e}")
+        logger.warning(f"Error parsing request: {request_str}. {e}")
         return None
 
     messages = []
@@ -82,7 +83,7 @@ async def parse_output(output_str: str) -> Tuple[Optional[str], Optional[str]]:
     try:
         output = json.loads(output_str)
     except Exception as e:
-        logger.exception(f"Error parsing request: {e}")
+        logger.warning(f"Error parsing output: {output_str}. {e}")
         return None, None
 
     output_message = ""
@@ -107,9 +108,9 @@ async def parse_output(output_str: str) -> Tuple[Optional[str], Optional[str]]:
     return output_message, chat_id
 
 
-async def parse_get_prompt_with_output(
-    row: GetPromptWithOutputsRow,
-) -> Optional[PartialConversation]:
+async def _get_question_answer(
+    row: Union[GetPromptWithOutputsRow, GetAlertsWithPromptAndOutputRow]
+) -> Tuple[Optional[QuestionAnswer], Optional[str]]:
     """
     Parse a row from the get_prompt_with_outputs query and return a PartialConversation
 
@@ -124,7 +125,7 @@ async def parse_get_prompt_with_output(
 
     # If we couldn't parse the request or output, return None
     if not request_msg_str or not output_msg_str or not chat_id:
-        return None
+        return None, None
 
     request_message = ChatMessage(
         message=request_msg_str,
@@ -136,10 +137,20 @@ async def parse_get_prompt_with_output(
         timestamp=row.output_timestamp,
         message_id=row.output_id,
     )
-    question_answer = QuestionAnswer(
-        question=request_message,
-        answer=output_message,
-    )
+    return QuestionAnswer(question=request_message, answer=output_message), chat_id
+
+
+async def parse_get_prompt_with_output(
+    row: GetPromptWithOutputsRow,
+) -> Optional[PartialConversation]:
+    """
+    Parse a row from the get_prompt_with_outputs query and return a PartialConversation
+
+    The row contains the raw request and output strings from the pipeline.
+    """
+    question_answer, chat_id = await _get_question_answer(row)
+    if not question_answer or not chat_id:
+        return None
     return PartialConversation(
         question_answer=question_answer,
         provider=row.provider,
@@ -187,3 +198,34 @@ async def match_conversations(
         )
 
     return conversations
+
+
+async def parse_get_alert_conversation(
+    row: GetAlertsWithPromptAndOutputRow,
+) -> Optional[AlertConversation]:
+    """
+    Parse a row from the get_alerts_with_prompt_and_output query and return a Conversation
+
+    The row contains the raw request and output strings from the pipeline.
+    """
+    question_answer, chat_id = await _get_question_answer(row)
+    if not question_answer or not chat_id:
+        return None
+
+    conversation = Conversation(
+        question_answers=[question_answer],
+        provider=row.provider,
+        type=row.type,
+        chat_id=chat_id or "chat-id-not-found",
+        conversation_timestamp=row.timestamp,
+    )
+    code_snippet = json.loads(row.code_snippet) if row.code_snippet else None
+    return AlertConversation(
+        conversation=conversation,
+        alert_id=row.id,
+        code_snippet=code_snippet,
+        trigger_string=row.trigger_string,
+        trigger_type=row.trigger_type,
+        trigger_category=row.trigger_category,
+        timestamp=row.timestamp,
+    )

src/codegate/dashboard/request_models.py

@@ -3,6 +3,8 @@
 
 from pydantic import BaseModel
 
+from codegate.pipeline.base import CodeSnippet
+
 
 class ChatMessage(BaseModel):
     """
@@ -45,3 +47,17 @@ class Conversation(BaseModel):
     type: str
     chat_id: str
     conversation_timestamp: datetime.datetime
+
+
+class AlertConversation(BaseModel):
+    """
+    Represents an alert with it's respective conversation.
+    """
+
+    conversation: Conversation
+    alert_id: str
+    code_snippet: Optional[CodeSnippet]
+    trigger_string: Optional[str]
+    trigger_type: str
+    trigger_category: Optional[str]
+    timestamp: datetime.datetime