Allow complex output types (#218)

* Allow complex output types

- Removes an assertion during schema introspection that prevented list types as a field.

- Adds a test verifying that excessively complex output types work correctly

* I think moving out of the bad_predictors dir stops evaluating ERROR?

* add schema

* go e2e test

* fix encoding

* regenerate type stubs, fix command in CI check error

Author: michaeldwan

.github/workflows/ci.yml

@@ -62,7 +62,7 @@ jobs:
           echo "Generated stubs differ from committed versions:"
           git diff python/
           echo ""
-          echo "Run 'script/check.sh stubs' locally and commit the updated .pyi files"
+          echo "Run 'script/check.sh python' locally and commit the updated .pyi files"
           exit 1
         fi
         echo "All stub files are up to date"

internal/tests/harness_test.go

@@ -506,8 +506,17 @@ func writeCogConfig(t *testing.T, tempDir, predictorClass string, concurrencyMax
 // FIXME: this is a hack to provide compatibility with the `cog_test` test harness while we migrate to in-process testing.
 func linkPythonModule(t *testing.T, basePath, tempDir, module string) {
 	t.Helper()
+
+	// Try runners directory first (for backward compatibility)
 	runnersPath := path.Join(basePath, "python", "tests", "runners")
 	srcPath := path.Join(runnersPath, fmt.Sprintf("%s.py", module))
+
+	// If not found in runners, try schemas directory
+	if _, err := os.Stat(srcPath); os.IsNotExist(err) {
+		schemasPath := path.Join(basePath, "python", "tests", "schemas")
+		srcPath = path.Join(schemasPath, fmt.Sprintf("%s.py", module))
+	}
+
 	dstPath := path.Join(tempDir, "predict.py")
 
 	// Debug logging

internal/tests/output_test.go

@@ -51,3 +51,102 @@ func TestPredictionOutputSucceeded(t *testing.T) {
 	}
 	assert.Equal(t, expectedOutput, predictionResponse.Output)
 }
+
+func TestComplexOutputTypes(t *testing.T) {
+	t.Parallel()
+	if *legacyCog {
+		t.Skip("legacy Cog does not support complex output types")
+	}
+	runtimeServer := setupCogRuntime(t, cogRuntimeServerConfig{
+		procedureMode:    false,
+		explicitShutdown: true,
+		uploadURL:        "",
+		module:           "output_complex_types",
+		predictorClass:   "Predictor",
+	})
+	waitForSetupComplete(t, runtimeServer, runner.StatusReady, runner.SetupSucceeded)
+
+	input := map[string]any{"s": "test"}
+	req := httpPredictionRequest(t, runtimeServer, runner.PredictionRequest{Input: input})
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+	body, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	var predictionResponse server.PredictionResponse
+	err = json.Unmarshal(body, &predictionResponse)
+	require.NoError(t, err)
+
+	// Create expected output using JSON round-trip to match server serialization
+	expectedOutputs := []map[string]any{
+		{
+			"strings": []string{"hello", "world"},
+			"numbers": []int{1, 2, 3},
+			"single_item": map[string]any{
+				"name":  "item1",
+				"value": 42,
+			},
+			"items": []map[string]any{
+				{"name": "item1", "value": 42},
+				{"name": "item2", "value": 84},
+			},
+			"container": map[string]any{
+				"items": []map[string]any{
+					{"name": "item1", "value": 42},
+					{"name": "item2", "value": 84},
+				},
+				"tags": []string{"tag1", "tag2"},
+				"nested": map[string]any{
+					"item":        map[string]any{"name": "item1", "value": 42},
+					"description": "nested description",
+				},
+				"optional_list": []string{"opt1", "opt2"},
+				"count":         2,
+			},
+			"nested_items": []map[string]any{
+				{
+					"item":        map[string]any{"name": "item1", "value": 42},
+					"description": "nested description",
+				},
+			},
+		},
+		{
+			"strings": []string{"foo", "bar"},
+			"numbers": []int{4, 5, 6},
+			"single_item": map[string]any{
+				"name":  "item2",
+				"value": 84,
+			},
+			"items": []map[string]any{
+				{"name": "item2", "value": 84},
+			},
+			"container": map[string]any{
+				"items": []map[string]any{
+					{"name": "item1", "value": 42},
+					{"name": "item2", "value": 84},
+				},
+				"tags": []string{"tag1", "tag2"},
+				"nested": map[string]any{
+					"item":        map[string]any{"name": "item1", "value": 42},
+					"description": "nested description",
+				},
+				"optional_list": []string{"opt1", "opt2"},
+				"count":         2,
+			},
+			"nested_items": []map[string]any{
+				{
+					"item":        map[string]any{"name": "item1", "value": 42},
+					"description": "nested description",
+				},
+			},
+		},
+	}
+	expectedJSON, err := json.Marshal(expectedOutputs)
+	require.NoError(t, err)
+	var expectedOutput []any
+	err = json.Unmarshal(expectedJSON, &expectedOutput)
+	require.NoError(t, err)
+	assert.Equal(t, expectedOutput, predictionResponse.Output)
+	assert.Equal(t, runner.PredictionSucceeded, predictionResponse.Status)
+}

python/coglet/inspector.py

@@ -187,9 +187,6 @@ def _output_adt(tpe: type) -> adt.Output:
         fields = {}
         for name, t in tpe.__annotations__.items():
             ft = adt.FieldType.from_type(t)
-            assert ft.repetition is not adt.Repetition.REPEATED, (
-                f'output field must not be list: {name}: {ft.python_type()}'
-            )
             fields[name] = ft
         return adt.Output(kind=adt.Kind.OBJECT, fields=fields)
 

python/coglet/util.py

@@ -19,6 +19,9 @@ def output_json(obj):
     elif tpe is api.Secret:
         # Encode Secret('foobar') as '**********'
         return str(obj)
+    elif hasattr(obj, '__dict__') and hasattr(obj, '__dataclass_fields__'):
+        # Handle dataclass objects (including BaseModel)
+        return {field: getattr(obj, field) for field in obj.__dataclass_fields__}
     else:
         raise TypeError(f'Object of type {tpe} is not JSON serializable')
 

python/coglet/util.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 def now_iso() -> str:
     ...
 
-def output_json(obj): # -> str:
+def output_json(obj): # -> str | dict[Any, Any]:
     ...
 
 def schema_json(obj): # -> str: