Go: Support for streaming formatters

go/ai/format_array.go

@@ -15,6 +15,7 @@
 package ai
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -68,6 +69,12 @@ func (a arrayHandler) Config() ModelOutputConfig {
 	return a.config
 }
 
+func (a arrayHandler) StreamCallback(cb ModelStreamCallback) ModelStreamCallback {
+	return func(ctx context.Context, mrc *ModelResponseChunk) error {
+		return cb(ctx, mrc)
+	}
+}
+
 // ParseMessage parses the message and returns the formatted message.
 func (a arrayHandler) ParseMessage(m *Message) (*Message, error) {
 	if a.config.Format == OutputFormatArray {
@@ -103,3 +110,8 @@ func (a arrayHandler) ParseMessage(m *Message) (*Message, error) {
 
 	return m, nil
 }
+
+// ParseChunk parse the chunk and returns a new formatted chunk.
+func (a arrayHandler) ParseChunk(c *ModelResponseChunk) (*ModelResponseChunk, error) {
+	return c, nil
+}

go/ai/format_enum.go

@@ -15,6 +15,7 @@
 package ai
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"regexp"
@@ -67,6 +68,12 @@ func (e enumHandler) Config() ModelOutputConfig {
 	return e.config
 }
 
+func (e enumHandler) StreamCallback(cb ModelStreamCallback) ModelStreamCallback {
+	return func(ctx context.Context, mrc *ModelResponseChunk) error {
+		return cb(ctx, mrc)
+	}
+}
+
 // ParseMessage parses the message and returns the formatted message.
 func (e enumHandler) ParseMessage(m *Message) (*Message, error) {
 	if e.config.Format == OutputFormatEnum {
@@ -100,6 +107,11 @@ func (e enumHandler) ParseMessage(m *Message) (*Message, error) {
 	return m, nil
 }
 
+// ParseChunk parse the chunk and returns a new formatted chunk.
+func (e enumHandler) ParseChunk(c *ModelResponseChunk) (*ModelResponseChunk, error) {
+	return c, nil
+}
+
 // Get enum strings from json schema
 func objectEnums(schema map[string]any) []string {
 	var enums []string

go/ai/format_json.go

@@ -15,10 +15,13 @@
 package ai
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 
+	partialparser "github.com/blaze2305/partial-json-parser"
+	"github.com/blaze2305/partial-json-parser/options"
 	"github.com/firebase/genkit/go/internal/base"
 )
 
@@ -55,8 +58,9 @@ func (j jsonFormatter) Handler(schema map[string]any) (FormatHandler, error) {
 
 // jsonHandler is a handler for the JSON formatter.
 type jsonHandler struct {
-	instructions string
-	config       ModelOutputConfig
+	instructions  string
+	config        ModelOutputConfig
+	previousParts []*Part
 }
 
 // Instructions returns the instructions for the formatter.
@@ -69,6 +73,21 @@ func (j jsonHandler) Config() ModelOutputConfig {
 	return j.config
 }
 
+// StreamCallback handler for streaming formatted responses
+func (j jsonHandler) StreamCallback(cb ModelStreamCallback) ModelStreamCallback {
+	return func(ctx context.Context, mrc *ModelResponseChunk) error {
+		j.previousParts = append(j.previousParts, mrc.Content...)
+		mrc.Content = j.previousParts
+
+		parsed, err := j.ParseChunk(mrc)
+		if err != nil {
+			return err
+		}
+
+		return cb(ctx, parsed)
+	}
+}
+
 // ParseMessage parses the message and returns the formatted message.
 func (j jsonHandler) ParseMessage(m *Message) (*Message, error) {
 	if j.config.Format == OutputFormatJSON {
@@ -85,7 +104,6 @@ func (j jsonHandler) ParseMessage(m *Message) (*Message, error) {
 			}
 
 			text := base.ExtractJSONFromMarkdown(part.Text)
-
 			if j.config.Schema != nil {
 				var schemaBytes []byte
 				schemaBytes, err := json.Marshal(j.config.Schema)
@@ -107,3 +125,34 @@ func (j jsonHandler) ParseMessage(m *Message) (*Message, error) {
 
 	return m, nil
 }
+
+// ParseChunk parse the chunk and returns a new formatted chunk.
+func (j jsonHandler) ParseChunk(c *ModelResponseChunk) (*ModelResponseChunk, error) {
+	if j.config.Format == OutputFormatJSON {
+		if c == nil {
+			return nil, errors.New("chunk is empty")
+		}
+
+		if len(c.Content) == 0 {
+			return nil, errors.New("message has no content")
+		}
+
+		// Get all chunks streamed so far
+		text := c.Text()
+		text = base.ExtractJSONFromMarkdown(text)
+		// Try and extract a json object
+		text = base.GetJsonObject(text)
+		if text != "" {
+			var err error
+			text, err = partialparser.ParseMalformedString(text, options.ALL, false)
+			if err != nil {
+				return nil, errors.New("message is not a valid JSON")
+			}
+		} else {
+			return nil, nil
+		}
+
+		c.Content = []*Part{NewJSONPart(text)}
+	}
+	return c, nil
+}

go/ai/format_jsonl.go

@@ -15,10 +15,14 @@
 package ai
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"strings"
 
+	partialparser "github.com/blaze2305/partial-json-parser"
+	"github.com/blaze2305/partial-json-parser/options"
 	"github.com/firebase/genkit/go/internal/base"
 )
 
@@ -55,8 +59,9 @@ func (j jsonlFormatter) Handler(schema map[string]any) (FormatHandler, error) {
 }
 
 type jsonlHandler struct {
-	instructions string
-	config       ModelOutputConfig
+	instructions  string
+	config        ModelOutputConfig
+	previousParts []*Part
 }
 
 // Instructions returns the instructions for the formatter.
@@ -69,6 +74,20 @@ func (j jsonlHandler) Config() ModelOutputConfig {
 	return j.config
 }
 
+func (j jsonlHandler) StreamCallback(cb ModelStreamCallback) ModelStreamCallback {
+	return func(ctx context.Context, mrc *ModelResponseChunk) error {
+		j.previousParts = append(j.previousParts, mrc.Content...)
+		mrc.Content = j.previousParts
+
+		parsed, err := j.ParseChunk(mrc)
+		if err != nil {
+			return err
+		}
+
+		return cb(ctx, parsed)
+	}
+}
+
 // ParseMessage parses the message and returns the formatted message.
 func (j jsonlHandler) ParseMessage(m *Message) (*Message, error) {
 	if j.config.Format == OutputFormatJSONL {
@@ -106,3 +125,58 @@ func (j jsonlHandler) ParseMessage(m *Message) (*Message, error) {
 
 	return m, nil
 }
+
+// ParseChunk parse the chunk and returns a new formatted chunk.
+func (j jsonlHandler) ParseChunk(c *ModelResponseChunk) (*ModelResponseChunk, error) {
+	if j.config.Format == OutputFormatJSONL {
+		if c == nil {
+			return nil, errors.New("message is empty")
+		}
+		if len(c.Content) == 0 {
+			return nil, errors.New("message has no content")
+		}
+
+		// Get all chunks streamed so far
+		text := c.Text()
+
+		startIndex := 0
+		// If there are previous chunks, adjust startIndex based on the last newline
+		// in the previous text to ensure complete lines are processed from the accumulatedText.
+		noParts := len(c.Content)
+		if c.Content != nil && noParts > 1 {
+			var sb strings.Builder
+			i := 0
+			for i < noParts-1 {
+				sb.WriteString(c.Content[i].Text)
+				i++
+			}
+
+			previousText := sb.String()
+			lastNewline := strings.LastIndex(previousText, `\n`)
+
+			if lastNewline != -1 {
+				// Exclude the newline
+				startIndex = lastNewline + 2
+			}
+		}
+
+		text = text[startIndex:]
+
+		var newParts []*Part
+		lines := base.GetJsonObjectLines(text)
+		for _, line := range lines {
+			if line != "" {
+				var err error
+				line, err = partialparser.ParseMalformedString(line, options.ALL, false)
+				if err != nil {
+					return nil, errors.New("message is not a valid JSON")
+				}
+
+				newParts = append(newParts, NewJSONPart(line))
+			}
+		}
+
+		c.Content = newParts
+	}
+	return c, nil
+}

go/ai/format_text.go

@@ -14,6 +14,10 @@
 
 package ai
 
+import (
+	"context"
+)
+
 type textFormatter struct{}
 
 // Name returns the name of the formatter.
@@ -47,7 +51,18 @@ func (t textHandler) Instructions() string {
 	return t.instructions
 }
 
+func (t textHandler) StreamCallback(cb ModelStreamCallback) ModelStreamCallback {
+	return func(ctx context.Context, mrc *ModelResponseChunk) error {
+		return cb(ctx, mrc)
+	}
+}
+
 // ParseMessage parses the message and returns the formatted message.
 func (t textHandler) ParseMessage(m *Message) (*Message, error) {
 	return m, nil
 }
+
+// ParseChunk parse the chunk and returns a new formatted chunk.
+func (t textHandler) ParseChunk(c *ModelResponseChunk) (*ModelResponseChunk, error) {
+	return c, nil
+}

go/ai/formatter.go

@@ -50,10 +50,14 @@ type Formatter interface {
 type FormatHandler interface {
 	// ParseMessage parses the message and returns a new formatted message.
 	ParseMessage(message *Message) (*Message, error)
+	// ParseChunk parse the chunk and returns a new formatted chunk.
+	ParseChunk(chunk *ModelResponseChunk) (*ModelResponseChunk, error)
 	// Instructions returns the formatter instructions to embed in the prompt.
 	Instructions() string
 	// Config returns the output config for the model request.
 	Config() ModelOutputConfig
+	// Stream callback returns a ModelStreamCallback
+	StreamCallback(cb ModelStreamCallback) ModelStreamCallback
 }
 
 // ConfigureFormats registers default formats in the registry