Use string format for jsonSchema (#73)

Author: sgoedecke

cmd/eval/eval_test.go

@@ -569,21 +569,7 @@ name: JSON Schema Evaluation
 description: Testing responseFormat and jsonSchema in eval
 model: openai/gpt-4o
 responseFormat: json_schema
-jsonSchema:
-  name: response_schema
-  strict: true
-  schema:
-    type: object
-    properties:
-      message:
-        type: string
-        description: The response message
-      confidence:
-        type: number
-        description: Confidence score
-    required:
-      - message
-    additionalProperties: false
+jsonSchema: '{"name": "response_schema", "strict": true, "schema": {"type": "object", "properties": {"message": {"type": "string", "description": "The response message"}, "confidence": {"type": "number", "description": "Confidence score"}}, "required": ["message"], "additionalProperties": false}}'
 testData:
   - input: "hello"
     expected: "hello world"

cmd/run/run_test.go

@@ -341,22 +341,7 @@ name: JSON Schema Test
 description: Test responseFormat and jsonSchema
 model: openai/test-model
 responseFormat: json_schema
-jsonSchema:
-  name: person_schema
-  strict: true
-  schema:
-    type: object
-    properties:
-      name:
-        type: string
-        description: The name
-      age:
-        type: integer
-        description: The age
-    required:
-      - name
-      - age
-    additionalProperties: false
+jsonSchema: '{"name": "person_schema", "strict": true, "schema": {"type": "object", "properties": {"name": {"type": "string", "description": "The name"}, "age": {"type": "integer", "description": "The age"}}, "required": ["name", "age"], "additionalProperties": false}}'
 messages:
   - role: system
     content: You are a helpful assistant.

examples/json_schema_prompt.yml

@@ -1,64 +1,52 @@
 name: JSON Schema Response Example
 description: Example prompt demonstrating responseFormat and jsonSchema usage
-model: openai/gpt-4o
+model: openai/gpt-4o-mini
 responseFormat: json_schema
-jsonSchema:
-  name: Person Information Schema
-  strict: true
-  schema:
-    type: object
-    description: A structured response containing person information
-    properties:
-      name:
-        type: string
-        description: The full name of the person
-      age:
-        type: integer
-        description: The age of the person in years
-        minimum: 0
-        maximum: 150
-      email:
-        type: string
-        description: The email address of the person
-        format: email
-      skills:
-        type: array
-        description: A list of skills the person has
-        items:
-          type: string
-      address:
-        type: object
-        description: The person's address
-        properties:
-          street:
-            type: string
-            description: Street address
-          city:
-            type: string
-            description: City name
-          country:
-            type: string
-            description: Country name
-        required:
-          - city
-          - country
-    required:
-      - name
-      - age
+jsonSchema: |-
+  {
+    "name": "animal_description",
+    "strict": true,
+    "schema": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string",
+          "description": "The name of the animal"
+        },
+        "habitat": {
+          "type": "string",
+          "description": "The habitat where the animal lives"
+        },
+        "diet": {
+          "type": "string",
+          "description": "What the animal eats",
+          "enum": ["carnivore", "herbivore", "omnivore"]
+        },
+        "characteristics": {
+          "type": "array",
+          "description": "Key characteristics of the animal",
+          "items": {
+            "type": "string"
+          }
+        }
+      },
+      "required": ["name", "habitat", "diet"],
+      "additionalProperties": false
+    }
+  }
 messages:
   - role: system
-    content: You are a helpful assistant that provides structured information about people.
+    content: You are a helpful assistant that provides detailed information about animals.
   - role: user
-    content: "Generate information for a person named {{name}} who is {{age}} years old."
+    content: "Describe a {{animal}} in detail."
 testData:
-  - name: "Alice Johnson"
-    age: "30"
-  - name: "Bob Smith"
-    age: "25"
+  - animal: "dog"
+  - animal: "cat"
+  - animal: "elephant"
 evaluators:
-  - name: has-required-fields
+  - name: has-name
     string:
       contains: "name"
-  - name: valid-json-structure
+  - name: has-habitat
     string:
-      contains: "age"
+      contains: "habitat"

pkg/prompt/prompt.go

@@ -2,6 +2,7 @@
 package prompt
 
 import (
+	"encoding/json"
 	"fmt"
 	"os"
 	"strings"
@@ -69,9 +70,32 @@ type Choice struct {
 
 // JsonSchema represents a JSON schema for structured responses
 type JsonSchema struct {
-	Name   string                 `yaml:"name" json:"name"`
-	Strict *bool                  `yaml:"strict,omitempty" json:"strict,omitempty"`
-	Schema map[string]interface{} `yaml:"schema" json:"schema"`
+	Raw    string
+	Parsed map[string]interface{}
+}
+
+// UnmarshalYAML implements custom YAML unmarshaling for JsonSchema
+// Only supports JSON string format
+func (js *JsonSchema) UnmarshalYAML(node *yaml.Node) error {
+	// Only support string nodes (JSON format)
+	if node.Kind != yaml.ScalarNode {
+		return fmt.Errorf("jsonSchema must be a JSON string")
+	}
+
+	var jsonStr string
+	if err := node.Decode(&jsonStr); err != nil {
+		return err
+	}
+
+	// Parse and validate the JSON schema
+	var parsed map[string]interface{}
+	if err := json.Unmarshal([]byte(jsonStr), &parsed); err != nil {
+		return fmt.Errorf("invalid JSON in jsonSchema: %w", err)
+	}
+
+	js.Raw = jsonStr
+	js.Parsed = parsed
+	return nil
 }
 
 // LoadFromFile loads and parses a prompt file from the given path
@@ -105,16 +129,18 @@ func (f *File) validateResponseFormat() error {
 		return fmt.Errorf("invalid responseFormat: %s. Must be 'text', 'json_object', or 'json_schema'", *f.ResponseFormat)
 	}
 
-	// If responseFormat is "json_schema", jsonSchema must be provided with required fields
+	// If responseFormat is "json_schema", jsonSchema must be provided
 	if *f.ResponseFormat == "json_schema" {
 		if f.JsonSchema == nil {
 			return fmt.Errorf("jsonSchema is required when responseFormat is 'json_schema'")
 		}
-		if f.JsonSchema.Name == "" {
-			return fmt.Errorf("jsonSchema.name is required when responseFormat is 'json_schema'")
+
+		// Check for required fields in the already parsed schema
+		if _, ok := f.JsonSchema.Parsed["name"]; !ok {
+			return fmt.Errorf("jsonSchema must contain 'name' field")
 		}
-		if f.JsonSchema.Schema == nil {
-			return fmt.Errorf("jsonSchema.schema is required when responseFormat is 'json_schema'")
+		if _, ok := f.JsonSchema.Parsed["schema"]; !ok {
+			return fmt.Errorf("jsonSchema must contain 'schema' field")
 		}
 	}
 
@@ -176,7 +202,6 @@ func (f *File) BuildChatCompletionOptions(messages []azuremodels.ChatMessage) az
 		Stream:   false,
 	}
 
-	// Apply model parameters
 	if f.ModelParameters.MaxTokens != nil {
 		req.MaxTokens = f.ModelParameters.MaxTokens
 	}
@@ -187,20 +212,12 @@ func (f *File) BuildChatCompletionOptions(messages []azuremodels.ChatMessage) az
 		req.TopP = f.ModelParameters.TopP
 	}
 
-	// Apply response format
 	if f.ResponseFormat != nil {
 		responseFormat := &azuremodels.ResponseFormat{
 			Type: *f.ResponseFormat,
 		}
 		if f.JsonSchema != nil {
-			// Convert JsonSchema to map[string]interface{}
-			schemaMap := make(map[string]interface{})
-			schemaMap["name"] = f.JsonSchema.Name
-			if f.JsonSchema.Strict != nil {
-				schemaMap["strict"] = *f.JsonSchema.Strict
-			}
-			schemaMap["schema"] = f.JsonSchema.Schema
-			responseFormat.JsonSchema = &schemaMap
+			responseFormat.JsonSchema = &f.JsonSchema.Parsed
 		}
 		req.ResponseFormat = responseFormat
 	}

pkg/prompt/prompt_test.go

@@ -1,6 +1,7 @@
 package prompt
 
 import (
+	"encoding/json"
 	"os"
 	"path/filepath"
 	"testing"
@@ -139,27 +140,35 @@ messages:
 		require.Nil(t, promptFile.JsonSchema)
 	})
 
-	t.Run("loads prompt file with responseFormat json_schema and jsonSchema", func(t *testing.T) {
+	t.Run("loads prompt file with responseFormat json_schema and jsonSchema as JSON string", func(t *testing.T) {
 		const yamlBody = `
-name: JSON Schema Response Format Test
-description: Test with JSON schema response format
+name: JSON Schema String Format Test
+description: Test with JSON schema as JSON string
 model: openai/gpt-4o
 responseFormat: json_schema
-jsonSchema:
-  name: person_info
-  strict: true
-  schema:
-    type: object
-    properties:
-      name:
-        type: string
-        description: The name of the person
-      age:
-        type: integer
-        description: The age of the person
-    required:
-      - name
-    additionalProperties: false
+jsonSchema: |-
+  {
+    "name": "describe_animal",
+    "strict": true,
+    "schema": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string",
+          "description": "The name of the animal"
+        },
+        "habitat": {
+          "type": "string",
+          "description": "The habitat the animal lives in"
+        }
+      },
+      "additionalProperties": false,
+      "required": [
+        "name",
+        "habitat"
+      ]
+    }
+  }
 messages:
   - role: user
     content: "Hello"
@@ -175,10 +184,26 @@ messages:
 		require.NotNil(t, promptFile.ResponseFormat)
 		require.Equal(t, "json_schema", *promptFile.ResponseFormat)
 		require.NotNil(t, promptFile.JsonSchema)
-		require.Equal(t, "person_info", promptFile.JsonSchema.Name)
-		require.True(t, *promptFile.JsonSchema.Strict)
-		require.Contains(t, promptFile.JsonSchema.Schema, "type")
-		require.Contains(t, promptFile.JsonSchema.Schema, "properties")
+
+		// Verify the schema contents using the already parsed data
+		schema := promptFile.JsonSchema.Parsed
+		require.Equal(t, "describe_animal", schema["name"])
+		require.Equal(t, true, schema["strict"])
+		require.Contains(t, schema, "schema")
+
+		// Verify the nested schema structure
+		nestedSchema := schema["schema"].(map[string]interface{})
+		require.Equal(t, "object", nestedSchema["type"])
+		require.Contains(t, nestedSchema, "properties")
+		require.Contains(t, nestedSchema, "required")
+
+		properties := nestedSchema["properties"].(map[string]interface{})
+		require.Contains(t, properties, "name")
+		require.Contains(t, properties, "habitat")
+
+		required := nestedSchema["required"].([]interface{})
+		require.Contains(t, required, "name")
+		require.Contains(t, required, "habitat")
 	})
 
 	t.Run("validates invalid responseFormat", func(t *testing.T) {
@@ -224,23 +249,32 @@ messages:
 	})
 
 	t.Run("BuildChatCompletionOptions includes responseFormat and jsonSchema", func(t *testing.T) {
+		jsonSchemaStr := `{
+			"name": "test_schema",
+			"strict": true,
+			"schema": {
+				"type": "object",
+				"properties": {
+					"name": {
+						"type": "string",
+						"description": "The name"
+					}
+				},
+				"required": ["name"]
+			}
+		}`
+
 		promptFile := &File{
 			Model:          "openai/gpt-4o",
 			ResponseFormat: func() *string { s := "json_schema"; return &s }(),
-			JsonSchema: &JsonSchema{
-				Name:   "test_schema",
-				Strict: func() *bool { b := true; return &b }(),
-				Schema: map[string]interface{}{
-					"type": "object",
-					"properties": map[string]interface{}{
-						"name": map[string]interface{}{
-							"type":        "string",
-							"description": "The name",
-						},
-					},
-					"required": []string{"name"},
-				},
-			},
+			JsonSchema: func() *JsonSchema {
+				js := &JsonSchema{Raw: jsonSchemaStr}
+				err := json.Unmarshal([]byte(jsonSchemaStr), &js.Parsed)
+				if err != nil {
+					t.Fatal(err)
+				}
+				return js
+			}(),
 		}
 
 		messages := []azuremodels.ChatMessage{