fix(go): allow multiple schemas using picoschema

go/plugins/dotprompt/picoschema.go

@@ -64,15 +64,16 @@ func parsePico(val any) (*jsonschema.Schema, error) {
 	case string:
 		typ, desc, found := strings.Cut(val, ",")
 		switch typ {
-		case "string", "boolean", "null", "number", "integer", "any":
+		case "string", "boolean", "null", "number", "integer":
+		case "any":
+			typ = ""
 		default:
 			return nil, fmt.Errorf("picoschema: unsupported scalar type %q", typ)
 		}
-		if typ == "any" {
-			typ = ""
-		}
-		ret := &jsonschema.Schema{
-			Type: typ,
+
+		ret := &jsonschema.Schema{}
+		if typ != "" {
+			ret.Type = typ
 		}
 		if found {
 			ret.Description = strings.TrimSpace(desc)
@@ -100,40 +101,41 @@ func parsePico(val any) (*jsonschema.Schema, error) {
 				return nil, err
 			}
 
-			if !found {
-				ret.Properties.Set(propertyName, property)
-				continue
-			}
-
-			typ = strings.TrimSuffix(typ, ")")
-			typ, desc, found := strings.Cut(strings.TrimSuffix(typ, ")"), ",")
-			switch typ {
-			case "array":
-				property = &jsonschema.Schema{
-					Type:  "array",
-					Items: property,
-				}
-			case "object":
-				// Use property unchanged.
-			case "enum":
-				if property.Enum == nil {
-					return nil, fmt.Errorf("picoschema: enum value %v is not an array", property)
-				}
-				if isOptional {
-					property.Enum = append(property.Enum, nil)
+			if found {
+				typ = strings.TrimSuffix(typ, ")")
+				typ, desc, found := strings.Cut(strings.TrimSuffix(typ, ")"), ",")
+				switch typ {
+				case "array":
+					property = &jsonschema.Schema{
+						Type:  "array",
+						Items: property,
+					}
+				case "object":
+					// Use property unchanged.
+				case "enum":
+					if property.Enum == nil {
+						return nil, fmt.Errorf("picoschema: enum value %v is not an array", property)
+					}
+
+					if isOptional {
+						property.Enum = append(property.Enum, nil)
+					}
+
+				case "*":
+					ret.AdditionalProperties = property
+					continue
+				default:
+					return nil, fmt.Errorf("picoschema: parenthetical type %q is none of %q", typ,
+						[]string{"object", "array", "enum", "*"})
 				}
 
-			case "*":
-				ret.AdditionalProperties = property
-				continue
-			default:
-				return nil, fmt.Errorf("picoschema: parenthetical type %q is none of %q", typ,
-					[]string{"object", "array", "enum", "*"})
-
+				if found {
+					property.Description = strings.TrimSpace(desc)
+				}
 			}
 
-			if found {
-				property.Description = strings.TrimSpace(desc)
+			if isOptional {
+				property = makeNullable(property)
 			}
 
 			ret.Properties.Set(propertyName, property)
@@ -268,3 +270,41 @@ func mapToJSONSchema(m map[string]any) (*jsonschema.Schema, error) {
 
 	return &ret, nil
 }
+
+func makeNullable(schema *jsonschema.Schema) *jsonschema.Schema {
+	// Do not wrap enums in anyOf
+	if len(schema.Enum) > 0 {
+		return schema
+	}
+	// If the schema is empty (represents 'any'), do not wrap it
+	if schema.Type == "" &&
+		(schema.Properties == nil || schema.Properties.Len() == 0) &&
+		schema.Items == nil &&
+		len(schema.AnyOf) == 0 &&
+		len(schema.AllOf) == 0 &&
+		len(schema.OneOf) == 0 &&
+		len(schema.Enum) == 0 {
+		return schema
+	}
+	// Check if the schema already allows null
+	if schema.Type == "null" {
+		return schema
+	}
+	if len(schema.AnyOf) > 0 {
+		for _, s := range schema.AnyOf {
+			if s.Type == "null" {
+				return schema
+			}
+		}
+	}
+	// Wrap the original schema to allow null
+	return &jsonschema.Schema{
+		AnyOf: []*jsonschema.Schema{
+			schema,
+			{Type: "null"},
+		},
+		Description:          schema.Description,
+		AdditionalProperties: schema.AdditionalProperties,
+		Properties:           schema.Properties,
+	}
+}

go/plugins/dotprompt/picoschema_test.go

@@ -15,8 +15,10 @@
 package dotprompt
 
 import (
+	"encoding/json"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -41,16 +43,8 @@ func TestPicoschema(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	skip := map[string]bool{
-		"required field":                 true,
-		"nested object in array and out": true,
-	}
-
 	for _, test := range tests {
 		t.Run(test.Description, func(t *testing.T) {
-			if skip[test.Description] {
-				t.Skip("no support for type as an array")
-			}
 			var val any
 			if err := yaml.Unmarshal([]byte(test.YAML), &val); err != nil {
 				t.Fatalf("YAML unmarshal failure: %v", err)
@@ -67,8 +61,17 @@ func TestPicoschema(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
+			gotData, err := json.Marshal(got)
+			if err != nil {
+				t.Fatal(err)
+			}
+			var gotMap map[string]any
+			if err := json.Unmarshal(gotData, &gotMap); err != nil {
+				t.Fatal(err)
+			}
+			replaceAnyOfWithTypeArray(gotMap)
 			want := replaceEmptySchemas(test.Want)
-			if diff := cmp.Diff(want, got); diff != "" {
+			if diff := cmp.Diff(want, gotMap); diff != "" {
 				t.Errorf("mismatch (-want, +got):\n%s", diff)
 			}
 		})
@@ -97,3 +100,54 @@ func replaceEmptySchemas(m map[string]any) any {
 	}
 	return m
 }
+
+func replaceAnyOfWithTypeArray(schema map[string]any) {
+	// Check if 'anyOf' is present
+	if anyOf, ok := schema["anyOf"].([]any); ok && len(anyOf) > 0 {
+		types := []any{}
+		descriptions := []string{}
+		otherKeysExist := false
+
+		for _, item := range anyOf {
+			if subSchema, ok := item.(map[string]any); ok {
+				// Collect 'type' and 'description' from sub-schemas
+				if t, hasType := subSchema["type"]; hasType {
+					types = append(types, t)
+				} else {
+					otherKeysExist = true
+					break
+				}
+				if desc, hasDesc := subSchema["description"]; hasDesc {
+					descriptions = append(descriptions, desc.(string))
+				}
+			} else {
+				otherKeysExist = true
+				break
+			}
+		}
+
+		// Replace 'anyOf' with 'type' array if no other keys exist
+		if !otherKeysExist && len(types) > 0 {
+			schema["type"] = types
+			delete(schema, "anyOf")
+			// Combine descriptions if necessary
+			if len(descriptions) > 0 && schema["description"] == nil {
+				schema["description"] = strings.Join(descriptions, "; ")
+			}
+		}
+	}
+
+	// Recursively process nested schemas
+	for _, value := range schema {
+		switch v := value.(type) {
+		case map[string]any:
+			replaceAnyOfWithTypeArray(v)
+		case []any:
+			for _, item := range v {
+				if m, ok := item.(map[string]any); ok {
+					replaceAnyOfWithTypeArray(m)
+				}
+			}
+		}
+	}
+}