Refactor PSQL incident store into ORM (SQLAlchemy)

Makefile

@@ -10,7 +10,7 @@ HUB_URL ?= ""
 IMPORTER_ARGS ?= ""
 
 run-postgres:
-	$(CONTAINER_RUNTIME) run -it -v data:/var/lib/postgresql/data -e POSTGRES_USER=kai -e POSTGRES_PASSWORD=dog8code -e POSTGRES_DB=kai -p 5432:5432 docker.io/pgvector/pgvector:pg15
+	$(CONTAINER_RUNTIME) run -it -v data:/var/lib/postgresql/data -e POSTGRES_USER=kai -e POSTGRES_PASSWORD=dog8code -e POSTGRES_DB=kai -p 5432:5432 docker.io/library/postgres:16.3
 
 run-server:
 	PYTHONPATH=$(KAI_PYTHON_PATH) LOGLEVEL=$(LOGLEVEL) DEMO_MODE=$(DEMO_MODE) gunicorn --timeout 3600 -w $(NUM_WORKERS) --bind localhost:8080 --worker-class aiohttp.GunicornWebWorker 'kai.server:app()'

kai/config.toml

@@ -2,6 +2,24 @@ log_level = "info"
 demo_mode = false
 trace_enabled = true
 
+# **Postgresql incident store**
+# [incident_store]
+# provider = "postgresql"
+
+# [incident_store.args]
+# host = "127.0.0.1"
+# database = "kai"
+# user = "kai"
+# password = "dog8code"
+
+# **In-memory sqlite incident store**
+# ```
+# [incident_store]
+# provider = "sqlite"
+#
+# [incident_store.args]
+# connection_string = "sqlite:///:memory:"
+
 [incident_store]
 provider = "postgresql"
 

kai/evaluation.py

@@ -12,10 +12,10 @@
 from kai.model_provider import ModelProvider
 from kai.models.analyzer_types import Incident
 from kai.models.file_solution import guess_language, parse_file_solution_content
-from kai.models.kai_config import KaiConfig
+from kai.models.kai_config import KaiConfig, KaiConfigIncidentStoreSQLiteArgs
 from kai.report import Report
-from kai.service.incident_store.in_memory import InMemoryIncidentStore
 from kai.service.incident_store.incident_store import Application
