feat(go/genkit/prompt): add support for code defined schemas

go/ai/option.go

@@ -877,7 +877,8 @@ type promptExecutionOptions struct {
 	commonGenOptions
 	executionOptions
 	documentOptions
-	Input any // Input fields for the prompt. If not nil this should be a struct that matches the prompt's input schema.
+	Input  any // Input fields for the prompt. If not nil this should be a struct that matches the prompt's input schema.
+	Output any // Output fields for the prompt. If not nil, this should be a struct that matches the prompt's desired output schema
 }
 
 // PromptExecuteOption is an option for executing a prompt. It applies only to [prompt.Execute].
@@ -914,3 +915,9 @@ func (o *promptExecutionOptions) applyPromptExecute(pgOpts *promptExecutionOptio
 func WithInput(input any) PromptExecuteOption {
 	return &promptExecutionOptions{Input: input}
 }
+
+// WithOutput sets the output for the prompt desired output. Output must conform to the
+// prompt's output schema and should be a map[string]any or a struct of the same type.
+func WithOutput(output any) PromptExecuteOption {
+	return &promptExecutionOptions{Output: output}
+}

go/ai/prompt.go

@@ -50,6 +50,18 @@ type prompt struct {
 	registry *registry.Registry
 }
 
+// SchemaRef is a reference to a prompt schema
+type SchemaRef interface {
+	Name() string
+}
+
+type SchemaName string
+
+// Name returns the name of the prompt schema
+func (s SchemaName) Name() string {
+	return (string)(s)
+}
+
 // DefinePrompt creates a new [Prompt] and registers it.
 func DefinePrompt(r *registry.Registry, name string, opts ...PromptOption) Prompt {
 	if name == "" {
@@ -120,6 +132,12 @@ func LookupPrompt(r *registry.Registry, name string) Prompt {
 	}
 }
 
+// DefineSchema defines a schema in the registry
+// It panics if an error was encountered
+func DefineSchema(r *registry.Registry, name string, schema any) {
+	// TODO: marshal the schema and store it into the registry
+}
+
 // Execute renders a prompt, does variable substitution and
 // passes the rendered template to the AI model specified by the prompt.
 func (p *prompt) Execute(ctx context.Context, opts ...PromptExecuteOption) (*ModelResponse, error) {
@@ -184,6 +202,8 @@ func (p *prompt) Render(ctx context.Context, input any) (*GenerateActionOptions,
 		input = p.Desc().Metadata["prompt"].(map[string]any)["defaultInput"]
 	}
 
+	// TODO: should we consider output atp?
+
 	return p.Run(ctx, input, nil)
 }
 
@@ -543,6 +563,8 @@ func LoadPrompt(r *registry.Registry, dir, filename, namespace string) Prompt {
 		toolRefs[i] = ToolName(tool)
 	}
 
+	// TODO: get input/output schemas from metadata
+
 	promptMetadata := map[string]any{
 		"template": parsedPrompt.Template,
 	}

go/genkit/genkit.go

@@ -631,6 +631,11 @@ func LookupPrompt(g *Genkit, name string) ai.Prompt {
 	return ai.LookupPrompt(g.reg, name)
 }
 
+// DefineSchema defines a prompt schema programmatically
+func DefineSchema(g *Genkit, name string, schema any) {
+	return ai.DefineSchema(g.reg, name, schema)
+}
+
 // GenerateWithRequest performs a model generation request using explicitly provided
 // [ai.GenerateActionOptions]. This function is typically used in conjunction with
 // prompts defined via [DefinePrompt], where [ai.prompt.Render] produces the

go/samples/prompts/main.go

@@ -50,6 +50,7 @@ func main() {
 	PromptWithFunctions(ctx, g)
 	PromptWithOutputTypeDotprompt(ctx, g)
 	PromptWithMediaType(ctx, g)
+	PromptWithSchema(ctx, g)
 
 	mux := http.NewServeMux()
 	for _, a := range genkit.ListFlows(g) {
@@ -328,6 +329,45 @@ func PromptWithMediaType(ctx context.Context, g *genkit.Genkit) {
 	fmt.Println(resp.Text())
 }
 
+/*
+Inline schema equivalence
+
+	  title: string, recipe title
+	  ingredients(array):
+		name: string
+		quantity: string
+	  steps(array, the steps required to complete the recipe): string
+*/
+type Ingredient struct {
+	Name     string `json:"name" description:"ingredient name"`
+	Quantity string `json:"quantity" description:"ingredient quantity"`
+}
+
+type RecipeSchema struct {
+	Title       string       `json:"title" description:"Recipe name"`
+	Ingredients []Ingredient `json:"ingredients" description:"Recipe ingredients"`
+	Steps       []string     `json:"steps" description:"Recipe steps"`
+}
+
+func PromptWithSchema(ctx context.Context, g *genkit.Genkit) {
+	// prompt schemas can be referenced at any time
+	genkit.DefineSchema(g, "recipe", RecipeSchema{})
+
+	prompt := genkit.LoadPrompt(g, "./prompts/recipe.prompt", "recipes")
+	if prompt == nil {
+		log.Fatal("empty prompt")
+	}
+
+	resp, err := prompt.Execute(ctx,
+		ai.WithModelName("vertexai/gemini-2.0-flash"),
+		ai.WithOutput(RecipeSchema{}),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+	fmt.Println(resp.Text())
+}
+
 func fetchImgAsBase64() (string, error) {
 	imgUrl := "https://pd.w.org/2025/07/58268765f177911d4.13750400-2048x1365.jpg"
 	resp, err := http.Get(imgUrl)

go/samples/prompts/prompts/recipe.prompt

@@ -0,0 +1,18 @@
+---
+model: vertexai/gemini-2.0-flash
+input:
+  schema:
+    food: string
+    ingredients?(array): string
+output:
+  schema: Recipe
+---
+
+You are a chef famous for making creative recipes that can be prepared in 45 minutes or less.
+
+Generate a recipe for {{food}}.
+
+{{#if ingredients}}
+Make sure to include the following ingredients:
+{{list ingredients}}
+{{/if}}