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"`
 }

api/types_test.go

@@ -231,3 +231,83 @@ func TestMessage_UnmarshalJSON(t *testing.T) {
 		}
 	}
 }
+
+func TestPropertyType_UnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected PropertyType
+	}{
+		{
+			name:     "string type",
+			input:    `"string"`,
+			expected: PropertyType{"string"},
+		},
+		{
+			name:     "array of types",
+			input:    `["string", "number"]`,
+			expected: PropertyType{"string", "number"},
+		},
+		{
+			name:     "array with single type",
+			input:    `["string"]`,
+			expected: PropertyType{"string"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var pt PropertyType
+			if err := json.Unmarshal([]byte(test.input), &pt); err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+
+			if len(pt) != len(test.expected) {
+				t.Errorf("Length mismatch: got %v, expected %v", len(pt), len(test.expected))
+			}
+
+			for i, v := range pt {
+				if v != test.expected[i] {
+					t.Errorf("Value mismatch at index %d: got %v, expected %v", i, v, test.expected[i])
+				}
+			}
+		})
+	}
+}
+
+func TestPropertyType_MarshalJSON(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    PropertyType
+		expected string
+	}{
+		{
+			name:     "single type",
+			input:    PropertyType{"string"},
+			expected: `"string"`,
+		},
+		{
+			name:     "multiple types",
+			input:    PropertyType{"string", "number"},
+			expected: `["string","number"]`,
+		},
+		{
+			name:     "empty type",
+			input:    PropertyType{},
+			expected: `[]`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			data, err := json.Marshal(test.input)
+			if err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+
+			if string(data) != test.expected {
+				t.Errorf("Marshaled data mismatch: got %v, expected %v", string(data), test.expected)
+			}
+		})
+	}
+}

openai/openai_test.go

@@ -283,24 +283,24 @@ func TestChatMiddleware(t *testing.T) {
 								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        api.PropertyType `json:"type"`
+									Description string           `json:"description"`
+									Enum        []string         `json:"enum,omitempty"`
 								} `json:"properties"`
 							}{
 								Type:     "object",
 								Required: []string{"location"},
 								Properties: map[string]struct {
-									Type        string   `json:"type"`
-									Description string   `json:"description"`
-									Enum        []string `json:"enum,omitempty"`
+									Type        api.PropertyType `json:"type"`
+									Description string           `json:"description"`
+									Enum        []string         `json:"enum,omitempty"`
 								}{
 									"location": {
-										Type:        "string",
+										Type:        api.PropertyType{"string"},
 										Description: "The city and state",
 									},
 									"unit": {
-										Type: "string",
+										Type: api.PropertyType{"string"},
 										Enum: []string{"celsius", "fahrenheit"},
 									},
 								},

server/routes_generate_test.go

@@ -372,24 +372,24 @@ func TestGenerateChat(t *testing.T) {
 						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        api.PropertyType `json:"type"`
+							Description string           `json:"description"`
+							Enum        []string         `json:"enum,omitempty"`
 						} `json:"properties"`
 					}{
 						Type:     "object",
 						Required: []string{"location"},
 						Properties: map[string]struct {
-							Type        string   `json:"type"`
-							Description string   `json:"description"`
-							Enum        []string `json:"enum,omitempty"`
+							Type        api.PropertyType `json:"type"`
+							Description string           `json:"description"`
+							Enum        []string         `json:"enum,omitempty"`
 						}{
 							"location": {
-								Type:        "string",
+								Type:        api.PropertyType{"string"},
 								Description: "The city and state",
 							},
 							"unit": {
-								Type: "string",
+								Type: api.PropertyType{"string"},
 								Enum: []string{"celsius", "fahrenheit"},
 							},
 						},
@@ -469,24 +469,24 @@ func TestGenerateChat(t *testing.T) {
 						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        api.PropertyType `json:"type"`
+							Description string           `json:"description"`
+							Enum        []string         `json:"enum,omitempty"`
 						} `json:"properties"`
 					}{
 						Type:     "object",
 						Required: []string{"location"},
 						Properties: map[string]struct {
-							Type        string   `json:"type"`
-							Description string   `json:"description"`
-							Enum        []string `json:"enum,omitempty"`
+							Type        api.PropertyType `json:"type"`
+							Description string           `json:"description"`
+							Enum        []string         `json:"enum,omitempty"`
 						}{
 							"location": {
-								Type:        "string",
+								Type:        api.PropertyType{"string"},
 								Description: "The city and state",
 							},
 							"unit": {
-								Type: "string",
+								Type: api.PropertyType{"string"},
 								Enum: []string{"celsius", "fahrenheit"},
 							},
 						},