Implementing Low-Latency Structured JSON Outputs from LLMs with Finite State Machine Guided Decoding in Go — code snippets

snippet-1.go

package main

import (
	"fmt"
	"reflect"
	"strings"
)

// Compiler handles the conversion of Go structs into GBNF (GGML Backus-Naur Form) grammars.
type Compiler struct {
	rules map[string]string
}

// NewCompiler initializes a new GBNF compiler.
func NewCompiler() *Compiler {
	return &Compiler{
		rules: make(map[string]string),
	}
}

// Compile takes a Go struct and returns the complete GBNF grammar as a string.
func (c *Compiler) Compile(v interface{}) (string, error) {
	t := reflect.TypeOf(v)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	if t.Kind() != reflect.Struct {
		return "", fmt.Errorf("only struct types are supported, got %s", t.Kind())
	}

	c.rules = make(map[string]string)
	c.rules["ws"] = `ws ::= ([ \t\n\r]+)?`
	c.rules["string"] = `string ::= "\"" ([^\"\\] | "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]))* "\""`
	c.rules["number"] = `number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [+-]? [0-9]+)?`
	c.rules["boolean"] = `boolean ::= "true" | "false"`

	rootRule, err := c.compileStruct(t)
	if err != nil {
		return "", err
	}

	c.rules["root"] = fmt.Sprintf("root ::= %s", rootRule)

	var sb strings.Builder
	// Ensure root is written first
	sb.WriteString(c.rules["root"] + "\n")
	for name, rule := range c.rules {
		if name == "root" || name == "ws" || name == "string" || name == "number" || name == "boolean" {
			continue
		}
		sb.WriteString(rule + "\n")
	}
	sb.WriteString(c.rules["ws"] + "\n")
	sb.WriteString(c.rules["string"] + "\n")
	sb.WriteString(c.rules["number"] + "\n")
	sb.WriteString(c.rules["boolean"] + "\n")

	return sb.String(), nil
}

func (c *Compiler) compileStruct(t reflect.Type) (string, error) {
	structName := strings.ToLower(t.Name())
	if structName == "" {
		structName = "anonstruct"
	}

	var fields []string
	for i := 0; i < t.NumField(); i++ {
		field := t.Field(i)
		jsonTag := field.Tag.Get("json")
		if jsonTag == "-" {
			continue
		}

		fieldName := jsonTag
		if idx := strings.Index(jsonTag, ","); idx != -1 {
			fieldName = jsonTag[:idx]
		}
		if fieldName == "" {
			fieldName = field.Name
		}

		fieldRuleName := fmt.Sprintf("%s-%s", structName, strings.ToLower(field.Name))
		fieldGbnf, err := c.compileFieldType(field.Type, fieldRuleName)
		if err != nil {
			return "", err
		}

		// Field format: "field_name" ws ":" ws field_value
		fields = append(fields, fmt.Sprintf(`"\"%s\"" ws ":" ws %s`, fieldName, fieldGbnf))
	}

	// Build GBNF rule for the object structure: "{" ws field1 ws "," ws field2 ... ws "}"
	var ruleBuilder strings.Builder
	ruleBuilder.WriteString(`"{" ws `)
	for i, f := range fields {
		if i > 0 {
			ruleBuilder.WriteString(`"," ws `)
		}
		ruleBuilder.WriteString(f)
		ruleBuilder.WriteString(" ws ")
	}
	ruleBuilder.WriteString(`"}"`)

	ruleName := fmt.Sprintf("%s-obj", structName)
	c.rules[ruleName] = fmt.Sprintf("%s ::= %s", ruleName, ruleBuilder.String())
	return ruleName, nil
}

func (c *Compiler) compileFieldType(t reflect.Type, ruleName string) (string, error) {
	switch t.Kind() {
	case reflect.String:
		return "string", nil
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
		reflect.Float32, reflect.Float64:
		return "number", nil
	case reflect.Bool:
		return "boolean", nil
	case reflect.Struct:
		return c.compileStruct(t)
	case reflect.Slice:
		elemRuleName := fmt.Sprintf("%s-elem", ruleName)
		elemType, err := c.compileFieldType(t.Elem(), elemRuleName)
		if err != nil {
			return "", err
		}
		// Array format: "[" ws (elem (ws "," ws elem)*)? ws "]"
		listRule := fmt.Sprintf(`"[" ws (%s (ws "," ws %s)*)? ws "]"`, elemType, elemType)
		c.rules[ruleName] = fmt.Sprintf("%s ::= %s", ruleName, listRule)
		return ruleName, nil
	default:
		return "", fmt.Errorf("unsupported type kind %s for FSM compilation", t.Kind())
	}
}

snippet-2.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

// CompletionRequest represents the payload for llama.cpp server completions.
type CompletionRequest struct {
	Prompt      string   `json:"prompt"`
	Grammar     string   `json:"grammar,omitempty"`
	Temperature float64  `json:"temperature"`
	Stream      bool     `json:"stream"`
	Stop        []string `json:"stop,omitempty"`
	NPredict    int      `json:"n_predict,omitempty"`
}

// CompletionResponse represents a single completion chunk or full response from llama.cpp.
type CompletionResponse struct {
	Content string `json:"content"`
	Stop    bool   `json:"stop"`
}

// LlamaClient is a high-performance client for interacting with llama.cpp or compatible engines.
type LlamaClient struct {
	endpoint   string
	httpClient *http.Client
}

// NewLlamaClient initializes a client with optimized connection pool settings.
func NewLlamaClient(endpoint string) *LlamaClient {
	return &LlamaClient{
		endpoint: endpoint,
		httpClient: &http.Client{
			Transport: &http.Transport{
				MaxIdleConns:        100,
				MaxIdleConnsPerHost: 100,
				IdleConnTimeout:     90 * time.Second,
			},
		},
	}
}

// GenerateStructured fires a request to llama.cpp with a GBNF grammar parameter, blocking until completion.
func (c *LlamaClient) GenerateStructured(ctx context.Context, prompt, grammar string, temp float64) (string, error) {
	reqBody := CompletionRequest{
		Prompt:      prompt,
		Grammar:     grammar,
		Temperature: temp,
		Stream:      false,
		NPredict:    1024,
	}

	jsonBytes, err := json.Marshal(reqBody)
	if err != nil {
		return "", fmt.Errorf("failed to marshal request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint+"/completion", bytes.NewReader(jsonBytes))
	if err != nil {
		return "", fmt.Errorf("failed to create http request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("http request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("server returned error code %d: %s", resp.StatusCode, string(body))
	}

	var completion CompletionResponse
	if err := json.NewDecoder(resp.Body).Decode(&completion); err != nil {
		return "", fmt.Errorf("failed to decode response: %w", err)
	}

	return completion.Content, nil
}

snippet-3.go

package main

import (
	"errors"
	"fmt"
	"strings"
)

type ParserState int

const (
	StateStart ParserState = iota
	StateObjectStart
	StateKeyStart
	StateInKey
	StateKeyEnd
	StateColon
	StateValueStart
	StateInStringValue
	StateInNumericValue
	StateInBoolValue
	StateInArray
	StateInArrayStringValue
	StateArrayValueEnd
	StateValueEnd
	StateComma
	StateObjectEnd
)

// FieldEmitter represents a callback invoked whenever a key-value pair is successfully parsed.
type FieldEmitter func(key string, val interface{})

// StreamFSM handles state transitions for partial streaming JSON structures.
type StreamFSM struct {
	state       ParserState
	buffer      strings.Builder
	currentKey  string
	escapeNext  bool
	emitter     FieldEmitter
	arrayValues []string
}

func NewStreamFSM(emitter FieldEmitter) *StreamFSM {
	return &StreamFSM{
		state:   StateStart,
		emitter: emitter,
	}
}

// Feed processes incoming character chunks and drives the FSM state.
func (f *StreamFSM) Feed(chunk string) error {
	for i := 0; i < len(chunk); i++ {
		c := chunk[i]
		if err := f.consume(c); err != nil {
			return fmt.Errorf("FSM error at char '%c' (state %d): %w", c, f.state, err)
		}
	}
	return nil
}

func (f *StreamFSM) consume(c byte) error {
	// Skip whitespace unless we are actively parsing inside a string key or string value
	if f.state != StateInKey && f.state != StateInStringValue && f.state != StateInArrayStringValue {
		if c == ' ' || c == '\t' || c == '\n' || c == '\r' {
			return nil
		}
	}

	switch f.state {
	case StateStart:
		if c == '{' {
			f.state = StateObjectStart
			return nil
		}
		return errors.New("expected '{' to start object")

	case StateObjectStart:
		if c == '"' {
			f.state = StateInKey
			f.buffer.Reset()
			return nil
		}
		if c == '}' {
			f.state = StateObjectEnd
			return nil
		}
		return errors.New("expected '\"' or '}'")

	case StateInKey:
		if f.escapeNext {
			f.buffer.WriteByte(c)
			f.escapeNext = false
			return nil
		}
		if c == '\\' {
			f.escapeNext = true
			return nil
		}
		if c == '"' {
			f.currentKey = f.buffer.String()
			f.state = StateKeyEnd
			return nil
		}
		f.buffer.WriteByte(c)

	case StateKeyEnd:
		if c == ':' {
			f.state = StateColon
			return nil
		}
		return errors.New("expected ':' after key")

	case StateColon:
		if c == '"' {
			f.state = StateInStringValue
			f.buffer.Reset()
			return nil
		}
		if (c >= '0' && c <= '9') || c == '-' {
			f.state = StateInNumericValue
			f.buffer.Reset()
			f.buffer.WriteByte(c)
			return nil
		}
		if c == 't' || c == 'f' {
			f.state = StateInBoolValue
			f.buffer.Reset()
			f.buffer.WriteByte(c)
			return nil
		}
		if c == '[' {
			f.state = StateInArray
			f.arrayValues = []string{}
			return nil
		}
		return errors.New("expected string, number, boolean, or array start")

	case StateInStringValue:
		if f.escapeNext {
			f.buffer.WriteByte(c)
			f.escapeNext = false
			return nil
		}
		if c == '\\' {
			f.escapeNext = true
			return nil
		}
		if c == '"' {
			f.emitter(f.currentKey, f.buffer.String())
			f.state = StateValueEnd
			return nil
		}
		f.buffer.WriteByte(c)

	case StateInNumericValue:
		if (c >= '0' && c <= '9') || c == '.' || c == 'e' || c == 'E' || c == '+' || c == '-' {
			f.buffer.WriteByte(c)
			return nil
		}
		f.emitter(f.currentKey, f.buffer.String())
		f.state = StateValueEnd
		return f.consume(c)

	case StateInBoolValue:
		if c >= 'a' && c <= 'z' {
			f.buffer.WriteByte(c)
			return nil
		}
		f.emitter(f.currentKey, f.buffer.String() == "true")
		f.state = StateValueEnd
		return f.consume(c)

	case StateInArray:
		if c == ']' {
			f.emitter(f.currentKey, f.arrayValues)
			f.state = StateValueEnd
			return nil
		}
		if c == '"' {
			f.state = StateInArrayStringValue
			f.buffer.Reset()
			return nil
		}
		return errors.New("expected '\"' or ']' inside string array")

	case StateInArrayStringValue:
		if f.escapeNext {
			f.buffer.WriteByte(c)
			f.escapeNext = false
			return nil
		}
		if c == '\\' {
			f.escapeNext = true
			return nil
		}
		if c == '"' {
			f.arrayValues = append(f.arrayValues, f.buffer.String())
			f.state = StateArrayValueEnd
			return nil
		}
		f.buffer.WriteByte(c)

	case StateArrayValueEnd:
		if c == ',' {
			f.state = StateInArray // transition back to expect another string
			return nil
		}
		if c == ']' {
			f.emitter(f.currentKey, f.arrayValues)
			f.state = StateValueEnd
			return nil
		}
		return errors.New("expected ',' or ']' inside array")

	case StateValueEnd:
		if c == ',' {
			f.state = StateComma
			return nil
		}
		if c == '}' {
			f.state = StateObjectEnd
			return nil
		}
		return errors.New("expected ',' or '}' after value")

	case StateComma:
		if c == '"' {
			f.state = StateInKey
			f.buffer.Reset()
			return nil
		}
		return errors.New("expected '\"' key after comma")

	case StateObjectEnd:
		return errors.New("cannot consume characters after root object closed")
	}

	return nil
}

snippet-4.go

package main

import (
	"math"
)

// LogitSampler wraps raw logit arrays and performs masking before sampling.
type LogitSampler struct {
	vocabulary []string // maps token ID -> string token value
}

func NewLogitSampler(vocab []string) *LogitSampler {
	return &LogitSampler{vocabulary: vocab}
}

// FilterLogits applies logit masking to enforce transitions.
// It sets logits of all tokens that do not match the expected state pattern to -infinity.
func (s *LogitSampler) FilterLogits(logits []float32, validTransitions []int) []float32 {
	// Create a fast lookup map for valid token IDs
	validMap := make(map[int]struct{}, len(validTransitions))
	for _, id := range validTransitions {
		validMap[id] = struct{}{}
	}

	// Apply logit masking
	maskedLogits := make([]float32, len(logits))
	for idx, val := range logits {
		if _, ok := validMap[idx]; ok {
			maskedLogits[idx] = val
		} else {
			// -Infinity representation in 32-bit float to prevent sampling
			maskedLogits[idx] = -math.MaxFloat32
		}
	}

	return maskedLogits
}

// SampleSoftmax computes standard softmax over logits and selects a token.
func (s *LogitSampler) SampleSoftmax(logits []float32, temperature float64) int {
	var maxLogit float32 = -math.MaxFloat32
	for _, val := range logits {
		if val > maxLogit {
			maxLogit = val
		}
	}

	var sumExp float64
	probabilities := make([]float64, len(logits))

	for idx, val := range logits {
		if val == -math.MaxFloat32 {
			probabilities[idx] = 0.0
			continue
		}
		// Apply temperature scaling
		scaledVal := float64(val-maxLogit) / temperature
		expVal := math.Exp(scaledVal)
		probabilities[idx] = expVal
		sumExp += expVal
	}

	// Normalize probabilities
	if sumExp > 0 {
		for idx := range probabilities {
			probabilities[idx] /= sumExp
		}
	}

	// Simple argmax selection (greedy sampling for demonstration)
	bestIdx := -1
	var maxProb float64 = -1.0
	for idx, prob := range probabilities {
		if prob > maxProb {
			maxProb = prob
			bestIdx = idx
		}
	}

	return bestIdx
}

snippet-5.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"time"
)

var ErrStateLoopDetected = errors.New("FSM guidance stuck in infinite state loop or token limit exceeded")

// GeneratorFunc represents the closure that executes a stream request to the LLM.
type GeneratorFunc func(ctx context.Context, temperature float64) (<-chan string, error)

// ResilientRunner manages FSM generation runs, monitoring for loops and timeouts.
type ResilientRunner struct {
	maxTokens      int
	timeout        time.Duration
	maxRetries     int
	fallbackEngine GeneratorFunc
}

func NewResilientRunner(maxTokens int, timeout time.Duration, maxRetries int) *ResilientRunner {
	return &ResilientRunner{
		maxTokens:  maxTokens,
		timeout:    timeout,
		maxRetries: maxRetries,
	}
}

// ExecuteWithFallback runs the generator and monitors the state stream.
// If a loop or timeout is encountered, it retries with lowered temperature or executes fallback logic.
func (r *ResilientRunner) ExecuteWithFallback(ctx context.Context, generator GeneratorFunc, validator func(string) error) (string, error) {
	temp := 0.7 // initial temperature
	var lastErr error

	for attempt := 0; attempt <= r.maxRetries; attempt++ {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		default:
		}

		log.Printf("Starting structured generation attempt %d (temp=%.2f)", attempt+1, temp)
		result, err := r.runAttempt(ctx, generator, temp, validator)
		if err == nil {
			return result, nil
		}

		lastErr = err
		log.Printf("Attempt %d failed: %v", attempt+1, err)

		if errors.Is(err, ErrStateLoopDetected) {
			// Lower the temperature for the next run to make it more deterministic
			temp = mathMax(0.0, temp-0.3)
		} else {
			// For non-loop errors, apply minor jitter or step down
			temp = mathMax(0.0, temp-0.2)
		}
	}

	return "", fmt.Errorf("all structured generation attempts failed. Last error: %w", lastErr)
}

func (r *ResilientRunner) runAttempt(ctx context.Context, generator GeneratorFunc, temp float64, validator func(string) error) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, r.timeout)
	defer cancel()

	stream, err := generator(ctx, temp)
	if err != nil {
		return "", fmt.Errorf("failed to initialize stream: %w", err)
	}

	var outputBuffer string
	tokenCount := 0
	consecutiveRepeats := 0
	var lastToken string

	for {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case token, ok := <-stream:
			if !ok {
				// Stream closed normally. Validate full output.
				if err := validator(outputBuffer); err != nil {
					return "", fmt.Errorf("final structure validation failed: %w", err)
				}
				return outputBuffer, nil
			}

			// Watchdog: detect simple repetitive token patterns
			if token == lastToken {
				consecutiveRepeats++
				if consecutiveRepeats > 15 {
					return "", ErrStateLoopDetected
				}
			} else {
				consecutiveRepeats = 0
			}
			lastToken = token

			outputBuffer += token
			tokenCount++

			if tokenCount > r.maxTokens {
				return "", fmt.Errorf("%w: generated %d tokens", ErrStateLoopDetected, tokenCount)
			}
		}
	}
}

func mathMax(a, b float64) float64 {
	if a > b {
		return a
	}
	return b
}

snippet-6.go

package main

import (
	"context"
	"fmt"
	"log"
	"time"
)

// PaymentRoute defines the strict schema we want the LLM to output.
type PaymentRoute struct {
	RouteID     string   `json:"route_id"`
	Provider    string   `json:"provider"`
	BackupPaths []string `json:"backup_paths"`
	Enabled     bool     `json:"enabled"`
	MaxLimit    float64  `json:"max_limit"`
}

func main() {
	// 1. Initialize GBNF Compiler
	comp := NewCompiler()
	grammar, err := comp.Compile(PaymentRoute{})
	if err != nil {
		log.Fatalf("Failed to compile schema: %v", err)
	}

	fmt.Println("=== Compiled GBNF Grammar ===")
	fmt.Println(grammar)
	fmt.Println("=============================")

	// 2. Setup Client
	_ = NewLlamaClient("http://localhost:8080")

	// 3. Setup Streaming Parser with emitter callback
	streamFSM := NewStreamFSM(func(key string, val interface{}) {
		fmt.Printf("[Stream Event] Parsed Key: %q = Value: %v (%T)\n", key, val, val)
	})

	// 4. Setup Resilient Middleware Runner
	runner := NewResilientRunner(512, 10*time.Second, 2)

	// Define generator closure
	generator := func(ctx context.Context, temp float64) (<-chan string, error) {
		tokenChan := make(chan string, 100)
		
		// Simulate a background stream reader
		go func() {
			defer close(tokenChan)
			// In production, you would stream from llamaClient and feed here.
			// For demonstration, we feed a pre-determined valid stream chunk sequence:
			chunks := []string{
				`{`,
				`"route_id"`, `: "route-use1-primary"`, `,`,
				` "provider":`, ` "stripe"`, `,`,
				` "backup_paths": [`, `"adyen"`, `, "checkout"`, `],`,
				` "enabled"`, `: true`, `,`,
				` "max_limit"`, `: 150000.50`,
				`}`,
			}
			for _, chunk := range chunks {
				select {
				case <-ctx.Done():
					return
				case tokenChan <- chunk:
					time.Sleep(10 * time.Millisecond) // simulate network delay
				}
			}
		}()

		return tokenChan, nil
	}

	// Validate function
	validator := func(rawJSON string) error {
		// Feed the FSM parser to verify syntactic and structural correctness
		return streamFSM.Feed(rawJSON)
	}

	ctx := context.Background()
	resultJSON, err := runner.ExecuteWithFallback(ctx, generator, validator)
	if err != nil {
		log.Fatalf("Execution failed: %v", err)
	}

	fmt.Printf("\n=== Final Validated JSON ===\n%s\n", resultJSON)
}