enforcement of the schema and clear error if schema is violated

Author: andrewheusser

app.py

@@ -1,8 +1,10 @@
 import json
+from jsonschema import validate, ValidationError
 from flask import Flask, request, jsonify
 from celery import Celery, Task
 from kombu import serialization
 
+
 import pandas as pd
 import arviz as az
 from pymc_marketing.mmm import (
@@ -102,6 +104,16 @@ def __call__(self, *args: object, **kwargs: object) -> object:
 # Ensure proper permissions (readable/writable by all users)
 os.chmod(DATA_DIR, 0o777)
 
+# Extract the request schema from the OpenAPI spec
+def get_mmm_request_schema():
+    try:
+        with open('gpt-agent/api_spec.json', 'r') as f:
+            api_spec = json.load(f)
+            return api_spec['paths']['/run_mmm_async']['post']['requestBody']['content']['application/json']['schema']
+    except Exception as e:
+        logging.error("Failed to load API spec: %s", str(e))
+        raise e
+
 
 @celery.task(bind=True)
 def run_mmm_task(self, data):
@@ -233,12 +245,19 @@ def run_mmm_async():
         data = request.get_json()
         logging.debug("run_mmm_async request data: %s", data)
 
-        logging.info("checking that the data has file_refs: %s", data)
-        if ("openaiFileIdRefs" not in data) or (len(data["openaiFileIdRefs"]) == 0): # TODO: do a more thorough schema check here
-            logging.error("Data does not have openaiFileIdRefs")
-            return jsonify({"error": "Request must include openaiFileIdRefs"}), 400
-        else:
-            logging.info("Data has openaiFileIdRefs")
+        try:
+            schema = get_mmm_request_schema()
+            validate(instance=data, schema=schema)
+        except ValidationError as e:
+            logging.error("Schema validation failed: %s", str(e))
+            return jsonify({
+                "error": "Invalid request format",
+                "details": {
+                    "message": str(e),
+                    "path": " -> ".join(str(p) for p in e.path),
+                    "schema_path": " -> ".join(str(p) for p in e.schema_path)
+                }
+            }), 400
 
         task = run_mmm_task.apply_async(args=[data])
         logging.info("Task submitted with ID: %s", task.id)

environment.yml

@@ -14,3 +14,4 @@ dependencies:
   - procps-ng
   - yq=2.12.0
   - dill=0.3.9
+  - jsonschema=4.23.0

gpt-agent/api_spec.json

@@ -22,7 +22,7 @@
             "application/json": {
               "schema": {
                 "type": "object",
-                "required": ["openaiFileIdRefs"],
+                "required": ["openaiFileIdRefs", "date_column", "channel_columns", "y_column"],
                 "properties": {
                   "openaiFileIdRefs": {
                     "type": "array",
@@ -59,7 +59,6 @@
                   },
                   "date_column": {
                     "type": "string",
-                    "default": "date",
                     "description": "Name of the date column in data."
                   },
                   "channel_columns": {
@@ -69,6 +68,10 @@
                     },
                     "description": "List of channel column names."
                   },
+                  "y_column": {
+                    "type": "string",
+                    "description": "Name of the y column in data."
+                  },
                   "adstock_max_lag": {
                     "type": "integer",
                     "default": 8,

test_mmm_async.py

@@ -50,8 +50,7 @@ def test_missing_file_refs(base_url):
     }
     response = requests.post(run_url, data=json.dumps(payload), headers=headers)
     assert response.status_code == 400
-    assert response.json()["error"] == "Request must include openaiFileIdRefs"
-
+    assert response.json()["error"] == "Invalid request format"
 
 def test_async_mmm_run(base_url):
     # Payload that includes data