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

package openai

import (
	"encoding/json"
	"reflect"
	"strings"
)

func normalizeNativeToolCallsForResponse(toolCalls []any, tools []any) []any {
	specs := requestedToolSpecs(tools)
	if len(toolCalls) == 0 || len(specs) == 0 {
		return toolCalls
	}
	var out []any
	for i, item := range toolCalls {
		normalized, changed := normalizeNativeToolCallForResponse(item, specs)
		if !changed {
			if out != nil {
				out[i] = item
			}
			continue
		}
		if out == nil {
			out = make([]any, len(toolCalls))
			copy(out, toolCalls)
		}
		out[i] = normalized
	}
	if out == nil {
		return toolCalls
	}
	return out
}

func normalizeNativeToolCallForResponse(item any, specs map[string]textToolSpec) (any, bool) {
	call, ok := item.(map[string]any)
	if !ok {
		return item, false
	}
	fn, ok := call["function"].(map[string]any)
	if !ok {
		return item, false
	}
	name, _ := fn["name"].(string)
	spec, ok := specs[strings.TrimSpace(name)]
	if !ok {
		return item, false
	}
	arguments, ok := normalizeNativeToolCallArguments(fn["arguments"], spec)
	if !ok {
		return item, false
	}
	nextFn := make(map[string]any, len(fn))
	for key, value := range fn {
		nextFn[key] = value
	}
	nextFn["arguments"] = arguments
	nextCall := make(map[string]any, len(call))
	for key, value := range call {
		nextCall[key] = value
	}
	nextCall["function"] = nextFn
	return nextCall, true
}

func normalizeNativeToolCallArguments(value any, spec textToolSpec) (string, bool) {
	var args map[string]any
	original := ""
	switch v := value.(type) {
	case string:
		original = v
		if err := json.Unmarshal([]byte(v), &args); err != nil {
			return "", false
		}
	case map[string]any:
		args = v
	default:
		return "", false
	}
	if args == nil {
		return "", false
	}
	decoded := decodeStringifiedSchemaValues(args, spec.parameters)
	normalized := normalizeTextToolCallArguments(spec.name, args, spec)
	changed := decoded || !reflect.DeepEqual(normalized, args) || original == ""
	if !changed {
		return "", false
	}
	encoded, err := json.Marshal(normalized)
	if err != nil {
		return "", false
	}
	if original != "" && string(encoded) == original {
		return "", false
	}
	return string(encoded), true
}

func decodeStringifiedSchemaValues(args map[string]any, parameters map[string]any) bool {
	props := schemaObjectProperties(parameters)
	if len(props) == 0 {
		return false
	}
	changed := false
	for key, value := range args {
		schema, ok := props[key]
		if !ok {
			continue
		}
		if decoded, ok := decodeStringifiedSchemaValue(value, schema); ok {
			args[key] = decoded
			changed = true
		}
	}
	return changed
}

func decodeStringifiedSchemaValue(value any, schema any) (any, bool) {
	switch v := value.(type) {
	case string:
		trimmed := strings.TrimSpace(v)
		if trimmed == "" || (!strings.HasPrefix(trimmed, "[") && !strings.HasPrefix(trimmed, "{")) {
			return nil, false
		}
		var decoded any
		if err := json.Unmarshal([]byte(trimmed), &decoded); err != nil {
			return nil, false
		}
		if !schemaAcceptsDecodedValue(decoded, schema) {
			return nil, false
		}
		return decoded, true
	case map[string]any:
		if changed := decodeStringifiedSchemaValues(v, mapValue(schema)); changed {
			return v, true
		}
	case []any:
		itemSchema := schemaArrayItemSchema(mapValue(schema))
		if itemSchema == nil {
			return nil, false
		}
		var out []any
		for i, item := range v {
			decoded, ok := decodeStringifiedSchemaValue(item, itemSchema)
			if !ok {
				if out != nil {
					out[i] = item
				}
				continue
			}
			if out == nil {
				out = make([]any, len(v))
				copy(out, v)
			}
			out[i] = decoded
		}
		if out != nil {
			return out, true
		}
	}
	return nil, false
}

func schemaAcceptsDecodedValue(value any, schema any) bool {
	m := mapValue(schema)
	switch value.(type) {
	case []any:
		return schemaAllowsType(m, "array") || schemaArrayItemSchema(m) != nil
	case map[string]any:
		return schemaAllowsType(m, "object") || len(schemaObjectProperties(m)) > 0
	default:
		return false
	}
}

func normalizeTextToolCallArguments(_ string, args map[string]any, spec textToolSpec) map[string]any {
	topProps := schemaObjectProperties(spec.parameters)
	commandsSchema := topProps["commands"]
	commandField := schemaSingleCommandField(topProps)
	commands, ok := args["commands"].([]any)
	var out map[string]any
	ensureOut := func() map[string]any {
		if out != nil {
			return out
		}
		out = make(map[string]any, len(args))
		for key, value := range args {
			if shouldStripToolArgumentKey(key, topProps) {
				continue
			}
			out[key] = value
		}
		return out
	}
	for key := range args {
		if shouldStripToolArgumentKey(key, topProps) {
			ensureOut()
			break
		}
	}
	if !ok {
		if out != nil {
			return out
		}
		return args
	}

	if commandField != "" && !schemaHasProperty(topProps, "commands") {
		if rendered, ok := firstCommandAsShellString(commands); ok {
			collapsed := ensureOut()
			delete(collapsed, "commands")
			collapsed[commandField] = rendered
			return collapsed
		}
	}

	itemSchema := schemaArrayItemSchema(commandsSchema)
	itemAllowsString := schemaAllowsType(itemSchema, "string")
	itemAllowsObject := schemaAllowsType(itemSchema, "object") || len(schemaObjectProperties(itemSchema)) > 0
	itemObjectProps := schemaObjectProperties(itemSchema)
	var normalized []any
	for i, command := range commands {
		cleaned, changed := normalizeCommandArgument(command, commandArgSchema{
			allowsString: itemAllowsString,
			allowsObject: itemAllowsObject,
			objectProps:  itemObjectProps,
		})
		if !changed {
			if normalized != nil {
				normalized[i] = command
			}
			continue
		}
		if normalized == nil {
			normalized = make([]any, len(commands))
			copy(normalized, commands)
		}
		normalized[i] = cleaned
	}
	if normalized == nil {
		if out != nil {
			return out
		}
		return args
	}
	ensureOut()["commands"] = normalized
	return out
}