always define components.schemas.Output

..even when the user-defined predict function returns a simple type like a string.

Co-Authored-By: Ben Firshman <[email protected]>
Signed-off-by: Zeke Sikelianos <[email protected]>

Author: zeke

python/cog/predictor.py

@@ -6,7 +6,7 @@
 import os.path
 from pathlib import Path
 import typing
-from pydantic import create_model
+from pydantic import create_model, BaseModel
 from pydantic.fields import FieldInfo
 
 # Added in Python 3.8. Can be from typing if we drop support for <3.8.
@@ -117,4 +117,8 @@ def get_output_type(predictor: Predictor):
     if get_origin(OutputType) is Generator:
         OutputType = get_args(OutputType)[0]
 
-    return OutputType
+    # Wrap the type in a model so Pydantic can document it in component schema
+    class Output(BaseModel):
+        __root__: OutputType
+
+    return Output

python/cog/server/http.py

@@ -37,6 +37,18 @@ class Request(BaseModel):
         input: InputType = None
         output_file_prefix: str = None
 
+    # response_model is purely for generating schema.
+    # We generate Response again in the request so we can set file output paths correctly, etc.
+    OutputType = get_output_type(predictor)
+
+    @app.post(
+        "/predictions",
+        response_model=get_response_type(OutputType),
+        response_model_exclude_unset=True,
+    )
+
+    # The signature of this function is used by FastAPI to generate the schema.
+    # The function body is not used to generate the schema.
     def predict(request: Request = Body(default=None)):
         if request is None or request.input is None:
             output = predictor.predict()
@@ -77,15 +89,6 @@ def predict(...) -> output_type:
         )
         return JSONResponse(content=encoded_response)
 
-    # response_model is purely for generating schema.
-    # We generate Response again in the request so we can set file output paths correctly, etc.
-    OutputType = get_output_type(predictor)
-    app.post(
-        "/predictions",
-        response_model=get_response_type(OutputType),
-        response_model_exclude_unset=True,
-    )(predict)
-
     return app
 
 

python/tests/server/test_http.py

@@ -6,6 +6,7 @@
 from unittest import mock
 
 from fastapi.testclient import TestClient
+from pydantic import BaseModel
 from PIL import Image
 import pytest
 
@@ -152,6 +153,10 @@ def predict(
                         "choices": {"$ref": "#/components/schemas/choices"},
                     },
                 },
+                "Output": {
+                    "title": "Output",
+                    "type": "string",
+                },
                 "Request": {
                     "title": "Request",
                     "type": "object",
@@ -169,7 +174,7 @@ def predict(
                     "type": "object",
                     "properties": {
                         "status": {"$ref": "#/components/schemas/Status"},
-                        "output": {"title": "Output", "type": "string"},
+                        "output": {"$ref": "#/components/schemas/Output"},
                         "error": {"title": "Error", "type": "string"},
                     },
                     "description": "The status of a prediction.",
@@ -203,6 +208,150 @@ def predict(
     }
 
 
+def test_openapi_specification_with_custom_user_defined_output_type():
+    # Calling this `MyOutput` to test if cog renames it to `Output` in the schema
+    class MyOutput(BaseModel):
+        foo_number: int = "42"
+        foo_string: str = "meaning of life"
+
+    class Predictor(cog.Predictor):
+        def predict(
+            self,
+        ) -> MyOutput:
+            pass
+
+    client = make_client(Predictor())
+    resp = client.get("/openapi.json")
+    assert resp.status_code == 200
+    print(resp.json())
+
+    assert resp.json() == {
+        "openapi": "3.0.2",
+        "info": {"title": "Cog", "version": "0.1.0"},
+        "paths": {
+            "/": {
+                "get": {
+                    "summary": "Root",
+                    "operationId": "root__get",
+                    "responses": {
+                        "200": {
+                            "description": "Successful Response",
+                            "content": {"application/json": {"schema": {}}},
+                        }
+                    },
+                }
+            },
+            "/predictions": {
+                "post": {
+                    "summary": "Predict",
+                    "operationId": "predict_predictions_post",
+                    "requestBody": {
+                        "content": {
+                            "application/json": {
+                                "schema": {"$ref": "#/components/schemas/Request"}
+                            }
+                        }
+                    },
+                    "responses": {
+                        "200": {
+                            "description": "Successful Response",
+                            "content": {
+                                "application/json": {
+                                    "schema": {"$ref": "#/components/schemas/Response"}
+                                }
+                            },
+                        },
+                        "422": {
+                            "description": "Validation Error",
+                            "content": {
+                                "application/json": {
+                                    "schema": {
+                                        "$ref": "#/components/schemas/HTTPValidationError"
+                                    }
+                                }
+                            },
+                        },
+                    },
+                }
+            },
+        },
+        "components": {
+            "schemas": {
+                "HTTPValidationError": {
+                    "title": "HTTPValidationError",
+                    "type": "object",
+                    "properties": {
+                        "detail": {
+                            "title": "Detail",
+                            "type": "array",
+                            "items": {"$ref": "#/components/schemas/ValidationError"},
+                        }
+                    },
+                },
+                "Input": {"title": "Input", "type": "object", "properties": {}},
+                "MyOutput": {
+                    "title": "MyOutput",
+                    "type": "object",
+                    "properties": {
+                        "foo_number": {
+                            "title": "Foo Number",
+                            "type": "integer",
+                            "default": "42",
+                        },
+                        "foo_string": {
+                            "title": "Foo String",
+                            "type": "string",
+                            "default": "meaning of life",
+                        },
+                    },
+                },
+                "Output": {"$ref": "#/components/schemas/MyOutput", "title": "Output"},
+                "Request": {
+                    "title": "Request",
+                    "type": "object",
+                    "properties": {
+                        "input": {"$ref": "#/components/schemas/Input"},
+                        "output_file_prefix": {
+                            "title": "Output File Prefix",
+                            "type": "string",
+                        },
+                    },
+                },
+                "Response": {
+                    "title": "Response",
+                    "required": ["status"],
+                    "type": "object",
+                    "properties": {
+                        "status": {"$ref": "#/components/schemas/Status"},
+                        "output": {"$ref": "#/components/schemas/Output"},
+                        "error": {"title": "Error", "type": "string"},
+                    },
+                    "description": "The status of a prediction.",
+                },
+                "Status": {
+                    "title": "Status",
+                    "enum": ["processing", "success", "failed"],
+                    "description": "An enumeration.",
+                },
+                "ValidationError": {
+                    "title": "ValidationError",
+                    "required": ["loc", "msg", "type"],
+                    "type": "object",
+                    "properties": {
+                        "loc": {
+                            "title": "Location",
+                            "type": "array",
+                            "items": {"type": "string"},
+                        },
+                        "msg": {"title": "Message", "type": "string"},
+                        "type": {"title": "Error Type", "type": "string"},
+                    },
+                },
+            }
+        },
+    }
+
+
 def test_yielding_strings_from_generator_predictors():
     class Predictor(cog.Predictor):
         def predict(self) -> Generator[str, None, None]: