readysite / pkg / assistant / tool.go
3.7 KB
tool.go
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.
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.
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.
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
}
← Back