iop/apps/edge/internal/openai/json_field_patch.go

package openai

import (
	"encoding/json"
	"fmt"
	"sort"
)

// topLevelJSONPatch is a request-local instruction for one top-level JSON
// member. A nil value removes the member; a non-nil value replaces it or is
// appended when the member is absent.
type topLevelJSONPatch struct {
	name  string
	value []byte
}

type topLevelJSONFieldSpan struct {
	name       string
	memberFrom int
	memberTo   int
	valueFrom  int
	valueTo    int
}

type jsonByteEdit struct {
	from        int
	to          int
	replacement []byte
}

type topLevelJSONPatchPlan struct {
	body       []byte
	edits      []jsonByteEdit
	outputSize int
}

func planTopLevelJSONPatches(body []byte, patches []topLevelJSONPatch) (topLevelJSONPatchPlan, error) {
	if len(patches) == 0 {
		return topLevelJSONPatchPlan{body: body, outputSize: len(body)}, nil
	}
	fields, objectEnd, err := scanTopLevelJSONObject(body)
	if err != nil {
		return topLevelJSONPatchPlan{}, err
	}

	patchByName := make(map[string]topLevelJSONPatch, len(patches))
	for _, patch := range patches {
		if patch.name == "" {
			return topLevelJSONPatchPlan{}, fmt.Errorf("JSON patch field is required")
		}
		if patch.value != nil && !json.Valid(patch.value) {
			return topLevelJSONPatchPlan{}, fmt.Errorf("invalid JSON patch for %s", patch.name)
		}
		patchByName[patch.name] = patch
	}

	found := make(map[string]bool, len(patchByName))
	edits := make([]jsonByteEdit, 0, len(patches))
	for index, field := range fields {
		patch, ok := patchByName[field.name]
		if !ok {
			continue
		}
		found[field.name] = true
		if patch.value != nil {
			edits = append(edits, jsonByteEdit{from: field.valueFrom, to: field.valueTo, replacement: patch.value})
			continue
		}

		// Removing a target member also removes exactly one adjacent comma. The
		// untouched members and all bytes outside this local span remain intact.
		from, to := field.memberFrom, field.memberTo
		if index+1 < len(fields) {
			to = fields[index+1].memberFrom
		} else if index > 0 {
			from = fields[index-1].memberTo
		}
		edits = append(edits, jsonByteEdit{from: from, to: to})
	}

	var insertion []byte
	memberCount := len(fields)
	for _, patch := range patches {
		if found[patch.name] || patch.value == nil {
			continue
		}
		nameJSON, err := json.Marshal(patch.name)
		if err != nil {
			return topLevelJSONPatchPlan{}, err
		}
		if memberCount > 0 {
			insertion = append(insertion, ',')
		}
		insertion = append(insertion, nameJSON...)
		insertion = append(insertion, ':')
		insertion = append(insertion, patch.value...)
		memberCount++
	}
	if len(insertion) > 0 {
		edits = append(edits, jsonByteEdit{from: objectEnd, to: objectEnd, replacement: insertion})
	}

	sort.SliceStable(edits, func(i, j int) bool { return edits[i].from < edits[j].from })
	outputSize := len(body)
	lastTo := 0
	for _, edit := range edits {
		if edit.from < lastTo || edit.from < 0 || edit.to < edit.from || edit.to > len(body) {
			return topLevelJSONPatchPlan{}, fmt.Errorf("overlapping JSON field patches")
		}
		outputSize += len(edit.replacement) - (edit.to - edit.from)
		if outputSize < 0 {
			return topLevelJSONPatchPlan{}, fmt.Errorf("JSON patch output size overflow")
		}
		lastTo = edit.to
	}
	return topLevelJSONPatchPlan{body: body, edits: edits, outputSize: outputSize}, nil
}

func (p topLevelJSONPatchPlan) apply() []byte {
	if len(p.edits) == 0 {
		return p.body
	}
	out := make([]byte, 0, p.outputSize)
	last := 0
	for _, edit := range p.edits {
		out = append(out, p.body[last:edit.from]...)
		out = append(out, edit.replacement...)
		last = edit.to
	}
	out = append(out, p.body[last:]...)
	return out
}

