Merge pull request #284 from stacklok/mount-codegate-volume

Create a `codegate_volume` directory.

Author: aponcedeleonch

Dockerfile

@@ -107,14 +107,17 @@ ENV CODEGATE_APP_LOG_LEVEL=WARNING
 ENV CODEGATE_LOG_FORMAT=TEXT
 
 # Copy the initial models in the image to default models
-RUN mkdir -p /app/default_models && cp /app/models/* /app/default_models/
+RUN mkdir -p /app/default_models && cp /app/codegate_volume/models/* /app/default_models/
 
 # Define volume for persistent data
-VOLUME ["/app/models"]
+VOLUME ["/app/codegate_volume/"]
 
-# give the right permissions
+# This has to be performed after copying from the builder stages.
+# Otherwise, the permissions will be reset to root.
 USER root
-RUN chown -R codegate /app/models
+RUN mkdir -p /app/codegate_volume/db
+# Make codegate user the owner of codegate_volume directory to allow writing to it
+RUN chown -R codegate /app/codegate_volume
 USER codegate
 
 # Set the container's default entrypoint

README.md

@@ -51,7 +51,6 @@ Unlike E.T., your code never phones home! 🛸 Codegate is designed with privacy
 Make sure you have these tools installed:
 
 - 🐳 [Docker](https://docs.docker.com/get-docker/)
-- 🔧 [Docker Compose](https://docs.docker.com/compose/install/)
 - 🛠️ [jq](https://stedolan.github.io/jq/download/)
 - 💻 [VSCode](https://code.visualstudio.com/download)
 
@@ -155,10 +154,12 @@ docker run -p 8989:8989 -p 8990:80 codegate:latest
 
 # With pre-built pulled image
 docker pull ghcr.io/stacklok/codegate/codegate:latest
-docker run -p 8989:8989 -p 8990:80 ghcr.io/stacklok/codegate/codegate:latest
+docker run --name codegate -d -p 8989:8989 -p 8990:80 ghcr.io/stacklok/codegate/codegate:latest
 
-# With persistent models
-docker run -p 8989:8989 -p 8990:80 -v /path/to/volume:/app/models ghcr.io/stacklok/codegate/codegate:latest
+# It will mount a volume to /app/codegate_volume
+# The directory supports storing Llama CPP models under subidrectoy /models
+# A sqlite DB with the messages and alerts is stored under the subdirectory /db
+docker run --name codegate -d -v /path/to/volume:/app/codegate_volume -p 8989:8989 -p 8990:80 ghcr.io/stacklok/codegate/codegate:latest
 ```
 
 ### Exposed parameters

models/all-minilm-L6-v2-q5_k_m.gguf renamed to codegate_volume/models/all-minilm-L6-v2-q5_k_m.gguf

@@ -19,7 +19,7 @@ log_level: "INFO"  # One of: ERROR, WARNING, INFO, DEBUG
 ##
 
 # Model to use for chatting
-model_base_path: "./models"
+model_base_path: "./codegate_volume/models"
 
 # Context length of the model
 chat_model_n_ctx: 32768

config.yaml

@@ -33,7 +33,7 @@ provider_urls:
 ##
 
 # Model to use for chatting
-chat_model_path: "./models/qwen2.5-coder-1.5b-instruct-q5_k_m.gguf"
+chat_model_path: "./codegate_volume/models/qwen2.5-coder-1.5b-instruct-q5_k_m.gguf"
 
 # Context length of the model
 chat_model_n_ctx: 32768

config.yaml.example

@@ -5,7 +5,8 @@ CODEGATE_VLLM_URL=${CODEGATE_VLLM_URL:-$DEFAULT_CODEGATE_VLLM_URL}
 # those are hardcoded on the image, will not change
 BACKUP_PATH="/tmp/weaviate_backup"
 BACKUP_NAME="backup"
-MODEL_BASE_PATH="/app/models"
+MODEL_BASE_PATH="/app/codegate_volume/models"
+CODEGATE_DB_FILE="/app/codegate_volume/db/codegate.db"
 
 # Function to restore backup if paths are provided
 restore_backup() {
@@ -30,8 +31,8 @@ start_dashboard() {
 # Function to start the main application
 start_application() {
     # first restore the models
-    cp /app/default_models/* /app/models/
-    CMD_ARGS="--port 8989 --host 0.0.0.0 --vllm-url $CODEGATE_VLLM_URL --model-base-path $MODEL_BASE_PATH"
+    cp /app/default_models/* /app/codegate_volume/models
+    CMD_ARGS="--port 8989 --host 0.0.0.0 --vllm-url $CODEGATE_VLLM_URL --model-base-path $MODEL_BASE_PATH --db-path $CODEGATE_DB_FILE"
 
     # Check and append additional URLs if they are set
     [ -n "$CODEGATE_OPENAI_URL" ] && CMD_ARGS+=" --openai-url $CODEGATE_OPENAI_URL"

scripts/entrypoint.sh

@@ -35,7 +35,7 @@ def __init__(self, jsonl_dir='data', take_backup=True, restore_backup=True):
         ]
         self.client.connect()
         self.inference_engine = LlamaCppInferenceEngine()
-        self.model_path = "./models/all-minilm-L6-v2-q5_k_m.gguf"
+        self.model_path = "./codegate_volume/models/all-minilm-L6-v2-q5_k_m.gguf"
 
     def restore_backup(self):
         if os.getenv("BACKUP_FOLDER"):

scripts/import_packages.py

@@ -118,7 +118,7 @@ def show_prompts(prompts: Optional[Path]) -> None:
 @click.option(
     "--model-base-path",
     type=str,
-    default="./models",
+    default="./codegate_volume/models",
     help="Path to the model base directory",
 )
 @click.option(
@@ -127,6 +127,12 @@ def show_prompts(prompts: Optional[Path]) -> None:
     default="all-minilm-L6-v2-q5_k_m.gguf",
     help="Name of the model to use for embeddings",
 )
+@click.option(
+    "--db-path",
+    type=str,
+    default=None,
+    help="Path to the SQLite database file",
+)
 def serve(
     port: Optional[int],
     host: Optional[str],
@@ -140,6 +146,7 @@ def serve(
     ollama_url: Optional[str],
     model_base_path: Optional[str],
     embedding_model: Optional[str],
+    db_path: Optional[str],
 ) -> None:
     """Start the codegate server."""
     logger = None
@@ -166,6 +173,7 @@ def serve(
             cli_provider_urls=cli_provider_urls,
             model_base_path=model_base_path,
             embedding_model=embedding_model,
+            db_path=db_path,
         )
 
         setup_logging(cfg.log_level, cfg.log_format)
@@ -181,10 +189,11 @@ def serve(
                 "provider_urls": cfg.provider_urls,
                 "model_base_path": cfg.model_base_path,
                 "embedding_model": cfg.embedding_model,
+                "db_path": cfg.db_path,
             },
         )
 
-        init_db_sync()
+        init_db_sync(cfg.db_path)
         app = init_app()
 
         import uvicorn

src/codegate/cli.py

@@ -37,10 +37,11 @@ class Config:
     log_format: LogFormat = LogFormat.JSON
     prompts: PromptConfig = field(default_factory=PromptConfig)
 
-    model_base_path: str = "./models"
+    model_base_path: str = "./codegate_volume/models"
     chat_model_n_ctx: int = 32768
     chat_model_n_gpu_layers: int = -1
     embedding_model: str = "all-minilm-L6-v2-q5_k_m.gguf"
+    db_path: Optional[str] = None
 
     # Provider URLs with defaults
     provider_urls: Dict[str, str] = field(default_factory=lambda: DEFAULT_PROVIDER_URLS.copy())
@@ -178,6 +179,7 @@ def load(
         cli_provider_urls: Optional[Dict[str, str]] = None,
         model_base_path: Optional[str] = None,
         embedding_model: Optional[str] = None,
+        db_path: Optional[str] = None,
     ) -> "Config":
         """Load configuration with priority resolution.
 
@@ -197,6 +199,7 @@ def load(
             cli_provider_urls: Optional dict of provider URLs from CLI
             model_base_path: Optional path to model base directory
             embedding_model: Optional name of the model to use for embeddings
+            db_path: Optional path to the SQLite database file
 
         Returns:
             Config: Resolved configuration
@@ -253,6 +256,8 @@ def load(
             config.model_base_path = model_base_path
         if embedding_model is not None:
             config.embedding_model = embedding_model
+        if db_path is not None:
+            config.db_path = db_path
 
         # Set the __config class attribute
         Config.__config = config

src/codegate/config.py

@@ -9,7 +9,7 @@
 import structlog
 from litellm import ChatCompletionRequest, ModelResponse
 from pydantic import BaseModel
-from sqlalchemy import create_engine, text
+from sqlalchemy import text
 from sqlalchemy.ext.asyncio import create_async_engine
 
 from codegate.db.models import Alert, Output, Prompt
@@ -29,21 +29,18 @@ def __init__(self, sqlite_path: Optional[str] = None):
         # Initialize SQLite database engine with proper async URL
         if not sqlite_path:
             current_dir = Path(__file__).parent
-            self._db_path = (current_dir.parent.parent.parent / "codegate.db").absolute()
-        else:
-            self._db_path = Path(sqlite_path).absolute()
-
-        # Initialize SQLite database engine with proper async URL
-        current_dir = Path(__file__).parent
-        self._db_path = (current_dir.parent.parent.parent / "codegate.db").absolute()
+            sqlite_path = (
+                current_dir.parent.parent.parent / "codegate_volume" / "db" / "codegate.db"
+            )
+        self._db_path = Path(sqlite_path).absolute()
+        self._db_path.parent.mkdir(parents=True, exist_ok=True)
         logger.debug(f"Initializing DB from path: {self._db_path}")
         engine_dict = {
             "url": f"sqlite+aiosqlite:///{self._db_path}",
             "echo": False,  # Set to False in production
             "isolation_level": "AUTOCOMMIT",  # Required for SQLite
         }
         self._async_db_engine = create_async_engine(**engine_dict)
-        self._db_engine = create_engine(**engine_dict)
 
     def does_db_exist(self):
         return self._db_path.is_file()
@@ -242,13 +239,9 @@ async def get_alerts_with_prompt_and_output(self) -> List[GetAlertsWithPromptAnd
         return prompts
 
 
-def init_db_sync():
+def init_db_sync(db_path: Optional[str] = None):
     """DB will be initialized in the constructor in case it doesn't exist."""
-    db = DbRecorder()
-    # Remove the DB file if exists for the moment to not cause issues at schema change.
-    # We can replace this in the future with migrations or something similar.
-    if db.does_db_exist():
-        db._db_path.unlink()
+    db = DbRecorder(db_path)
     asyncio.run(db.init_db())