[#263] move setting headers to prepare_response

legion/legion/http.py

@@ -22,6 +22,7 @@
 import urllib
 
 import legion.config
+import legion.containers.headers
 import legion.utils
 
 import flask
@@ -133,15 +134,26 @@ def parse_request(input_request):
         raise ValueError('Unexpected http method: {}'.format(input_request.method))
 
 
-def prepare_response(response):
+def prepare_response(response_data, model_id=None, model_version=None, model_endpoint=None):
     """
     Produce an HTTP response from dict/list
 
-    :param response: dict/list with data
-    :type response: dict[str, any] or list[any]
+    :param response_data: dict/list with data
+    :type response_data: dict[str, any] or list[any]
+    :param model_id: model id
+    :type model_id: str
+    :param model_version: model version
+    :type model_version: str
+    :param model_endpoint: model endpoint
+    :type model_endpoint: str
     :return: bytes
     """
-    return flask.jsonify(response)
+    response = flask.jsonify(response_data)
+    response.headers[legion.containers.headers.MODEL_ID] = model_id
+    response.headers[legion.containers.headers.MODEL_VERSION] = model_version
+    response.headers[legion.containers.headers.MODEL_ENDPOINT] = model_endpoint
+
+    return response
 
 
 def provide_json_response(method):

legion/legion/serving/pyserve.py

@@ -22,7 +22,6 @@
 import itertools
 
 import legion.config
-import legion.containers.headers
 import legion.http
 import legion.model
 import legion.pymodel
@@ -121,11 +120,8 @@ def model_invoke(model_id, model_version, endpoint='default'):
         raise Exception('Unknown endpoint {!r}'.format(endpoint))
 
     output = model.endpoints[endpoint].invoke(input_dict)
-    response = legion.http.prepare_response(output)
-    response.headers[legion.containers.headers.MODEL_ID] = model_id
-    response.headers[legion.containers.headers.MODEL_VERSION] = model_version
-    response.headers[legion.containers.headers.MODEL_ENDPOINT] = endpoint
-
+    response = legion.http.prepare_response(output, model_id=model_id, model_version=model_version,
+                                            model_endpoint=endpoint)
     return response