[Go] Make dotprompt take in jsonschema.Schema instead of `map[strin…

…g]any` for better ergonomics. (#645)
firebase · Jul 18, 2024 · 3318226 · 3318226
1 parent a7ad569
commit 3318226
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 25 deletions.
diff --git a/go/plugins/dotprompt/dotprompt.go b/go/plugins/dotprompt/dotprompt.go
@@ -98,14 +98,17 @@ type Config struct {
 	// Details for the model.
 	GenerationConfig *ai.GenerationCommonConfig
 
-	InputSchema      *jsonschema.Schema // schema for input variables
-	VariableDefaults map[string]any     // default input variable values
+	// Schema for input variables.
+	InputSchema *jsonschema.Schema
+
+	// Default input variable values
+	VariableDefaults map[string]any
 
 	// Desired output format.
 	OutputFormat ai.OutputFormat
 
 	// Desired output schema, for JSON output.
-	OutputSchema map[string]any // TODO: use *jsonschema.Schema
+	OutputSchema *jsonschema.Schema
 
 	// Arbitrary metadata.
 	Metadata map[string]any

diff --git a/go/plugins/dotprompt/dotprompt_test.go b/go/plugins/dotprompt/dotprompt_test.go
@@ -125,12 +125,8 @@ func TestPrompts(t *testing.T) {
 					t.Errorf("unexpected output schema: %v", prompt.OutputSchema)
 				}
 			} else {
-				var output map[string]any
-				if err := json.Unmarshal([]byte(test.output), &output); err != nil {
-					t.Fatalf("JSON unmarshal of %q failed: %v", test.output, err)
-				}
-				if diff := cmp.Diff(output, prompt.OutputSchema); diff != "" {
-					t.Errorf("output schema mismatch (-want, +got):\n%s", diff)
+				if diff := cmpSchema(t, prompt.OutputSchema, test.output); diff != "" {
+					t.Errorf("input schema mismatch (-want, +got):\n%s", diff)
 				}
 			}
 		})

diff --git a/go/plugins/dotprompt/genkit.go b/go/plugins/dotprompt/genkit.go
@@ -16,6 +16,7 @@ package dotprompt
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"reflect"
@@ -117,9 +118,21 @@ func (p *Prompt) buildRequest(ctx context.Context, input any) (*ai.GenerateReque
 
 	req.Config = p.GenerationConfig
 
+	var outputSchema map[string]any
+	if p.OutputSchema != nil {
+		jsonBytes, err := p.OutputSchema.MarshalJSON()
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal output schema JSON: %w", err)
+		}
+		err = json.Unmarshal(jsonBytes, &outputSchema)
+		if err != nil {
+			return nil, fmt.Errorf("failed to unmarshal output schema JSON: %w", err)
+		}
+	}
+
 	req.Output = &ai.GenerateRequestOutput{
 		Format: p.OutputFormat,
-		Schema: p.OutputSchema,
+		Schema: outputSchema,
 	}
 
 	req.Tools = p.Tools

diff --git a/go/samples/coffee-shop/main.go b/go/samples/coffee-shop/main.go
@@ -36,7 +36,6 @@ package main
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"log"
 
@@ -172,24 +171,12 @@ func main() {
 		return text, nil
 	})
 
-	schema := r.Reflect(simpleGreetingOutput{})
-	jsonBytes, err := schema.MarshalJSON()
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	var outputSchema map[string]any
-	err = json.Unmarshal(jsonBytes, &outputSchema)
-	if err != nil {
-		log.Fatal(err)
-	}
-
 	simpleStructuredGreetingPrompt, err := dotprompt.Define("simpleStructuredGreeting", simpleStructuredGreetingPromptTemplate,
 		dotprompt.Config{
 			Model:        g,
-			InputSchema:  jsonschema.Reflect(simpleGreetingInput{}),
+			InputSchema:  r.Reflect(simpleGreetingInput{}),
 			OutputFormat: ai.OutputFormatJSON,
-			OutputSchema: outputSchema,
+			OutputSchema: r.Reflect(simpleGreetingOutput{}),
 		},
 	)
 	if err != nil {