types: allow tool function parameters with a single type or an array of types (#9434)

api/types.go

@@ -166,6 +166,48 @@ type Tool struct {
 	Function ToolFunction `json:"function"`
 }
 
+// PropertyType can be either a string or an array of strings
+type PropertyType []string
+
+// UnmarshalJSON implements the json.Unmarshaler interface
+func (pt *PropertyType) UnmarshalJSON(data []byte) error {
+	// Try to unmarshal as a string first
+	var s string
+	if err := json.Unmarshal(data, &s); err == nil {
+		*pt = []string{s}
+		return nil
+	}
+
+	// If that fails, try to unmarshal as an array of strings
+	var a []string
+	if err := json.Unmarshal(data, &a); err != nil {
+		return err
+	}
+	*pt = a
+	return nil
+}
+
+// MarshalJSON implements the json.Marshaler interface
+func (pt PropertyType) MarshalJSON() ([]byte, error) {
+	if len(pt) == 1 {
+		// If there's only one type, marshal as a string
+		return json.Marshal(pt[0])
+	}
+	// Otherwise marshal as an array
+	return json.Marshal([]string(pt))
+}
+
+// String returns a string representation of the PropertyType
+func (pt PropertyType) String() string {
+	if len(pt) == 0 {
+		return ""
+	}
+	if len(pt) == 1 {
+		return pt[0]
+	}
+	return fmt.Sprintf("%v", []string(pt))
+}
+
 type ToolFunction struct {
 	Name        string `json:"name"`
 	Description string `json:"description"`
@@ -173,9 +215,9 @@ type ToolFunction struct {
 		Type       string   `json:"type"`
 		Required   []string `json:"required"`
 		Properties map[string]struct {
-			Type        string   `json:"type"`
-			Description string   `json:"description"`
-			Enum        []string `json:"enum,omitempty"`
+			Type        PropertyType `json:"type"`
+			Description string       `json:"description"`
+			Enum        []string     `json:"enum,omitempty"`
 		} `json:"properties"`
 	} `json:"parameters"`
 }