Repository: caikit/caikit
Issue: REST server cannot handle null in model response

## Describe the bug

If I have a module which does not set a field in its response data model, and that field is itself a type which accepts None as a default (e.g. another data model type), then response serialization fails in the rest server.

## Platform

- Interpreter version: 3.11
- Library version: 0.11.0

## Sample Code

The simplest thing to do is not set a `ProducerId` where a DM output has one specified

## Expected behavior

Either a `{producer_id: null}` or a `{producer_id: {name: "", version: ""}}` would probably be acceptable.

## Observed behavior

We get a cryptic error from uvicorn with not much help:

> {"channel": "uvicorn.error", "exception": "Traceback (most recent call last):
\\n File "/usr/local/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py", line 426, in run_asgi
\\n result = await app( # type: ignore[func-returns-value]
\\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
\\n File "/usr/local/lib/python3.11/site-packages/uvicorn/middleware/proxy_headers.py", line 84, in call
\\n return await self.app(scope, receive, send)
\\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
\\n File "/usr/local/lib/python3.11/site-packages/fastapi/applications.py", line 289, in call
\\n await super().call(scope, receive, send)
\\n File "/usr/local/lib/python3.11/site-packages/starlette/applications.py", line 122, in call
\\n await self.middleware_stack(scope, receive, send)
\\n File "/usr/local/lib/python3.11/site-packages/starlette/middleware/errors.py", line 184, in call
\\n raise exc
\\n File "/usr/local/lib/python3.11/site-packages/starlette/middleware/errors.py", line 162, in call
\\n await self.app(scope, receive, _send)
\\n File "/usr/local/lib/python3.11/site-packages/starlette/middleware/exceptions.py", line 79, in call
\\n raise exc
\\n File "/usr/local/lib/python3.11/site-packages/starlette/middleware/exceptions.py", line 68, in call
\\n await self.app(scope, receive, sender)
\\n File "/usr/local/lib/python3.11/site-packages/fastapi/middleware/asyncexitstack.py", line 20, in call
\\n raise e
\\n File "/usr/local/lib/python3.11/site-packages/fastapi/middleware/asyncexitstack.py", line 17, in call
\\n await self.app(scope, receive, send)
\\n File "/usr/local/lib/python3.11/site-packages/starlette/routing.py", line 718, in call
\\n await route.handle(scope, receive, send)
\\n File "/usr/local/lib/python3.11/site-packages/starlette/routing.py", line 276, in handle
\\n await self.app(scope, receive, send)
\\n File "/usr/local/lib/python3.11/site-packages/starlette/routing.py", line 66, in app
\\n response = await func(request)
\\n ^^^^^^^^^^^^^^^^^^^
\\n File "/usr/local/lib/python3.11/site-packages/fastapi/routing.py", line 291, in app
\\n content = await serialize_response(
    \\n ^^^^^^^^^^^^^^^^^^^^^^^^^
\\n File "/usr/local/lib/python3.11/site-packages/fastapi/routing.py", line 154, in serialize_response
\\n raise ResponseValidationError(
    \\nfastapi.exceptions.ResponseValidationError", "level": "error", "message": "Exception in ASGI application
\\n", "num_indent": 0, "thread_id": 281473881595936, "timestamp": "2023-07-20T18:16:14.088814"}
{"channel": "uvicorn.access", "exception": null, "level": "info", "message": "172.17.0.1:47910 - "POST /api/v1/foo/task/text-generation HTTP/1.1" 500", "num_indent": 0, "thread_id": 281473881595936, "timestamp": "2023-07-20T18:16:14.278682"}

## Additional context

Add any other context about the problem here.


Changed files:
tests/fixtures/sample_lib/modules/other_task/other_implementation.py, caikit/runtime/http_server.py, tests/runtime/test_http_server.py, tests/runtime/test_grpc_server.py, tests/runtime/servicers/test_global_train_servicer_impl.py, tests/fixtures/sample_lib/data_model/sample.py

diff --git a/caikit/runtime/http_server.py b/caikit/runtime/http_server.py
index 087661a..f744650 100644
--- a/caikit/runtime/http_server.py
+++ b/caikit/runtime/http_server.py
@@ -28,8 +28,10 @@ import threading
 import time
 
 # Third Party
-from fastapi import FastAPI, Request, Response
-from fastapi.responses import PlainTextResponse
+from fastapi import FastAPI, Request, Response, status
+from fastapi.encoders import jsonable_encoder
+from fastapi.exceptions import ResponseValidationError
+from fastapi.responses import JSONResponse, PlainTextResponse
 from grpc import StatusCode
 from sse_starlette import EventSourceResponse, ServerSentEvent
 import numpy as np
@@ -117,6 +119,14 @@ class RuntimeHTTPServer(RuntimeServerBase):
 
         self.app = FastAPI()
 
+        # Response validation
+        @self.app.exception_handler(ResponseValidationError)
+        async def validation_exception_handler(_, exc: ResponseValidationError):
+            return JSONResponse(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                content=jsonable_encoder({"detail": exc.errors(), "body": exc.body}),
+            )
+
         # Start metrics server
         RuntimeServerBase._start_metrics_server()
 
diff --git a/tests/fixtures/sample_lib/data_model/sample.py b/tests/fixtures/sample_lib/data_model/sample.py
index f31e098..f3200ee 100644
--- a/tests/fixtures/sample_lib/data_model/sample.py
+++ b/tests/fixtures/sample_lib/data_model/sample.py
@@ -7,6 +7,7 @@ import typing
 
 # Local
 from caikit.core import DataObjectBase, TaskBase, dataobject, task
+from caikit.core.data_model import ProducerId
 
 
 @dataobject(package="caikit_data_model.sample_lib")
@@ -29,6 +30,7 @@ class OtherOutputType(DataObjectBase):
     """A simple return type for the `other_task` task"""
 
     farewell: str
+    producer_id: ProducerId
 
 
 @dataobject(package="caikit_data_model.sample_lib")
diff --git a/tests/fixtures/sample_lib/modules/other_task/other_implementation.py b/tests/fixtures/sample_lib/modules/other_task/other_implementation.py
index f765f21..93fca08 100644
--- a/tests/fixtures/sample_lib/modules/other_task/other_implementation.py
+++ b/tests/fixtures/sample_lib/modules/other_task/other_implementation.py
@@ -7,6 +7,7 @@ from typing import Union
 # Local
 from ...data_model.sample import OtherOutputType, OtherTask, SampleInputType
 from caikit.core.data_model import DataStream
+from caikit.core.data_model.producer import ProducerId
 from caikit.core.modules import ModuleLoader, ModuleSaver
 import caikit.core
 
@@ -21,9 +22,11 @@ class OtherModule(caikit.core.ModuleBase):
         self.learning_rate = learning_rate
 
     def run(
-        self, sample_input: Union[SampleInputType, str]
+        self, sample_input: Union[SampleInputType, str], include_producer_id=True
     ) -> Union[OtherOutputType, str]:
-        return OtherOutputType(f"goodbye: {sample_input.name} {self.batch_size} times")
+        farewell = f"goodbye: {sample_input.name} {self.batch_size} times"
+        producer_id = ProducerId("other_mod", "1.1.1") if include_producer_id else None
+        return OtherOutputType(farewell, producer_id)
 
     @classmethod
     def load(cls, model_path, **kwargs):
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 0c4afdb..6f2756e 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -27,6 +27,7 @@ import pytest
 # Local
 from caikit.config import get_config
 from caikit.core import MODEL_MANAGER
+from caikit.core.data_model.producer import ProducerId
 from caikit.interfaces.common.data_model.stream_sources import S3Path
 from caikit.runtime.servicers.global_train_servicer import GlobalTrainServicer
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
@@ -194,6 +195,7 @@ def test_global_train_other_task(
         inference_response
         == OtherOutputType(
             farewell=f"goodbye: Gabe {batch_size} times",
+            producer_id=ProducerId("other_mod", "1.1.1"),
         ).to_proto()
     )
 
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index acc205b..e83ab15 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -43,6 +43,7 @@ import alog
 from caikit import get_config
 from caikit.core import MODEL_MANAGER
 from caikit.core.data_model.base import DataBase
+from caikit.core.data_model.producer import ProducerId
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
     TrainingJob,
@@ -471,7 +472,7 @@ def test_train_fake_module_does_not_change_another_instance_model_of_block(
         predict_request, metadata=[("mm-model-id", actual_response.model_name)]
     )
     expected_trained_inference_response = OtherOutputType(
-        farewell="goodbye: Gabe 100 times"
+        farewell="goodbye: Gabe 100 times", producer_id=ProducerId("other_mod", "1.1.1")
     ).to_proto()
     assert trained_inference_response == expected_trained_inference_response
 
@@ -480,7 +481,7 @@ def test_train_fake_module_does_not_change_another_instance_model_of_block(
         predict_request, metadata=[("mm-model-id", other_task_model_id)]
     )
     expected_original_inference_response = OtherOutputType(
-        farewell="goodbye: Gabe 42 times"
+        farewell="goodbye: Gabe 42 times", producer_id=ProducerId("other_mod", "1.1.1")
     ).to_proto()
     assert original_inference_response == expected_original_inference_response
 
diff --git a/tests/runtime/test_http_server.py b/tests/runtime/test_http_server.py
index 797dd3d..5839e5c 100644
--- a/tests/runtime/test_http_server.py
+++ b/tests/runtime/test_http_server.py
@@ -301,6 +301,69 @@ def test_inference_other_task(other_task_model_id, runtime_http_server):
         assert json_response["farewell"] == "goodbye: world 42 times"
 
 
+def test_inference_other_task_request_throws_missing_field(
+    other_task_model_id, runtime_http_server
+):
+    """Test to check 400 in case of missing request fields"""
+    with TestClient(runtime_http_server.app) as client:
+        json_input = {}
+        response = client.post(
+            f"/api/v1/{other_task_model_id}/task/other",
+            json=json_input,
+        )
+        assert response.status_code == 400
+        json_response = json.loads(response.content.decode(response.default_encoding))
+        assert (
+            "missing 1 required positional argument: 'sample_input'"
+            in json_response["details"]
+        )
+
+
+def test_inference_other_task_response_throws_missing_producer_id(
+    other_task_model_id, runtime_http_server
+):
+    """Test to check we throw a 500 if response has fields that are missing"""
+    with TestClient(runtime_http_server.app) as client:
+        json_input = {
+            "inputs": {"name": "world"},
+            # this intentionally skips sending back a producer-id in our module
+            "parameters": {"include_producer_id": False},
+        }
+        response = client.post(
+            f"/api/v1/{other_task_model_id}/task/other",
+            json=json_input,
+        )
+        assert response.status_code == 500
+        json_response = json.loads(response.content.decode(response.default_encoding))
+        assert json_response["detail"][0]["loc"] == ["response", "producer_id"]
+
+
+def test_inference_other_task_request_throws_incorrect_field(
+    other_task_model_id, runtime_http_server
+):
+    """Test to check we throw a 422 if request has an incorrect field type"""
+    with TestClient(runtime_http_server.app) as client:
+        json_input = {
+            "inputs": {"name": "world"},
+            "parameters": {"include_producer_id": "wrong_type"},
+        }
+        response = client.post(
+            f"/api/v1/{other_task_model_id}/task/other",
+            json=json_input,
+        )
+        json_response = json.loads(response.content.decode(response.default_encoding))
+        assert json_response["detail"][0]["loc"] == [
+            "body",
+            "parameters",
+            "include_producer_id",
+        ]
+        assert (
+            json_response["detail"][0]["msg"]
+            == "Input should be a valid boolean, unable to interpret input"
+        )
+        assert response.status_code == 422
+
+
 def test_inference_streaming_sample_module(sample_task_model_id, runtime_http_server):
     """Simple check for testing a happy path unary-stream case"""
     with TestClient(runtime_http_server.app) as client: