fix(go): fixed bad dotprompt output format parsing (#4109)

Author: apascal07

go/ai/prompt.go

@@ -776,7 +776,11 @@ func LoadPromptFromSource(r api.Registry, source, name, namespace string) (Promp
 	}
 
 	if inputSchema, ok := metadata.Input.Schema.(map[string]any); ok {
-		opts.InputSchema = inputSchema
+		if ref, ok := inputSchema["$ref"].(string); ok {
+			opts.InputSchema = core.SchemaRef(ref)
+		} else {
+			opts.InputSchema = inputSchema
+		}
 	}
 
 	if metadata.Output.Format != "" {
@@ -794,6 +798,17 @@ func LoadPromptFromSource(r api.Registry, source, name, namespace string) (Promp
 		}
 	}
 
+	if outputSchema, ok := metadata.Output.Schema.(map[string]any); ok {
+		if ref, ok := outputSchema["$ref"].(string); ok {
+			opts.OutputSchema = core.SchemaRef(ref)
+		} else {
+			opts.OutputSchema = outputSchema
+		}
+		if opts.OutputFormat == "" {
+			opts.OutputFormat = OutputFormatJSON
+		}
+	}
+
 	key := promptKey(name, variant, namespace)
 
 	prompt := DefinePrompt(r, key, opts, WithPrompt(parsedPrompt.Template))

go/ai/prompt_test.go

@@ -1368,6 +1368,68 @@ Simple prompt
 			t.Fatal("Prompt 'simple' was not registered")
 		}
 	})
+
+	t.Run("prompt with inline output schema", func(t *testing.T) {
+		reg := registry.New()
+		ConfigureFormats(reg)
+
+		source := `---
+model: test/chat
+output:
+  format: json
+  schema:
+    type: object
+    properties:
+      title:
+        type: string
+      description:
+        type: string
+    required:
+      - title
+      - description
+---
+Generate something
+`
+		prompt, err := LoadPromptFromSource(reg, source, "outputSchemaPrompt", "")
+		if err != nil {
+			t.Fatalf("LoadPromptFromRaw failed: %v", err)
+		}
+		if prompt == nil {
+			t.Fatal("LoadPromptFromRaw returned nil prompt")
+		}
+
+		actionOpts, err := prompt.Render(context.Background(), nil)
+		if err != nil {
+			t.Fatalf("Render failed: %v", err)
+		}
+
+		// Verify that the output config is set correctly
+		if actionOpts.Output == nil {
+			t.Fatal("Expected Output config to be set")
+		}
+		if actionOpts.Output.Format != OutputFormatJSON {
+			t.Errorf("Expected output format 'json', got %q", actionOpts.Output.Format)
+		}
+		if actionOpts.Output.JsonSchema == nil {
+			t.Fatal("Expected output JsonSchema to be set for inline schema")
+		}
+
+		// Verify the schema structure
+		schema := actionOpts.Output.JsonSchema
+		if schema["type"] != "object" {
+			t.Errorf("Expected schema type 'object', got %v", schema["type"])
+		}
+		properties, ok := schema["properties"].(map[string]any)
+		if !ok {
+			t.Fatal("Expected schema properties to be a map")
+		}
+		if _, ok := properties["title"]; !ok {
+			t.Error("Expected schema to have 'title' property")
+		}
+		if _, ok := properties["description"]; !ok {
+			t.Error("Expected schema to have 'description' property")
+		}
+	})
 }
 
 // TestDefinePartialAndHelperJourney demonstrates a complete user journey for defining

go/internal/base/json.go

@@ -138,18 +138,29 @@ func SchemaAsMap(s *jsonschema.Schema) map[string]any {
 	return m
 }
 
-// jsonMarkdownRegex specifically looks for "json" language identifier
-var jsonMarkdownRegex = regexp.MustCompile("(?s)```json(.*?)```")
+// jsonMarkdownRegex matches fenced code blocks with "json" language identifier (case-insensitive).
+var jsonMarkdownRegex = regexp.MustCompile("(?si)```json\\s*(.*?)```")
+
+// plainMarkdownRegex matches fenced code blocks without any language identifier.
+var plainMarkdownRegex = regexp.MustCompile("(?s)```\\s*\\n(.*?)```")
 
 // ExtractJSONFromMarkdown returns the contents of the first fenced code block in
-// the markdown text md. If there is none, it returns md.
+// the markdown text md. It matches code blocks with "json" identifier (case-insensitive)
+// or code blocks without any language identifier. If there is no matching block, it returns md.
 func ExtractJSONFromMarkdown(md string) string {
+	// First try to match explicit json code blocks
 	matches := jsonMarkdownRegex.FindStringSubmatch(md)
-	if len(matches) < 2 {
-		return md
+	if len(matches) >= 2 {
+		return strings.TrimSpace(matches[1])
+	}
+
+	// Fall back to plain code blocks (no language identifier)
+	matches = plainMarkdownRegex.FindStringSubmatch(md)
+	if len(matches) >= 2 {
+		return strings.TrimSpace(matches[1])
 	}
-	// capture group 1 matches the actual fenced JSON block
-	return strings.TrimSpace(matches[1])
+
+	return md
 }
 
 // GetJSONObjectLines splits a string by newlines, trims whitespace from each line,

go/internal/base/json_test.go

@@ -78,6 +78,31 @@ func TestExtractJSONFromMarkdown(t *testing.T) {
 			in:   "```json\n{\"a\": 1}\n``` ```yaml\nkey: 1\nanother-key: 2```",
 			want: "{\"a\": 1}",
 		},
+		{
+			desc: "uppercase JSON identifier",
+			in:   "```JSON\n{\"a\": 1}\n```",
+			want: "{\"a\": 1}",
+		},
+		{
+			desc: "mixed case Json identifier",
+			in:   "```Json\n{\"a\": 1}\n```",
+			want: "{\"a\": 1}",
+		},
+		{
+			desc: "plain code block without identifier",
+			in:   "```\n{\"a\": 1}\n```",
+			want: "{\"a\": 1}",
+		},
+		{
+			desc: "plain code block with text before",
+			in:   "Here is the result:\n\n```\n{\"title\": \"Pizza\"}\n```",
+			want: "{\"title\": \"Pizza\"}",
+		},
+		{
+			desc: "json block preferred over plain block",
+			in:   "```\n{\"plain\": true}\n``` then ```json\n{\"json\": true}\n```",
+			want: "{\"json\": true}",
+		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {