Merge pull request #12 from gatewayd-io/improvements

Fix training bugs, improve API, retrain model v3, and expand test coverage

Author: mostafa

.github/workflows/test.yaml

@@ -12,26 +12,22 @@ jobs:
   test:
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
-          python-version: 3.12
+          python-version: "3.12"
       - name: Install dependencies
         run: |
           pip install poetry
           cd training && poetry install --with dev --no-root
-      - name: Run formatter, linter and type checker
-        run: |
-          cd training && poetry run ruff check .
-        #   mypy --explicit-package-bases .
-        #   flake8 .
-        #   interrogate -vv --ignore-init-module --exclude sigma_api .
-      - name: Run tests
+      - name: Lint training code
+        run: cd training && poetry run ruff check .
+      - name: Lint API code
+        run: cd training && poetry run ruff check ../api/
+      - name: Run training tests
         run: cd training && poetry run pytest --cov=training --cov-report term --cov-report lcov:coverage.lcov -vv
-    #   - name: Submit coverage report to Coveralls
-    #     if: ${{ success() }}
-    #     uses: coverallsapp/github-action@1.1.3
-    #     with:
-    #       github-token: ${{ secrets.GITHUB_TOKEN }}
-    #       path-to-lcov: ./coverage.lcov
+      - name: Install API dependencies
+        run: cd training && poetry run pip install flask gunicorn
+      - name: Run API tests
+        run: cd training && poetry run pytest ../api/test_api.py -vv

Dockerfile

@@ -1,21 +1,20 @@
-FROM tensorflow/tensorflow:latest
+FROM tensorflow/tensorflow:2.16.1
 
-ENV dataset=sqli_dataset2.csv
 ENV KMP_AFFINITY=noverbose
 ENV TF_CPP_MIN_LOG_LEVEL=3
-ENV DATASET_PATH=/app/${dataset}
+ENV VOCAB_PATH=/app/sql_tokenizer_vocab.json
+ENV MODEL_PATH=/app/sqli_model/3/
 ENV WORKERS=4
 ENV HOST=0.0.0.0
 ENV PORT=8000
 
 WORKDIR /app
-COPY api/api.py /app
-COPY api/pyproject.toml /app
-COPY api/poetry.lock /app
-COPY dataset/${dataset} /app
+COPY api/api.py /app/
+COPY api/pyproject.toml /app/
+COPY api/poetry.lock /app/
 COPY training/sql_tokenizer.py /app/
 COPY training/sql_tokenizer_vocab.json /app/
-COPY sqli_model/ /app/sqli_model/
+COPY sqli_model/3/ /app/sqli_model/3/
 RUN pip install --disable-pip-version-check poetry
 RUN poetry install --no-root
 

api/api.py

@@ -1,34 +1,45 @@
-from flask import Flask, jsonify, request
-import tensorflow as tf
-import pandas as pd
+import logging
 import os
-from sql_tokenizer import SQLTokenizer  # Import SQLTokenizer
+
+import tensorflow as tf
+from flask import Flask, jsonify, request
+
+from sql_tokenizer import SQLTokenizer
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+)
+logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
 
-# Constants and configurations
 MAX_WORDS = 10000
 MAX_LEN = 100
-DATASET_PATH = os.getenv("DATASET_PATH", "dataset/sqli_dataset1.csv")
+VOCAB_PATH = os.getenv("VOCAB_PATH", "sql_tokenizer_vocab.json")
 MODEL_PATH = os.getenv("MODEL_PATH", "/app/sqli_model/3/")
 
-# Load dataset and initialize SQLTokenizer
-DATASET = pd.read_csv(DATASET_PATH)
 sql_tokenizer = SQLTokenizer(max_words=MAX_WORDS, max_len=MAX_LEN)
-sql_tokenizer.fit_on_texts(DATASET["Query"])  # Fit tokenizer on dataset
+sql_tokenizer.load_token_index(VOCAB_PATH)
+logger.info("Loaded tokenizer vocabulary from %s (%d tokens)", VOCAB_PATH, len(sql_tokenizer.token_index))
 
-# Load the model using tf.saved_model.load and get the serving signature
 loaded_model = tf.saved_model.load(MODEL_PATH)
 model_predict = loaded_model.signatures["serving_default"]
+logger.info("Loaded model from %s", MODEL_PATH)
 
 
 def warm_up_model():
-    """Sends a dummy request to the model to 'warm it up'."""
+    """Sends a dummy request to the model to initialize it."""
     dummy_query = "SELECT * FROM users WHERE id = 1"
     query_seq = sql_tokenizer.texts_to_sequences([dummy_query])
     input_tensor = tf.convert_to_tensor(query_seq, dtype=tf.float32)
