From 56ee5942578ac42cfb365f49be30a8d8294cb7c8 Mon Sep 17 00:00:00 2001
From: Alex Pascal <apascal07@gmail.com>
Date: Thu, 23 May 2024 09:04:40 -0700
Subject: [PATCH] Added input schema validation.

---
 go/genkit/action.go             | 12 ++++++++----
 go/genkit/flow.go               |  8 +++++++-
 go/samples/flow-sample1/main.go | 14 ++++++++++++++
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/go/genkit/action.go b/go/genkit/action.go
index 13661f02c6..d23af55bbf 100644
--- a/go/genkit/action.go
+++ b/go/genkit/action.go
@@ -195,10 +195,11 @@ func (a *Action[I, O, S]) desc() actionDesc {
 	}
 	// Required by genkit UI:
 	if ad.Metadata == nil {
-		ad.Metadata = map[string]any{}
+		ad.Metadata = map[string]any{
+			"inputSchema":  nil,
+			"outputSchema": nil,
+		}
 	}
-	ad.Metadata["inputSchema"] = nil
-	ad.Metadata["outputSchema"] = nil
 	return ad
 }
 
@@ -214,5 +215,8 @@ func inferJSONSchema(x any) (s *jsonschema.Schema) {
 		// instead of nested inside a "$defs" object.
 		r.ExpandedStruct = true
 	}
-	return r.Reflect(x)
+	s = r.Reflect(x)
+	// TODO: Unwind this change once Monaco Editor supports newer than JSON schema draft-07.
+	s.Version = ""
+	return s
 }
diff --git a/go/genkit/flow.go b/go/genkit/flow.go
index 5f43db1b65..d16a1c59f0 100644
--- a/go/genkit/flow.go
+++ b/go/genkit/flow.go
@@ -230,7 +230,13 @@ type FlowResult[O any] struct {
 
 // action creates an action for the flow. See the comment at the top of this file for more information.
 func (f *Flow[I, O, S]) action() *Action[*flowInstruction[I], *flowState[I, O], S] {
-	return NewStreamingAction(f.name, nil, func(ctx context.Context, inst *flowInstruction[I], cb StreamingCallback[S]) (*flowState[I, O], error) {
+	var i I
+	var o O
+	metadata := map[string]any{
+		"inputSchema":  inferJSONSchema(i),
+		"outputSchema": inferJSONSchema(o),
+	}
+	return NewStreamingAction(f.name, metadata, func(ctx context.Context, inst *flowInstruction[I], cb StreamingCallback[S]) (*flowState[I, O], error) {
 		spanMetaKey.fromContext(ctx).SetAttr("flow:wrapperAction", "true")
 		return f.runInstruction(ctx, inst, cb)
 	})
diff --git a/go/samples/flow-sample1/main.go b/go/samples/flow-sample1/main.go
index 4d00d3e40c..215d10130a 100644
--- a/go/samples/flow-sample1/main.go
+++ b/go/samples/flow-sample1/main.go
@@ -26,6 +26,7 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"strconv"
 
 	"github.com/firebase/genkit/go/genkit"
 )
@@ -43,6 +44,19 @@ func main() {
 		return genkit.RunFlow(ctx, basic, "foo")
 	})
 
+	type complex struct {
+		Key   string `json:"key"`
+		Value int    `json:"value"`
+	}
+
+	genkit.DefineFlow("complex", func(ctx context.Context, c complex, _ genkit.NoStream) (string, error) {
+		foo, err := genkit.Run(ctx, "call-llm", func() (string, error) { return c.Key + ": " + strconv.Itoa(c.Value), nil })
+		if err != nil {
+			return "", err
+		}
+		return foo, nil
+	})
+
 	type chunk struct {
 		Count int `json:"count"`
 	}