126 lines
3.6 KiB
Go
126 lines
3.6 KiB
Go
package basic
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"fmt"
|
|
"math/big"
|
|
"strings"
|
|
)
|
|
|
|
// random generates a cryptographically-random integer in [min,max] or a random
|
|
// string of a given length drawn from a charset. Providing min/max selects the
|
|
// integer mode; providing length selects the string mode.
|
|
type random struct{}
|
|
|
|
const (
|
|
// maxRandomLen bounds a generated string length.
|
|
maxRandomLen = 1024
|
|
// defaultCharset is the alphabet used when none is supplied.
|
|
defaultCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
|
)
|
|
|
|
// Name returns the tool identifier.
|
|
func (random) Name() string { return "random" }
|
|
|
|
// Description explains what the tool does for the model.
|
|
func (random) Description() string {
|
|
return "Generate a cryptographically-random integer in [min,max] or a random string of a given length."
|
|
}
|
|
|
|
// Parameters returns the tool input JSON Schema.
|
|
func (random) Parameters() (map[string]any, error) {
|
|
return map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"min": map[string]any{"type": "integer", "description": "inclusive lower bound for an integer"},
|
|
"max": map[string]any{"type": "integer", "description": "inclusive upper bound for an integer"},
|
|
"length": map[string]any{"type": "integer", "description": "length of the random string"},
|
|
"charset": map[string]any{"type": "string", "description": "characters the string is drawn from"},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// Execute produces either an integer or a string depending on which arguments
|
|
// are supplied.
|
|
func (random) Execute(ctx context.Context, args any) (any, error) {
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
m, err := asMap(args)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if hasKey(m, "min") || hasKey(m, "max") {
|
|
min, err := optInt(m, "min", 0)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
max, err := optInt(m, "max", 0)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
v, err := randInt(min, max)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return map[string]any{"value": v}, nil
|
|
}
|
|
if hasKey(m, "length") {
|
|
n, err := reqInt(m, "length")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
charset, err := optString(m, "charset", defaultCharset)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
s, err := randString(n, charset)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return map[string]any{"value": s}, nil
|
|
}
|
|
return nil, fmt.Errorf("basic: provide min/max for an integer or length for a string: %w", ErrInvalidArgument)
|
|
}
|
|
|
|
// randInt returns a uniformly-distributed integer in [min,max] inclusive, using
|
|
// crypto/rand so the result is not predictable from the process state.
|
|
func randInt(min, max int) (int, error) {
|
|
if max < min {
|
|
return 0, fmt.Errorf("basic: max must be >= min: %w", ErrInvalidArgument)
|
|
}
|
|
if min == max {
|
|
return min, nil
|
|
}
|
|
span := int64(max) - int64(min) + 1
|
|
n, err := rand.Int(rand.Reader, big.NewInt(span))
|
|
if err != nil {
|
|
return 0, fmt.Errorf("basic: random source: %w", err)
|
|
}
|
|
return int(n.Int64()) + min, nil
|
|
}
|
|
|
|
// randString returns a length-character string drawn uniformly from charset using
|
|
// crypto/rand (no modulo bias).
|
|
func randString(length int, charset string) (string, error) {
|
|
if length < 1 || length > maxRandomLen {
|
|
return "", fmt.Errorf("basic: length must be between 1 and %d: %w", maxRandomLen, ErrInvalidArgument)
|
|
}
|
|
cs := []rune(charset)
|
|
if len(cs) == 0 {
|
|
return "", fmt.Errorf("basic: charset must not be empty: %w", ErrInvalidArgument)
|
|
}
|
|
var b strings.Builder
|
|
b.Grow(length)
|
|
span := big.NewInt(int64(len(cs)))
|
|
for i := 0; i < length; i++ {
|
|
n, err := rand.Int(rand.Reader, span)
|
|
if err != nil {
|
|
return "", fmt.Errorf("basic: random source: %w", err)
|
|
}
|
|
b.WriteRune(cs[n.Int64()])
|
|
}
|
|
return b.String(), nil
|
|
}
|