From 6747099d715bfc8c9679c3f6365a08ffb950ad32 Mon Sep 17 00:00:00 2001
From: Parth Sareen <parth.sareen@ollama.com>
Date: Tue, 8 Apr 2025 15:05:38 -0700
Subject: [PATCH] types: add any type and validation for ToolFunction enum
 (#10166)

---
 api/types.go                   |  2 +-
 api/types_test.go              | 61 ++++++++++++++++++++++++++++++++++
 openai/openai_test.go          |  6 ++--
 server/routes_generate_test.go | 12 +++----
 4 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/api/types.go b/api/types.go
index 0e63ef8bb..4e8486acf 100644
--- a/api/types.go
+++ b/api/types.go
@@ -217,7 +217,7 @@ type ToolFunction struct {
 		Properties map[string]struct {
 			Type        PropertyType `json:"type"`
 			Description string       `json:"description"`
-			Enum        []string     `json:"enum,omitempty"`
+			Enum        []any        `json:"enum,omitempty"`
 		} `json:"properties"`
 	} `json:"parameters"`
 }
diff --git a/api/types_test.go b/api/types_test.go
index e22c047f0..1a6fc811c 100644
--- a/api/types_test.go
+++ b/api/types_test.go
@@ -232,6 +232,67 @@ func TestMessage_UnmarshalJSON(t *testing.T) {
 	}
 }
 
+func TestToolFunction_UnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		wantErr string
+	}{
+		{
+			name: "valid enum with same types",
+			input: `{
+				"name": "test",
+				"description": "test function",
+				"parameters": {
+					"type": "object",
+					"required": ["test"],
+					"properties": {
+						"test": {
+							"type": "string",
+							"description": "test prop",
+							"enum": ["a", "b", "c"]
+						}
+					}
+				}
+			}`,
+			wantErr: "",
+		},
+		{
+			name: "empty enum array",
+			input: `{
+				"name": "test",
+				"description": "test function",
+				"parameters": {
+					"type": "object",
+					"required": ["test"],
+					"properties": {
+						"test": {
+							"type": "string",
+							"description": "test prop",
+							"enum": []
+						}
+					}
+				}
+			}`,
+			wantErr: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var tf ToolFunction
+			err := json.Unmarshal([]byte(tt.input), &tf)
+
+			if tt.wantErr != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tt.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
 func TestPropertyType_UnmarshalJSON(t *testing.T) {
 	tests := []struct {
 		name     string
diff --git a/openai/openai_test.go b/openai/openai_test.go
index 6039cb652..46fce7c8d 100644
--- a/openai/openai_test.go
+++ b/openai/openai_test.go
@@ -285,7 +285,7 @@ func TestChatMiddleware(t *testing.T) {
 								Properties map[string]struct {
 									Type        api.PropertyType `json:"type"`
 									Description string           `json:"description"`
-									Enum        []string         `json:"enum,omitempty"`
+									Enum        []any            `json:"enum,omitempty"`
 								} `json:"properties"`
 							}{
 								Type:     "object",
@@ -293,7 +293,7 @@ func TestChatMiddleware(t *testing.T) {
 								Properties: map[string]struct {
 									Type        api.PropertyType `json:"type"`
 									Description string           `json:"description"`
-									Enum        []string         `json:"enum,omitempty"`
+									Enum        []any            `json:"enum,omitempty"`
 								}{
 									"location": {
 										Type:        api.PropertyType{"string"},
@@ -301,7 +301,7 @@ func TestChatMiddleware(t *testing.T) {
 									},
 									"unit": {
 										Type: api.PropertyType{"string"},
-										Enum: []string{"celsius", "fahrenheit"},
+										Enum: []any{"celsius", "fahrenheit"},
 									},
 								},
 							},
diff --git a/server/routes_generate_test.go b/server/routes_generate_test.go
index 2613978aa..00a50cc3d 100644
--- a/server/routes_generate_test.go
+++ b/server/routes_generate_test.go
@@ -374,7 +374,7 @@ func TestGenerateChat(t *testing.T) {
 						Properties map[string]struct {
 							Type        api.PropertyType `json:"type"`
 							Description string           `json:"description"`
-							Enum        []string         `json:"enum,omitempty"`
+							Enum        []any            `json:"enum,omitempty"`
 						} `json:"properties"`
 					}{
 						Type:     "object",
@@ -382,7 +382,7 @@ func TestGenerateChat(t *testing.T) {
 						Properties: map[string]struct {
 							Type        api.PropertyType `json:"type"`
 							Description string           `json:"description"`
-							Enum        []string         `json:"enum,omitempty"`
+							Enum        []any            `json:"enum,omitempty"`
 						}{
 							"location": {
 								Type:        api.PropertyType{"string"},
@@ -390,7 +390,7 @@ func TestGenerateChat(t *testing.T) {
 							},
 							"unit": {
 								Type: api.PropertyType{"string"},
-								Enum: []string{"celsius", "fahrenheit"},
+								Enum: []any{"celsius", "fahrenheit"},
 							},
 						},
 					},
@@ -471,7 +471,7 @@ func TestGenerateChat(t *testing.T) {
 						Properties map[string]struct {
 							Type        api.PropertyType `json:"type"`
 							Description string           `json:"description"`
-							Enum        []string         `json:"enum,omitempty"`
+							Enum        []any            `json:"enum,omitempty"`
 						} `json:"properties"`
 					}{
 						Type:     "object",
@@ -479,7 +479,7 @@ func TestGenerateChat(t *testing.T) {
 						Properties: map[string]struct {
 							Type        api.PropertyType `json:"type"`
 							Description string           `json:"description"`
-							Enum        []string         `json:"enum,omitempty"`
+							Enum        []any            `json:"enum,omitempty"`
 						}{
 							"location": {
 								Type:        api.PropertyType{"string"},
@@ -487,7 +487,7 @@ func TestGenerateChat(t *testing.T) {
 							},
 							"unit": {
 								Type: api.PropertyType{"string"},
-								Enum: []string{"celsius", "fahrenheit"},
+								Enum: []any{"celsius", "fahrenheit"},
 							},
 						},
 					},