-    _ = model_predict(input_tensor)  # Make a dummy prediction to initialize the model
-    print("Model warmed up and ready to serve requests.")
+    _ = model_predict(input_tensor)
+    logger.info("Model warmed up and ready to serve requests.")
+
+
+@app.route("/health", methods=["GET"])
+def health():
+    return jsonify({"status": "ok"})
 
 
 @app.route("/predict", methods=["POST"])
@@ -37,27 +48,20 @@ def predict():
         return jsonify({"error": "No query provided"}), 400
 
     try:
-        # Tokenize and pad the input query using SQLTokenizer
         query = request.json["query"]
         query_seq = sql_tokenizer.texts_to_sequences([query])
         input_tensor = tf.convert_to_tensor(query_seq, dtype=tf.float32)
 
-        # Use the loaded model's serving signature to make the prediction
         prediction = model_predict(input_tensor)
 
-        # Check for valid output and extract the result
         if "output_0" not in prediction or prediction["output_0"].get_shape() != [1, 1]:
             return jsonify({"error": "Invalid model output"}), 500
 
-        # Extract confidence and return the response
-        return jsonify(
-            {
-                "confidence": float("%.4f" % prediction["output_0"].numpy()[0][0]),
-            }
-        )
-    except Exception as e:
-        # Log the error and return a proper error message
-        return jsonify({"error": str(e)}), 500
+        confidence = float("%.4f" % prediction["output_0"].numpy()[0][0])
+        return jsonify({"confidence": confidence})
+    except Exception:
+        logger.exception("Prediction failed")
+        return jsonify({"error": "Internal server error"}), 500
 
 
 if __name__ == "__main__":

api/pyproject.toml

@@ -6,12 +6,12 @@ authors = ["Mostafa Moradian <mostafa@gatewayd.io>"]
 readme = "README.md"
 
 [tool.poetry.dependencies]
-python = ">=3.10,<3.11"
+python = "^3.12"
 Flask = "^3.0.2"
 gunicorn = "^21.2.0"
-pandas = "^2.2.1"
+pandas = "^2.2.2"
 numpy = "^1.26.4"
-tensorflow = "^2.15.0"
+tensorflow = "^2.16.1"
 
 [build-system]
 requires = ["poetry-core"]

api/test_api.py

@@ -0,0 +1,74 @@
+import importlib
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "training"))
+
+os.environ.setdefault("VOCAB_PATH", os.path.join(
+    os.path.dirname(__file__), "..", "training", "sql_tokenizer_vocab.json"))
+os.environ.setdefault("MODEL_PATH", os.path.join(
+    os.path.dirname(__file__), "..", "sqli_model", "3"))
+
+spec = importlib.util.spec_from_file_location(
+    "api_module", os.path.join(os.path.dirname(__file__), "api.py"))
+api_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(api_module)
+app = api_module.app
+
+
+@pytest.fixture
+def client():
+    app.config["TESTING"] = True
+    with app.test_client() as client:
+        yield client
+
+
+def test_health(client):
+    resp = client.get("/health")
+    assert resp.status_code == 200
+    assert resp.get_json() == {"status": "ok"}
+
+
+def test_predict_missing_body(client):
+    resp = client.post("/predict", content_type="application/json")
+    assert resp.status_code == 400
+
+
+def test_predict_missing_query_key(client):
+    resp = client.post("/predict", json={"foo": "bar"})
+    assert resp.status_code == 400
+    data = resp.get_json()
+    assert "error" in data
+
+
+def test_predict_sqli(client):
+    resp = client.post("/predict", json={"query": "SELECT * FROM users WHERE id=1 OR 1=1"})
+    assert resp.status_code == 200
+    data = resp.get_json()
+    assert "confidence" in data
+    assert isinstance(data["confidence"], float)
+
+
+def test_predict_legitimate(client):
+    resp = client.post("/predict", json={"query": "SELECT name FROM products"})
+    assert resp.status_code == 200
+    data = resp.get_json()
+    assert "confidence" in data
+    assert isinstance(data["confidence"], float)
+
+
+def test_predict_empty_query(client):
+    resp = client.post("/predict", json={"query": ""})
+    assert resp.status_code == 200
+    data = resp.get_json()
+    assert "confidence" in data
+
+
+def test_predict_error_not_leaked(client):
+    """Ensure internal error details are not exposed to the client."""
+    resp = client.post("/predict", json={"query": ""})
+    if resp.status_code == 500:
+        data = resp.get_json()
+        assert data["error"] == "Internal server error"

sqli_model/3/fingerprint.pb

@@ -1 +1 @@
-ℨ���־�
��鿶���
���月��
 �����Ϗ�
(�������
2

+��������
���ݺ���Y���月��
 ��������
(���վ����
2
:'306335063828443668507412436166038701185