Add coder support

Author: nevillelyh

internal/server/path.go

@@ -32,6 +32,12 @@ func handlePath(json any, paths *[]string, fn func(string, *[]string) (string, e
 					return nil, err
 				}
 				xs[i] = o
+			} else {
+				o, err := handlePath(xs[i], paths, fn)
+				if err != nil {
+					return nil, err
+				}
+				xs[i] = o
 			}
 		}
 		return xs, nil
@@ -43,6 +49,12 @@ func handlePath(json any, paths *[]string, fn func(string, *[]string) (string, e
 					return nil, err
 				}
 				m[key] = o
+			} else {
+				o, err := handlePath(m[key], paths, fn)
+				if err != nil {
+					return nil, err
+				}
+				m[key] = o
 			}
 		}
 		return m, nil

internal/tests/coder_test.go

@@ -0,0 +1,64 @@
+package tests
+
+import (
+	"testing"
+
+	"github.com/replicate/cog-runtime/internal/server"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPredictionDataclassCoderSucceeded(t *testing.T) {
+	if *legacyCog {
+		// Compat: legacy Cog does not support custom coder
+		t.SkipNow()
+	}
+	ct := NewCogTest(t, "dataclass")
+	assert.NoError(t, ct.Start())
+
+	hc := ct.WaitForSetup()
+	assert.Equal(t, server.StatusReady.String(), hc.Status)
+	assert.Equal(t, server.SetupSucceeded, hc.Setup.Status)
+
+	resp := ct.Prediction(map[string]any{
+		"account": map[string]any{
+			"id":          0,
+			"name":        "John",
+			"address":     map[string]any{"street": "Smith", "zip": 12345},
+			"credentials": map[string]any{"password": "foo", "pubkey": b64encode("bar")},
+		},
+	})
+
+	output := map[string]any{
+		"account": map[string]any{
+			"id":          100.0,
+			"name":        "JOHN",
+			"address":     map[string]any{"street": "SMITH", "zip": 22345.0},
+			"credentials": map[string]any{"password": "**********", "pubkey": b64encode("*bar*")},
+		},
+	}
+	ct.AssertResponse(resp, server.PredictionSucceeded, output, "")
+
+	ct.Shutdown()
+	assert.NoError(t, ct.Cleanup())
+}
+
+func TestPredictionChatCoderSucceeded(t *testing.T) {
+	if *legacyCog {
+		// Compat: legacy Cog does not support custom coder
+		t.SkipNow()
+	}
+	ct := NewCogTest(t, "chat")
+	assert.NoError(t, ct.Start())
+
+	hc := ct.WaitForSetup()
+	assert.Equal(t, server.StatusReady.String(), hc.Status)
+	assert.Equal(t, server.SetupSucceeded, hc.Setup.Status)
+
+	resp := ct.Prediction(map[string]any{"msg": map[string]any{"role": "assistant", "content": "bar"}})
+	output := map[string]any{"role": "assistant", "content": "*bar*"}
+	ct.AssertResponse(resp, server.PredictionSucceeded, output, "")
+
+	ct.Shutdown()
+	assert.NoError(t, ct.Cleanup())
+}

pyproject.toml

@@ -22,12 +22,17 @@ dev = [
 ]
 
 test = [
+    'openai',
     'pytest-cov',
     'pytest',
     'pytest-asyncio',
     'tqdm',
 ]
 
+provided = [
+    'pydantic',
+]
+
 [build-system]
 requires = ['setuptools', 'setuptools-scm']
 build-backend = 'setuptools.build_meta'

python/cog/__init__.py

@@ -3,6 +3,7 @@
     BaseModel,
     BasePredictor,
     CancelationException,
+    Coder,
     ConcatenateIterator,
     Input,
     Path,
@@ -16,6 +17,7 @@
     'BaseModel',
     'BasePredictor',
     'CancelationException',
+    'Coder',
     'ConcatenateIterator',
     'Input',
     'Path',

python/cog/coder/__init__.py

@@ -0,0 +1,56 @@
+import dataclasses
+from typing import Any, Optional, Type
+
+from coglet import api
+
+
+class DataclassCoder(api.Coder):
+    @staticmethod
+    def factory(cls: Type) -> Optional[api.Coder]:
+        if dataclasses.is_dataclass(cls):
+            return DataclassCoder(cls)
+        else:
+            return None
+
+    def __init__(self, cls: Type):
+        assert dataclasses.is_dataclass(cls)
+        self.cls = cls
+
+    def encode(self, x: Any) -> dict[str, Any]:
+        # Secret is a dataclass and dataclasses.asdict recursively converts its internals
+        return self._to_dict(self.cls, x)
+
+    def _to_dict(self, cls: Type, x: Any) -> dict[str, Any]:
+        r: dict[str, Any] = {}
+        for f in dataclasses.fields(cls):
+            v = getattr(x, f.name)
+            # Keep Path and Secret as is and let json.dumps(default=fn) handle them
+            if f.type is api.Path:
+                v = api.Path(v)
+            elif f.type is api.Secret:
+                v = api.Secret(v)
+            elif dataclasses.is_dataclass(v):
+                v = self._to_dict(f.type, v)  # type: ignore
+            r[f.name] = v
+        return r
+
+    def decode(self, x: dict[str, Any]) -> Any:
+        kwargs = self._from_dict(self.cls, x)
+        return self.cls(**kwargs)  # type: ignore
+
+    def _from_dict(self, cls: Type, x: dict[str, Any]) -> Any:
+        r: dict[str, Any] = {}
+        for f in dataclasses.fields(cls):
+            if f.name not in x:
+                continue
+            elif f.type is api.Path:
+                r[f.name] = api.Path(x[f.name])
+            # Secret is a dataclass and must be handled before other dataclasses
+            elif f.type is api.Secret:
+                r[f.name] = api.Secret(x[f.name])
+            elif dataclasses.is_dataclass(f.type):
+                kwargs = self._from_dict(f.type, x[f.name])  # type: ignore
+                r[f.name] = f.type(**kwargs)  # type: ignore
+            else:
+                r[f.name] = x[f.name]
+        return r

python/cog/coder/dataclass_coder.py

@@ -0,0 +1,19 @@
+import typing
+from typing import Any, Optional, Type
+
+from coglet import api
+
+
+class JsonCoder(api.Coder):
+    @staticmethod
+    def factory(cls: Type) -> Optional[api.Coder]:
+        if typing.get_origin(cls) is dict is dict and typing.get_args(cls)[0] is str:
+            return JsonCoder()
+        else:
+            return None
+
+    def encode(self, x: Any) -> dict[str, Any]:
+        return x
+
+    def decode(self, x: dict[str, Any]) -> Any:
+        return x

python/cog/coder/json_coder.py

@@ -0,0 +1,24 @@
+import inspect
+from typing import Any, Type
+
+from pydantic import BaseModel
+
+from coglet import api
+
+
+class BaseModelCoder(api.Coder):
+    @staticmethod
+    def factory(cls: Type):
+        if cls is not BaseModel and any(c is BaseModel for c in inspect.getmro(cls)):
+            return BaseModelCoder(cls)
+        else:
+            return None
+
+    def __init__(self, cls: Type[BaseModel]):
+        self.cls = cls
+
+    def encode(self, x: BaseModel) -> dict[str, Any]:
+        return x.model_dump(exclude_unset=True)
+
+    def decode(self, x: dict[str, Any]) -> BaseModel:
+        return self.cls.model_construct(**x)