struct for tools

v2/ollamaclient.go

@@ -87,7 +87,7 @@ type Config struct {
 	Verbose                   bool
 	ContextLength             int64
 	SystemPrompt              string
-	Tools                     []json.RawMessage
+	Tools                     []Tool
 }
 
 // Cache is used for caching reproducible results from Ollama (seed -1, temperature 0)
@@ -179,7 +179,7 @@ func (oc *Config) SetContextLength(contextLength int64) {
 }
 
 // SetTool sets the tools for this Ollama config
-func (oc *Config) SetTool(tool json.RawMessage) {
+func (oc *Config) SetTool(tool Tool) {
 	oc.Tools = append(oc.Tools, tool)
 }
 

v2/stream.go

@@ -10,12 +10,12 @@ import (
 
 // GenerateChatRequest represents the request payload for generating chat output
 type GenerateChatRequest struct {
-	Model    string            `json:"model"`
-	Messages []Message         `json:"messages,omitempty"`
-	Images   []string          `json:"images,omitempty"` // base64 encoded images
-	Stream   bool              `json:"stream"`
-	Tools    []json.RawMessage `json:"tools,omitempty"`
-	Options  RequestOptions    `json:"options,omitempty"`
+	Model    string         `json:"model"`
+	Messages []Message      `json:"messages,omitempty"`
+	Images   []string       `json:"images,omitempty"` // base64 encoded images
+	Stream   bool           `json:"stream"`
+	Tools    []Tool         `json:"tools,omitempty"`
+	Options  RequestOptions `json:"options,omitempty"`
 }
 
 // GenerateChatResponse represents the response data from the generate chat API call

v2/tools.go

@@ -1,17 +1,23 @@
 package ollamaclient
 
+// ToolParameters represents the parameters of a tool
+type ToolParameters struct {
+	Type       string                  `json:"type"`
+	Properties map[string]ToolProperty `json:"properties"`
+	Required   []string                `json:"required"`
+}
+
+// ToolFunction represents the function details within a tool
+type ToolFunction struct {
+	Name        string         `json:"name"`
+	Description string         `json:"description"`
+	Parameters  ToolParameters `json:"parameters"`
+}
+
 // Tool represents a tool or function that can be used by the Ollama client
 type Tool struct {
-	Type     string `json:"type"`
-	Function struct {
-		Name        string `json:"name"`
-		Description string `json:"description"`
-		Parameters  struct {
-			Type       string                  `json:"type"`
-			Properties map[string]ToolProperty `json:"properties"`
-			Required   []string                `json:"required"`
-		} `json:"parameters"`
-	} `json:"function"`
+	Type     string       `json:"type"`
+	Function ToolFunction `json:"function"`
 }
 
 // ToolProperty represents a property of a tool's parameter

v2/tools_test.go

@@ -20,7 +20,9 @@ func TestTools(t *testing.T) {
 
 	oc.SetSystemPrompt("You are a helpful assistant.")
 	oc.SetRandom()
-	oc.SetTool(json.RawMessage(`{
+
+	var toolGetCurrentWeather Tool
+	json.Unmarshal(json.RawMessage(`{
 		"type": "function",
 		"function": {
 		  "name": "get_current_weather",
@@ -41,7 +43,32 @@ func TestTools(t *testing.T) {
 			"required": ["location", "format"]
 		  }
 		}
-	  }`))
+	  }`), &toolGetCurrentWeather)
+
+	// toolGetCurrentWeather := Tool{
+	// 	Type: "function",
+	// 	Function: ToolFunction{
+	// 		Name:        "get_current_weather",
+	// 		Description: "Get the current weather for a location",
+	// 		Parameters: ToolParameters{
+	// 			Type: "object",
+	// 			Properties: map[string]ToolProperty{
+	// 				"location": {
+	// 					Type:        "string",
+	// 					Description: "The location to get the weather for, e.g. San Francisco, CA",
+	// 				},
+	// 				"format": {
+	// 					Type:        "string",
+	// 					Description: "The format to return the weather in, e.g. 'celsius' or 'fahrenheit'",
+	// 					Enum:        []string{"celsius", "fahrenheit"},
+	// 				},
+	// 			},
+	// 			Required: []string{"location", "format"},
+	// 		},
+	// 	},
+	// }
+
+	oc.SetTool(toolGetCurrentWeather)
 
 	prompt := "What is the weather in Toronto?"
 	generatedOutput := oc.MustOutputChat(prompt)