From d5b55f7446bbeb821ce0013b7d82c9aa9732461b Mon Sep 17 00:00:00 2001
From: Alvaro Lopez Garcia <aloga@ifca.unican.es>
Date: Mon, 25 Oct 2021 13:04:47 +0200
Subject: [PATCH] Include example of marshalling of complex predictions

We can encapsulate complex predictions inside our JSON response, this
change produces a simple example that illustrates it.
---
 deepaas/api/v2/predict.py |  2 +-
 deepaas/model/v2/test.py  | 29 ++++++++++++++++++++---------
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/deepaas/api/v2/predict.py b/deepaas/api/v2/predict.py
index 1965bed7..8c43b9ec 100644
--- a/deepaas/api/v2/predict.py
+++ b/deepaas/api/v2/predict.py
@@ -78,7 +78,7 @@ async def post(self, request, wsk_args=None):
                 ret = open(ret.filename, 'rb')
 
             accept = args.get("accept", "application/json")
-            if accept != "application/json":
+            if accept not in ["application/json", "*/*"]:
                 response = web.Response(
                     body=ret,
                     content_type=accept,
diff --git a/deepaas/model/v2/test.py b/deepaas/model/v2/test.py
index b299748a..bc4cad70 100644
--- a/deepaas/model/v2/test.py
+++ b/deepaas/model/v2/test.py
@@ -14,6 +14,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import base64
 import time
 
 from oslo_log import log
@@ -45,6 +46,7 @@ class TestModel(base.BaseModel):
                 },
             )
         ),
+        "data": fields.Str(),
     }
 
     _logo = (
@@ -82,15 +84,22 @@ def warm(self):
 
     def predict(self, **kwargs):
         LOG.debug("Got the following kw arguments: %s", kwargs)
-        d = {
-            "date": "2019-01-1",
-            "labels": [{"label": "foo", "probability": 1.0}]
-        }
-        if kwargs.get("accept") == "text/plain":
-            return str(d)
-        elif kwargs.get("accept") == "image/png":
+        if kwargs.get("accept") == "image/png":
             return self._logo
-        return d
+
+        encoded = base64.b64encode(self._logo)
+        b64_str = str(encoded, 'utf-8')
+        b64_str = f"data:image/png;base64,{b64_str}"
+
+        if kwargs.get("accept") == "text/plain":
+            return b64_str
+        else:
+            d = {
+                "date": "2019-01-1",
+                "labels": [{"label": "foo", "probability": 1.0}],
+                "data": b64_str,
+            }
+            return d
 
     def train(self, *args, **kwargs):
         sleep = kwargs.get("sleep", 1)
@@ -120,7 +129,9 @@ def get_predict_args(self):
             "accept": fields.Str(
                 description=("Media type(s) that is/are acceptable for the "
                              "response."),
-                validate=validate.OneOf(["text/plain", "image/png"]),
+                validate=validate.OneOf(["application/json",
+                                         "text/plain",
+                                         "image/png"]),
                 location="headers",
             )
         }