func scanTopLevelJSONObject(body []byte) ([]topLevelJSONFieldSpan, int, error) {
	if !json.Valid(body) {
		return nil, 0, fmt.Errorf("invalid JSON request")
	}
	i := skipJSONSpace(body, 0)
	if i >= len(body) || body[i] != '{' {
		return nil, 0, fmt.Errorf("JSON request must be an object")
	}
	i++
	fields := make([]topLevelJSONFieldSpan, 0, 8)
	for {
		i = skipJSONSpace(body, i)
		if i >= len(body) {
			return nil, 0, fmt.Errorf("invalid JSON request")
		}
		if body[i] == '}' {
			return fields, i, nil
		}
		memberFrom := i
		keyTo, err := scanJSONString(body, i)
		if err != nil {
			return nil, 0, fmt.Errorf("invalid JSON request")
		}
		var name string
		if err := json.Unmarshal(body[i:keyTo], &name); err != nil {
			return nil, 0, fmt.Errorf("invalid JSON request")
		}
		i = skipJSONSpace(body, keyTo)
		if i >= len(body) || body[i] != ':' {
			return nil, 0, fmt.Errorf("invalid JSON request")
		}
		i = skipJSONSpace(body, i+1)
		valueFrom := i
		valueTo, err := scanJSONValue(body, i)
		if err != nil {
			return nil, 0, fmt.Errorf("invalid JSON request")
		}
		i = skipJSONSpace(body, valueTo)
		memberTo := i
		fields = append(fields, topLevelJSONFieldSpan{
			name: name, memberFrom: memberFrom, memberTo: memberTo,
			valueFrom: valueFrom, valueTo: valueTo,
		})
		if i >= len(body) {
			return nil, 0, fmt.Errorf("invalid JSON request")
		}
		switch body[i] {
		case ',':
			i++
		case '}':
			return fields, i, nil
		default:
			return nil, 0, fmt.Errorf("invalid JSON request")
		}
	}
}

func scanJSONString(body []byte, start int) (int, error) {
	if start >= len(body) || body[start] != '"' {
		return 0, fmt.Errorf("JSON string expected")
	}
	escaped := false
	for i := start + 1; i < len(body); i++ {
		if escaped {
			escaped = false
			continue
		}
		switch body[i] {
		case '\\':
			escaped = true
		case '"':
			return i + 1, nil
		}
	}
	return 0, fmt.Errorf("unterminated JSON string")
}

func scanJSONValue(body []byte, start int) (int, error) {
	if start >= len(body) {
		return 0, fmt.Errorf("JSON value expected")
	}
	if body[start] == '"' {
		return scanJSONString(body, start)
	}
	if body[start] == '{' || body[start] == '[' {
		stack := []byte{body[start]}
		inString, escaped := false, false
		for i := start + 1; i < len(body); i++ {
			ch := body[i]
			if inString {
				if escaped {
					escaped = false
				} else if ch == '\\' {
					escaped = true
				} else if ch == '"' {
					inString = false
				}
				continue
			}
			switch ch {
			case '"':
				inString = true
			case '{', '[':
				stack = append(stack, ch)
			case '}', ']':
				open := stack[len(stack)-1]
				if (open == '{' && ch != '}') || (open == '[' && ch != ']') {
					return 0, fmt.Errorf("mismatched JSON container")
				}
				stack = stack[:len(stack)-1]
				if len(stack) == 0 {
					return i + 1, nil
				}
			}
		}
		return 0, fmt.Errorf("unterminated JSON container")
	}
	for i := start; i < len(body); i++ {
		switch body[i] {
		case ' ', '\t', '\r', '\n', ',', '}', ']':
			if i == start {
				return 0, fmt.Errorf("JSON value expected")
			}
			return i, nil
		}
	}
	return len(body), nil
}

func skipJSONSpace(body []byte, i int) int {
	for i < len(body) {
		switch body[i] {
		case ' ', '\t', '\r', '\n':
			i++
		default:
			return i
		}
	}
	return i
}