diff --git a/api/types.go b/api/types.go index a38b335b7..814271ada 100644 --- a/api/types.go +++ b/api/types.go @@ -165,6 +165,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"` @@ -172,9 +214,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"` } diff --git a/api/types_test.go b/api/types_test.go index a9de5a9a9..6aedb3815 100644 --- a/api/types_test.go +++ b/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) + } + }) + } +} diff --git a/openai/openai_test.go b/openai/openai_test.go index d8c821d39..6e08d3ab5 100644 --- a/openai/openai_test.go +++ b/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"}, }, }, diff --git a/server/routes_generate_test.go b/server/routes_generate_test.go index aa263bf97..29f3d4c9a 100644 --- a/server/routes_generate_test.go +++ b/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"}, }, },