From bedde58b2d50b9cd6bf376558e3893a77f812597 Mon Sep 17 00:00:00 2001 From: fernandoalonso Date: Tue, 26 Nov 2024 00:03:50 +0000 Subject: [PATCH 1/4] Fix: allow multiple schemas using picoschema --- go/plugins/dotprompt/picoschema.go | 112 +++++++++++++++++------- go/plugins/dotprompt/picoschema_test.go | 72 +++++++++++++-- 2 files changed, 141 insertions(+), 43 deletions(-) diff --git a/go/plugins/dotprompt/picoschema.go b/go/plugins/dotprompt/picoschema.go index ca5b7c00c..46a1192cf 100644 --- a/go/plugins/dotprompt/picoschema.go +++ b/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,45 @@ func parsePico(val any) (*jsonschema.Schema, error) { return nil, err } - if !found { - ret.Properties.Set(propertyName, property) - continue + 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", "*"}) + + } } - 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", "*"}) + 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 +274,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, + } +} diff --git a/go/plugins/dotprompt/picoschema_test.go b/go/plugins/dotprompt/picoschema_test.go index c4cf66937..24e61dcf3 100644 --- a/go/plugins/dotprompt/picoschema_test.go +++ b/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) + } + } + } + } +} \ No newline at end of file From d8f052b6cfa9089f88928afbffa43e1df235bee5 Mon Sep 17 00:00:00 2001 From: fernandoalonso Date: Tue, 26 Nov 2024 00:18:21 +0000 Subject: [PATCH 2/4] Fix: attempt to fix unit test --- go/plugins/dotprompt/picoschema.go | 1 + 1 file changed, 1 insertion(+) diff --git a/go/plugins/dotprompt/picoschema.go b/go/plugins/dotprompt/picoschema.go index 46a1192cf..075b5f85c 100644 --- a/go/plugins/dotprompt/picoschema.go +++ b/go/plugins/dotprompt/picoschema.go @@ -146,6 +146,7 @@ func parsePico(val any) (*jsonschema.Schema, error) { } return ret, nil } + return nil, fmt.Errorf("picoschema: value %v of type %[1]T is not an object, slice or string", val) } // mapToJSONSchema converts a YAML value to a JSONSchema. From ff44e91ffade012514e35bc7eefd82a4800a2010 Mon Sep 17 00:00:00 2001 From: fernandoalonso Date: Tue, 26 Nov 2024 00:43:19 +0000 Subject: [PATCH 3/4] Fix: attempt to fix unit test --- go/plugins/dotprompt/picoschema.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/go/plugins/dotprompt/picoschema.go b/go/plugins/dotprompt/picoschema.go index 075b5f85c..765e04e39 100644 --- a/go/plugins/dotprompt/picoschema.go +++ b/go/plugins/dotprompt/picoschema.go @@ -127,10 +127,6 @@ func parsePico(val any) (*jsonschema.Schema, error) { default: return nil, fmt.Errorf("picoschema: parenthetical type %q is none of %q", typ, []string{"object", "array", "enum", "*"}) - - } - } - } if found { @@ -146,7 +142,6 @@ func parsePico(val any) (*jsonschema.Schema, error) { } return ret, nil } - return nil, fmt.Errorf("picoschema: value %v of type %[1]T is not an object, slice or string", val) } // mapToJSONSchema converts a YAML value to a JSONSchema. From f67c3a076d4008b027a5fc2280354612567e7b4a Mon Sep 17 00:00:00 2001 From: alonsopec89 Date: Fri, 10 Jan 2025 11:31:17 -0600 Subject: [PATCH 4/4] Fix: test action_test --- go/core/action_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/go/core/action_test.go b/go/core/action_test.go index 5bf24b5a0..47a4000f3 100644 --- a/go/core/action_test.go +++ b/go/core/action_test.go @@ -17,6 +17,7 @@ package core import ( "bytes" "context" + "fmt" "slices" "testing" @@ -124,7 +125,7 @@ func TestActionTracing(t *testing.T) { // The same trace store is used for all tests, so there might be several traces. // Look for this one, which has a unique name. for _, td := range tc.Traces { - if td.DisplayName == actionName { + if td.DisplayName == actionName || td.DisplayName == fmt.Sprintf("test/%s", actionName) { // Spot check: expect a single span. if g, w := len(td.Spans), 1; g != w { t.Errorf("got %d spans, want %d", g, w)