Fix dereferencing in Function and AsyncFunction

We were accidentally stripping out all components in the schemas due to
the Input/Output being referenced by PredictionRequest and
PredictionResponse and those in turn being referenced by the various
`path` entries. This commit ensures we retain Input and Output and
cleans up the object to only contain relevant fields.

Author: aron

replicate/use.py

@@ -190,7 +190,17 @@ def _resolve_ref(obj: Any) -> Any:
 
     result = _resolve_ref(dereferenced)
 
-    # Filter out any references that have now been referenced.
+    # Remove "paths" as these aren't relevant to models.
+    result["paths"] = {}
+
+    # Retain Input and Output schemas as these are important.
+    dereferenced_refs.discard("Input")
+    dereferenced_refs.discard("Output")
+
+    dereferenced_refs.discard("TrainingInput")
+    dereferenced_refs.discard("TrainingOutput")
+
+    # Filter out any remaining references that have been inlined.
     result["components"]["schemas"] = {
         k: v
         for k, v in result["components"]["schemas"].items()

tests/test_use.py

@@ -78,6 +78,82 @@ def create_mock_version(version_overrides=None, version_id="xyz123"):
                         "required": ["prompt"],
                     },
                     "Output": {"type": "string", "title": "Output"},
+                    "PredictionResponse": {
+                        "type": "object",
+                        "title": "PredictionResponse",
+                        "properties": {
+                            "id": {"type": "string", "title": "Id"},
+                            "logs": {"type": "string", "title": "Logs", "default": ""},
+                            "error": {"type": "string", "title": "Error"},
+                            "input": {"$ref": "#/components/schemas/Input"},
+                            "output": {"$ref": "#/components/schemas/Output"},
+                            "status": {"$ref": "#/components/schemas/Status"},
+                            "metrics": {"type": "object", "title": "Metrics"},
+                            "version": {"type": "string", "title": "Version"},
+                            "created_at": {
+                                "type": "string",
+                                "title": "Created At",
+                                "format": "date-time",
+                            },
+                            "started_at": {
+                                "type": "string",
+                                "title": "Started At",
+                                "format": "date-time",
+                            },
+                            "completed_at": {
+                                "type": "string",
+                                "title": "Completed At",
+                                "format": "date-time",
+                            },
+                        },
+                    },
+                    "PredictionRequest": {
+                        "type": "object",
+                        "title": "PredictionRequest",
+                        "properties": {
+                            "id": {"type": "string", "title": "Id"},
+                            "input": {"$ref": "#/components/schemas/Input"},
+                            "webhook": {
+                                "type": "string",
+                                "title": "Webhook",
+                                "format": "uri",
+                                "maxLength": 65536,
+                                "minLength": 1,
+                            },
+                            "created_at": {
+                                "type": "string",
+                                "title": "Created At",
+                                "format": "date-time",
+                            },
+                            "output_file_prefix": {
+                                "type": "string",
+                                "title": "Output File Prefix",
+                            },
+                            "webhook_events_filter": {
+                                "type": "array",
+                                "items": {"$ref": "#/components/schemas/WebhookEvent"},
+                                "default": ["start", "output", "logs", "completed"],
+                            },
+                        },
+                    },
+                    "Status": {
+                        "enum": [
+                            "starting",
+                            "processing",
+                            "succeeded",
+                            "canceled",
+                            "failed",
+                        ],
+                        "type": "string",
+                        "title": "Status",
+                        "description": "An enumeration.",
+                    },
+                    "WebhookEvent": {
+                        "enum": ["start", "output", "logs", "completed"],
+                        "type": "string",
+                        "title": "WebhookEvent",
+                        "description": "An enumeration.",
+                    },
                 }
             },
         },
@@ -345,6 +421,7 @@ async def test_use_function_openapi_schema_dereferenced(client_mode):
                     "openapi_schema": {
                         "components": {
                             "schemas": {
+                                "Extra": {"type": "object"},
                                 "Output": {"$ref": "#/components/schemas/ModelOutput"},
                                 "ModelOutput": {
                                     "type": "object",
@@ -374,6 +451,12 @@ async def test_use_function_openapi_schema_dereferenced(client_mode):
     else:
         schema = hotdog_detector.openapi_schema()
 
+    assert schema["components"]["schemas"]["Extra"] == {"type": "object"}
+    assert schema["components"]["schemas"]["Input"] == {
+        "type": "object",
+        "properties": {"prompt": {"type": "string", "title": "Prompt"}},
+        "required": ["prompt"],
+    }
     assert schema["components"]["schemas"]["Output"] == {
         "type": "object",
         "properties": {
@@ -386,7 +469,14 @@ async def test_use_function_openapi_schema_dereferenced(client_mode):
         },
     }
 
+    # Assert everything else is stripped out
+    assert schema["paths"] == {}
+
+    assert "PredictionRequest" not in schema["components"]["schemas"]
+    assert "PredictionResponse" not in schema["components"]["schemas"]
     assert "ModelOutput" not in schema["components"]["schemas"]
+    assert "Status" not in schema["components"]["schemas"]
+    assert "WebhookEvent" not in schema["components"]["schemas"]
 
 
 @pytest.mark.asyncio