+from kai.service.incident_store.sqlite import SQLiteIncidentStore
 
 """
 The point of this file is to automatically see if certain prompts make the
@@ -172,7 +172,12 @@ def evaluate(
 
                     created_git_repo = True
 
-                incident_store = InMemoryIncidentStore(None)
+                incident_store = SQLiteIncidentStore(
+                    KaiConfigIncidentStoreSQLiteArgs(
+                        connection_string="sqlite:///:memory:"
+                    ),
+                    None,
+                )
                 incident_store.load_report(example.application, example.report)
 
                 pb_incidents = []

kai/llm_io_handler.py

@@ -82,7 +82,7 @@ def get_prompt(
         if not fallback:
             raise e
 
-        KAI_LOG.error(f"Template '{e.name}' not found. Falling back to main.jinja")
+        KAI_LOG.warning(f"Template '{e.name}' not found. Falling back to main.jinja")
         template = jinja_env.get_template("main.jinja")
 
     KAI_LOG.debug(f"Template {template.filename} loaded")

kai/models/kai_config.py

@@ -4,40 +4,77 @@
 from typing import Literal, Optional, Union
 
 import yaml
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, validator
 
 # Incident store providers
 
 
 class KaiConfigIncidentStoreProvider(Enum):
     POSTGRESQL = "postgresql"
-    IN_MEMORY = "in_memory"
+    SQLITE = "sqlite"
 
 
 class KaiConfigIncidentStorePostgreSQLArgs(BaseModel):
-    host: str
-    database: str
-    user: str
-    password: str
+    host: Optional[str] = None
+    database: Optional[str] = None
+    user: Optional[str] = None
+    password: Optional[str] = None
+
+    connection_string: Optional[str] = None
+
+    @validator("connection_string", always=True)
+    def validate_connection_string(cls, v, values):
+        connection_string_present = v is not None
+        parameters_present = all(
+            values.get(key) is not None
+            for key in ["host", "database", "user", "password"]
+        )
+
+        if connection_string_present == parameters_present:
+            raise ValueError(
+                "Must provide one of connection_string or [host, database, user, password]"
+            )
+
+        return v
 
 
 class KaiConfigIncidentStorePostgreSQL(BaseModel):
     provider: Literal["postgresql"]
     args: KaiConfigIncidentStorePostgreSQLArgs
 
 
-class KaiConfigIncidentStoreInMemoryArgs(BaseModel):
-    dummy: bool
+class KaiConfigIncidentStoreSQLiteArgs(BaseModel):
+    host: Optional[str] = None
+    database: Optional[str] = None
+    user: Optional[str] = None
+    password: Optional[str] = None
+
+    connection_string: Optional[str] = None
+
+    @validator("connection_string", always=True)
+    def validate_connection_string(cls, v, values):
+        connection_string_present = v is not None
+        parameters_present = all(
+            values.get(key) is not None
+            for key in ["host", "database", "user", "password"]
+        )
+
+        if connection_string_present == parameters_present:
+            raise ValueError(
+                "Must provide one of connection_string or [host, database, user, password]"
+            )
+
+        return v
 
 
-class KaiConfigIncidentStoreInMemory(BaseModel):
-    provider: Literal["in_memory"]
-    args: KaiConfigIncidentStoreInMemoryArgs
+class KaiConfigIncidentStoreSQLIte(BaseModel):
+    provider: Literal["sqlite"]
+    args: KaiConfigIncidentStoreSQLiteArgs
 
 
 KaiConfigIncidentStore = Union[
     KaiConfigIncidentStorePostgreSQL,
-    KaiConfigIncidentStoreInMemory,
+    KaiConfigIncidentStoreSQLIte,
 ]
 
 # Model providers

kai/report.py

@@ -1,5 +1,7 @@
 __all__ = ["Report"]
 
+import hashlib
+import json
 import os
 import shutil
 from io import StringIO
@@ -12,9 +14,10 @@
 class Report:
     report: dict = None
 
-    def __init__(self, report_data: dict):
+    def __init__(self, report_data: dict, report_id: str):
         self.workaround_counter_for_missing_ruleset_name = 0
         self.report = self._reformat_report(report_data)
+        self.report_id = report_id
 
     def __str__(self):
         return str(self.report)
@@ -35,19 +38,25 @@ def __len__(self):
         return len(self.report)
 
     @classmethod
-    def load_report_from_object(cls, report_data: dict):
+    def load_report_from_object(cls, report_data: dict, report_id: str):
         """
         Class method to create a Report instance directly from a Python dictionary object.
         """
-        return cls(report_data=report_data)
+        return cls(report_data=report_data, report_id=report_id)
 
     @classmethod
     def load_report_from_file(cls, file_name: str):
         KAI_LOG.info(f"Reading report from {file_name}")
         with open(file_name, "r") as f:
             report: dict = yaml.safe_load(f)
         report_data = report
-        return cls(report_data)
+        # report_id is the hash of the json.dumps of the report_data
+        return cls(
+            report_data,
+            hashlib.sha256(
+                json.dumps(report_data, sort_keys=True).encode()
+            ).hexdigest(),
+        )
 
     @staticmethod
     def get_cleaned_file_path(uri: str):

kai/server.py

@@ -295,13 +295,15 @@ def app():
     if config.demo_mode:
         KAI_LOG.info("DEMO_MODE is enabled. LLM responses will be cached")
 
-    webapp["incident_store"] = IncidentStore.from_config(config.incident_store)
-    KAI_LOG.info(f"Selected incident store: {config.incident_store.provider}")
-
     webapp["model_provider"] = ModelProvider(config.models)
     KAI_LOG.info(f"Selected provider: {config.models.provider}")
     KAI_LOG.info(f"Selected model: {webapp['model_provider'].model_id}")
 
+    webapp["incident_store"] = IncidentStore.from_config(
+        config.incident_store, webapp["model_provider"]
+    )
+    KAI_LOG.info(f"Selected incident store: {config.incident_store.provider}")
+
     webapp.add_routes(routes)
 
     return webapp

kai/service/incident_store/__init__.py

@@ -1,11 +1,11 @@
-from .in_memory import InMemoryIncidentStore
 from .incident_store import Application, IncidentStore, Solution
 from .psql import PSQLIncidentStore
+from .sqlite import SQLiteIncidentStore
 
 __all__ = [
     "IncidentStore",
     "Solution",
     "PSQLIncidentStore",
-    "InMemoryIncidentStore",
+    "SQLiteIncidentStore",
     "Application",
 ]