tool.go

140 lines
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
package assistant

import "encoding/json"

// Tool represents a function the model can call.
type Tool struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	Parameters  *Parameters `json:"parameters,omitempty"`
}

// Parameters describes the JSON schema for tool parameters.
// This is provider-facing: used by provider subpackages to serialize tool definitions.
type Parameters struct {
	Type       string              `json:"type"` // Always "object"
	Properties map[string]Property `json:"properties,omitempty"`
	Required   []string            `json:"required,omitempty"`
}

// Property describes a single parameter.
// This is provider-facing: used by provider subpackages to serialize tool definitions.
type Property struct {
	Type        string   `json:"type"` // "string", "number", "boolean", "array", "object"
	Description string   `json:"description,omitempty"`
	Enum        []string `json:"enum,omitempty"`  // For string enums
	Items       *Items   `json:"items,omitempty"` // For arrays
}

// Items describes array item type.
// This is provider-facing: used by provider subpackages to serialize tool definitions.
type Items struct {
	Type string `json:"type"`
}

// ToolCall represents a tool invocation requested by the model.
type ToolCall struct {
	ID        string `json:"id"`
	Name      string `json:"name"`
	Arguments string `json:"arguments"` // JSON string
}

// ParseArguments parses the JSON arguments into the provided struct.
func (tc *ToolCall) ParseArguments(v any) error {
	return json.Unmarshal([]byte(tc.Arguments), v)
}

// --- ToolBuilder for fluent API ---

// ToolBuilder provides a fluent interface for building tools.
type ToolBuilder struct {
	name        string
	description string
	properties  map[string]Property
	required    []string
}

// NewTool creates a new ToolBuilder.
func NewTool(name, description string) *ToolBuilder {
	return &ToolBuilder{
		name:        name,
		description: description,
		properties:  make(map[string]Property),
	}
}

// String adds a string parameter.
func (b *ToolBuilder) String(name, description string, required bool) *ToolBuilder {
	b.properties[name] = Property{Type: "string", Description: description}
	if required {
		b.required = append(b.required, name)
	}
	return b
}

// Int adds an integer parameter.
func (b *ToolBuilder) Int(name, description string, required bool) *ToolBuilder {
	b.properties[name] = Property{Type: "integer", Description: description}
	if required {
		b.required = append(b.required, name)
	}
	return b
}

// Number adds a number parameter.
func (b *ToolBuilder) Number(name, description string, required bool) *ToolBuilder {
	b.properties[name] = Property{Type: "number", Description: description}
	if required {
		b.required = append(b.required, name)
	}
	return b
}

// Bool adds a boolean parameter.
func (b *ToolBuilder) Bool(name, description string, required bool) *ToolBuilder {
	b.properties[name] = Property{Type: "boolean", Description: description}
	if required {
		b.required = append(b.required, name)
	}
	return b
}

// Enum adds a string enum parameter.
func (b *ToolBuilder) Enum(name, description string, values []string, required bool) *ToolBuilder {
	b.properties[name] = Property{Type: "string", Description: description, Enum: values}
	if required {
		b.required = append(b.required, name)
	}
	return b
}

// Array adds an array parameter.
func (b *ToolBuilder) Array(name, description, itemType string, required bool) *ToolBuilder {
	b.properties[name] = Property{
		Type:        "array",
		Description: description,
		Items:       &Items{Type: itemType},
	}
	if required {
		b.required = append(b.required, name)
	}
	return b
}

// Build creates the Tool.
func (b *ToolBuilder) Build() Tool {
	tool := Tool{
		Name:        b.name,
		Description: b.description,
	}

	if len(b.properties) > 0 {
		tool.Parameters = &Parameters{
			Type:       "object",
			Properties: b.properties,
			Required:   b.required,
		}
	}

	return tool
}