fix: invalid components: schema "Output": extra sibling fields: [titl…

…e] (#1579)

ref: #1275

Signed-off-by: Dmitri Khokhlov <dkhokhlov@gmail.com>

python/cog/command/openapi_schema.py

@@ -5,13 +5,34 @@
 """
 
 import json
+from typing import Any, Dict, List, Union
 
 from ..errors import CogError, ConfigDoesNotExist, PredictorNotSet
 from ..predictor import load_config
 from ..schema import Status
 from ..server.http import create_app
 from ..suppress_output import suppress_output
 
+
+def remove_title_next_to_ref(
+    schema_node: Union[Dict[str, Any], List[Any]],
+) -> Union[Dict[str, Any], List[Any]]:
+    """
+    Recursively remove 'title' from schema components that have a '$ref'.
+    This function addresses a non-compliance issue in FastAPI's OpenAPI schema generation, where
+    'title' fields adjacent to '$ref' fields can cause validation problems with some OpenAPI tools.
+    """
+    if isinstance(schema_node, dict):
+        if "$ref" in schema_node and "title" in schema_node:
+            del schema_node["title"]
+        for _key, value in schema_node.items():
+            remove_title_next_to_ref(value)
+    elif isinstance(schema_node, list):  # type: ignore[reportUnnecessaryIsInstance]
+        for i, item in enumerate(schema_node):
+            schema_node[i] = remove_title_next_to_ref(item)
+    return schema_node
+
+
 if __name__ == "__main__":
     schema = {}
     try:
@@ -23,7 +44,7 @@
                 and app.state.setup_result.status == Status.FAILED
             ):
                 raise CogError(app.state.setup_result.logs)
-            schema = app.openapi()
+            schema = remove_title_next_to_ref(app.openapi())
     except (ConfigDoesNotExist, PredictorNotSet):
         # If there is no cog.yaml or 'predict' has not been set, then there is no type signature.
         # Not an error, there just isn't anything.

test-integration/test_integration/fixtures/complex_output_project/cog.yaml

@@ -0,0 +1,3 @@
+build:
+  python_version: "3.11"
+predict: "predict.py:Predictor"

test-integration/test_integration/fixtures/complex_output_project/predict.py

@@ -0,0 +1,17 @@
+import io
+
+from cog import BasePredictor, Path
+from typing import Optional
+from pydantic import BaseModel
+
+
+class ModelOutput(BaseModel):
+    success: bool
+    error: Optional[str]
+    segmentedImage: Optional[Path]
+
+
+class Predictor(BasePredictor):
+    # setup code
+    def predict(self, msg: str) -> ModelOutput:
+       return ModelOutput(success=False, error=msg, segmentedImage=None)

test-integration/test_integration/test_build.py

@@ -197,6 +197,17 @@ def test_build_with_cog_init_templates(tmpdir, docker_image):
     assert "Image built as cog-" in build_process.stderr.decode()
 
 
+def test_build_with_complex_output(tmpdir, docker_image):
+    project_dir = Path(__file__).parent / "fixtures/complex_output_project"
+    build_process = subprocess.run(
+        ["cog", "build", "-t", docker_image],
+        cwd=project_dir,
+        capture_output=True,
+    )
+    assert build_process.returncode == 0
+    assert "Image built as cog-" in build_process.stderr.decode()
+
+
 def test_python_37_depricated(docker_image):
     def test_build_invalid_schema(docker_image):
         project_dir = Path(__file__).parent / "fixtures/python_37"