Define Secret type (#1546)

Signed-off-by: Mattt Zmuda <mattt@replicate.com>

python/cog/__init__.py

@@ -2,7 +2,14 @@
 
 from .predictor import BasePredictor
 from .server.worker import emit_metric
-from .types import AsyncConcatenateIterator, ConcatenateIterator, File, Input, Path
+from .types import (
+    AsyncConcatenateIterator,
+    ConcatenateIterator,
+    File,
+    Input,
+    Path,
+    Secret,
+)
 
 try:
     from ._version import __version__
@@ -20,4 +27,5 @@
     "Input",
     "Path",
     "emit_metric",
+    "Secret",
 ]

python/cog/predictor.py

@@ -40,8 +40,17 @@
 from .types import (
     Path as CogPath,
 )
-
-ALLOWED_INPUT_TYPES: List[Type[Any]] = [str, int, float, bool, CogFile, CogPath]
+from .types import Secret as CogSecret
+
+ALLOWED_INPUT_TYPES: List[Type[Any]] = [
+    str,
+    int,
+    float,
+    bool,
+    CogFile,
+    CogPath,
+    CogSecret,
+]
 
 
 class BasePredictor(ABC):

python/cog/types.py

@@ -11,7 +11,7 @@
 
 import httpx
 import requests
-from pydantic import Field
+from pydantic import Field, SecretStr
 
 FILENAME_ILLEGAL_CHARS = set("\u0000/")
 
@@ -44,6 +44,19 @@ def Input(
     )
 
 
+class Secret(SecretStr):
+    @classmethod
+    def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
+        """Defines what this type should be in openapi.json"""
+        field_schema.update(
+            {
+                "type": "string",
+                "format": "password",
+                "x-cog-secret": True,
+            }
+        )
+
+
 class File(io.IOBase):
     validate_always = True
 

python/tests/server/fixtures/input_secret.py

@@ -0,0 +1,6 @@
+from cog import BasePredictor, Secret
+
+
+class Predictor(BasePredictor):
+    def predict(self, secret: Secret) -> str:
+        return secret.get_secret_value()

python/tests/server/test_http_input.py

@@ -247,6 +247,16 @@ def test_union_integers(client):
     assert resp.status_code == 422
 
 
+@uses_predictor("input_secret")
+def test_secret_str(client, match):
+    resp = client.post("/predictions", json={"input": {"secret": "foo"}})
+    assert resp.status_code == 200
+    assert resp.json() == match({"output": "foo", "status": "succeeded"})
+
+    resp = client.post("/predictions", json={"input": {"secret": {}}})
+    assert resp.status_code == 422
+
+
 def test_untyped_inputs():
     config = {"predict": _fixture_path("input_untyped")}
     app = create_app(

python/tests/test_types.py

@@ -4,7 +4,7 @@
 
 import pytest
 import responses
-from cog.types import URLFile, get_filename_from_url, get_filename_from_urlopen
+from cog.types import Secret, URLFile, get_filename_from_url, get_filename_from_urlopen
 
 
 @responses.activate
@@ -123,3 +123,11 @@ def test_get_filename(url, filename):
 def test_get_filename_from_urlopen(url, filename):
     resp = urllib.request.urlopen(url)  # noqa: S310
     assert get_filename_from_urlopen(resp) == filename
+
+
+def test_secret_type():
+    secret_value = "sw0rdf1$h"  # noqa: S105
+    secret = Secret(secret_value)
+
+    assert secret.get_secret_value() == secret_value
+    assert str(secret) == "**********"