318 lines
7.4 KiB
Go
318 lines
7.4 KiB
Go
package basic
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
// maxExprLen bounds an expression length so a request cannot force a parse/scan
|
|
// beyond a trivial cost and cannot smuggle a larger payload through the tool.
|
|
const maxExprLen = 256
|
|
|
|
// allowedChars lists the bytes permitted in an expression besides ASCII letters
|
|
// (letters are only admitted to spell the sqrt function; the lexer rejects any
|
|
// other identifier).
|
|
const allowedChars = "0123456789.+-*/^() \t\n"
|
|
|
|
// calculator evaluates a safe arithmetic expression. It recognizes signed numeric
|
|
// literals, + - * /, parentheses, the right-associative ^ power operator and the
|
|
// sqrt(x) function. Every other character or identifier is rejected, so the
|
|
// expression can never run arbitrary code. Division by zero and non-real results
|
|
// are errors.
|
|
type calculator struct{}
|
|
|
|
// Name returns the tool identifier.
|
|
func (calculator) Name() string { return "calculator" }
|
|
|
|
// Description explains what the tool does for the model.
|
|
func (calculator) Description() string {
|
|
return "Safely evaluate an arithmetic expression using + - * /, parentheses, ^ (power) and sqrt(x)."
|
|
}
|
|
|
|
// Parameters returns the tool input JSON Schema.
|
|
func (calculator) Parameters() (map[string]any, error) {
|
|
return map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"expression": map[string]any{"type": "string", "description": "the arithmetic expression to evaluate"},
|
|
},
|
|
"required": []any{"expression"},
|
|
}, nil
|
|
}
|
|
|
|
// Execute evaluates expression and returns {"expression":..., "result":...}.
|
|
func (calculator) 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
|
|
}
|
|
expr, err := reqString(m, "expression")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
val, err := eval(expr)
|
|
if err != nil {
|
|
log().Error("calculator failed", "expression", expr, "reason", err.Error())
|
|
return nil, err
|
|
}
|
|
return map[string]any{"expression": expr, "result": val}, nil
|
|
}
|
|
|
|
// eval validates, tokenizes and parses expr, returning its numeric value.
|
|
func eval(expr string) (float64, error) {
|
|
if expr == "" {
|
|
return 0, fmt.Errorf("%w: expression is empty", ErrEvaluation)
|
|
}
|
|
if len(expr) > maxExprLen {
|
|
return 0, fmt.Errorf("%w: expression exceeds %d characters", ErrEvaluation, maxExprLen)
|
|
}
|
|
for _, c := range expr {
|
|
if strings.ContainsRune(allowedChars, c) {
|
|
continue
|
|
}
|
|
if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' {
|
|
continue // letter: the lexer admits only "sqrt"
|
|
}
|
|
return 0, fmt.Errorf("%w: disallowed character %q", ErrEvaluation, c)
|
|
}
|
|
toks, err := lex(expr)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
p := &parser{toks: toks}
|
|
return p.parse()
|
|
}
|
|
|
|
func isAsciiLetter(c byte) bool {
|
|
return c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z'
|
|
}
|
|
|
|
// tokKind describes a token category produced by the lexer.
|
|
type tokKind int
|
|
|
|
const (
|
|
tokEOF tokKind = iota
|
|
tokNum
|
|
tokOp
|
|
tokLParen
|
|
tokRParen
|
|
tokIdent
|
|
)
|
|
|
|
type token struct {
|
|
kind tokKind
|
|
text string
|
|
val float64
|
|
}
|
|
|
|
// lex scans s into a token stream terminated by tokEOF. Any ASCII letter run that
|
|
// is not exactly "sqrt" is rejected.
|
|
func lex(s string) ([]token, error) {
|
|
toks := make([]token, 0, len(s)/2+1)
|
|
for i := 0; i < len(s); {
|
|
c := s[i]
|
|
switch {
|
|
case c == ' ' || c == '\t' || c == '\n':
|
|
i++
|
|
case c == '(':
|
|
toks = append(toks, token{kind: tokLParen, text: "("})
|
|
i++
|
|
case c == ')':
|
|
toks = append(toks, token{kind: tokRParen, text: ")"})
|
|
i++
|
|
case c == '+' || c == '-' || c == '*' || c == '/' || c == '^':
|
|
toks = append(toks, token{kind: tokOp, text: string(c)})
|
|
i++
|
|
case isAsciiLetter(c):
|
|
start := i
|
|
for i < len(s) && isAsciiLetter(s[i]) {
|
|
i++
|
|
}
|
|
name := s[start:i]
|
|
if name != "sqrt" {
|
|
return nil, fmt.Errorf("%w: unexpected identifier %q", ErrEvaluation, name)
|
|
}
|
|
toks = append(toks, token{kind: tokIdent, text: name})
|
|
case c >= '0' && c <= '9' || c == '.':
|
|
start := i
|
|
dot := false
|
|
for i < len(s) {
|
|
d := s[i]
|
|
if d >= '0' && d <= '9' {
|
|
i++
|
|
continue
|
|
}
|
|
if d == '.' && !dot {
|
|
dot = true
|
|
i++
|
|
continue
|
|
}
|
|
break
|
|
}
|
|
raw := s[start:i]
|
|
val, err := strconv.ParseFloat(raw, 64)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: invalid number %q", ErrEvaluation, raw)
|
|
}
|
|
toks = append(toks, token{kind: tokNum, text: raw, val: val})
|
|
default:
|
|
return nil, fmt.Errorf("%w: disallowed character %q", ErrEvaluation, c)
|
|
}
|
|
}
|
|
toks = append(toks, token{kind: tokEOF, text: ""})
|
|
return toks, nil
|
|
}
|
|
|
|
// parser is a recursive-descent parser over the token stream. Grammar:
|
|
//
|
|
// expr := term (('+'|'-') term)*
|
|
// term := unary (('*'|'/') unary)*
|
|
// unary := ('+'|'-') unary | power
|
|
// power := primary ('^' unary)?
|
|
//
|
|
// ^ binds tighter than unary minus (so -2^2 == -(2^2)) and is right-associative.
|
|
type parser struct {
|
|
toks []token
|
|
pos int
|
|
}
|
|
|
|
func (p *parser) cur() token { return p.toks[p.pos] }
|
|
|
|
func (p *parser) parse() (float64, error) {
|
|
v, err := p.expr()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if p.cur().kind != tokEOF {
|
|
return 0, fmt.Errorf("%w: unexpected token %q", ErrEvaluation, p.cur().text)
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
func (p *parser) expr() (float64, error) {
|
|
v, err := p.term()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
for p.cur().kind == tokOp && (p.cur().text == "+" || p.cur().text == "-") {
|
|
op := p.cur().text
|
|
p.pos++
|
|
r, err := p.term()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if op == "+" {
|
|
v += r
|
|
} else {
|
|
v -= r
|
|
}
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
func (p *parser) term() (float64, error) {
|
|
v, err := p.unary()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
for p.cur().kind == tokOp && (p.cur().text == "*" || p.cur().text == "/") {
|
|
op := p.cur().text
|
|
p.pos++
|
|
r, err := p.unary()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if op == "*" {
|
|
v *= r
|
|
} else {
|
|
if r == 0 {
|
|
return 0, fmt.Errorf("%w: division by zero", ErrEvaluation)
|
|
}
|
|
v /= r
|
|
}
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
func (p *parser) unary() (float64, error) {
|
|
if p.cur().kind == tokOp && (p.cur().text == "+" || p.cur().text == "-") {
|
|
neg := p.cur().text == "-"
|
|
p.pos++
|
|
v, err := p.unary()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if neg {
|
|
return -v, nil
|
|
}
|
|
return v, nil
|
|
}
|
|
return p.power()
|
|
}
|
|
|
|
func (p *parser) power() (float64, error) {
|
|
base, err := p.primary()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if p.cur().kind == tokOp && p.cur().text == "^" {
|
|
p.pos++
|
|
exp, err := p.unary()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
r := math.Pow(base, exp)
|
|
if math.IsNaN(r) {
|
|
return 0, fmt.Errorf("%w: non-real result", ErrEvaluation)
|
|
}
|
|
return r, nil
|
|
}
|
|
return base, nil
|
|
}
|
|
|
|
func (p *parser) primary() (float64, error) {
|
|
t := p.cur()
|
|
switch t.kind {
|
|
case tokNum:
|
|
p.pos++
|
|
return t.val, nil
|
|
case tokLParen:
|
|
p.pos++
|
|
v, err := p.expr()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if p.cur().kind != tokRParen {
|
|
return 0, fmt.Errorf("%w: missing closing parenthesis", ErrEvaluation)
|
|
}
|
|
p.pos++
|
|
return v, nil
|
|
case tokIdent:
|
|
p.pos++
|
|
if p.cur().kind != tokLParen {
|
|
return 0, fmt.Errorf("%w: sqrt requires parentheses", ErrEvaluation)
|
|
}
|
|
p.pos++
|
|
v, err := p.expr()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if p.cur().kind != tokRParen {
|
|
return 0, fmt.Errorf("%w: missing closing parenthesis in sqrt", ErrEvaluation)
|
|
}
|
|
p.pos++
|
|
if v < 0 {
|
|
return 0, fmt.Errorf("%w: sqrt of a negative number", ErrEvaluation)
|
|
}
|
|
return math.Sqrt(v), nil
|
|
default:
|
|
return 0, fmt.Errorf("%w: unexpected token %q", ErrEvaluation, t.text)
|
|
}
|
|
}
|