Files
enkly-ai/pkg/mathkit/solve.go
enkly 79ce8dc4c0
Some checks failed
Go CI Pipeline / build_and_test (push) Failing after 22m0s
Init Project
2026-10-30 18:06:55 +01:00

836 lines
22 KiB
Go

package mathkit
import (
"fmt"
"math"
"sort"
"strings"
)
// Solve dispatches a math problem by string operation name against a
// loosely-typed params map — exactly the shape a JSON object (an LLM
// tool-call argument blob, decoded with encoding/json into
// map[string]any) takes once decoded. It exists as one layer above
// the typed functions in linalg.go/calculus.go/statistics.go/
// optimize.go/ml.go specifically so a caller at the JSON boundary
// (services/agent/internal/infrastructure/toolexec/mlmath.go) has one
// function to call and one error type to handle, instead of a type
// switch per operation.
//
// Returns a human-readable report (not raw JSON or a bare float) on
// success, because the caller here is ultimately an LLM deciding what
// to do next — a labeled, explained answer is directly usable in its
// response, where a bare number would have to be re-explained.
func Solve(operation string, params map[string]any) (string, error) {
op := strings.ToLower(strings.TrimSpace(operation))
switch op {
// --- general ---
case "evaluate":
return solveEvaluate(params)
// --- linear algebra ---
case "matrix_add", "matrix_sub", "matrix_mul":
return solveMatrixBinOp(op, params)
case "matrix_transpose":
return solveMatrixTranspose(params)
case "matrix_determinant":
return solveMatrixDeterminant(params)
case "matrix_inverse":
return solveMatrixInverse(params)
case "matrix_rank":
return solveMatrixRank(params)
case "matrix_trace":
return solveMatrixTrace(params)
case "dot_product":
return solveDotProduct(params)
case "vector_norm":
return solveVectorNorm(params)
case "dominant_eigen":
return solveDominantEigen(params)
// --- calculus ---
case "derivative":
return solveDerivative(params)
case "gradient":
return solveGradient(params)
case "integral":
return solveIntegral(params)
// --- optimization ---
case "gradient_descent":
return solveGradientDescent(params)
case "newton_root":
return solveNewtonRoot(params)
// --- statistics & probability ---
case "mean":
return solveReduce1(params, "mean", func(d []float64) (float64, error) { return Mean(d) })
case "variance":
return solveReduce1Pop(params, "variance", Variance)
case "std_dev":
return solveReduce1Pop(params, "std_dev", StdDev)
case "median":
return solveReduce1(params, "median", Median)
case "covariance":
return solveReduce2Pop(params, "covariance", Covariance)
case "correlation":
return solveReduce2(params, "correlation", Correlation)
case "z_score_normalize":
return solveNormalize(params, "z_score_normalize", ZScoreNormalize)
case "min_max_normalize":
return solveNormalize(params, "min_max_normalize", MinMaxNormalize)
case "normal_pdf":
return solveNormalDist(params, "normal_pdf", NormalPDF)
case "normal_cdf":
return solveNormalDist(params, "normal_cdf", NormalCDF)
case "binomial_pmf":
return solveBinomialPMF(params)
// --- ML primitives ---
case "sigmoid":
return solveElementwise(params, "sigmoid", Sigmoid)
case "relu":
return solveElementwise(params, "relu", ReLU)
case "tanh":
return solveElementwise(params, "tanh", math.Tanh)
case "softmax":
return solveSoftmax(params)
case "mse_loss":
return solveLoss(params, "mse_loss", MSELoss)
case "cross_entropy_loss":
return solveLoss(params, "cross_entropy_loss", CrossEntropyLoss)
case "linear_regression":
return solveLinearRegression(params)
default:
return "", fmt.Errorf("unknown operation %q — see the tool description for the full list of supported operations", operation)
}
}
// ---------- param extraction helpers ----------
//
// Every helper below reports a specific, actionable error (which
// field was wrong and why) rather than a generic type-assertion
// panic recovery message, because these errors are fed straight back
// to the calling model as the tool result — see toolexec's rpcFail /
// IsError convention — and "field X must be a number, got a string"
// is something the model can act on; "interface conversion failed"
// is not.
func reqFloat(params map[string]any, key string) (float64, error) {
v, ok := params[key]
if !ok {
return 0, fmt.Errorf("missing required field %q", key)
}
f, ok := toFloat(v)
if !ok {
return 0, fmt.Errorf("field %q must be a number", key)
}
return f, nil
}
func optFloat(params map[string]any, key string, fallback float64) float64 {
v, ok := params[key]
if !ok {
return fallback
}
if f, ok := toFloat(v); ok {
return f
}
return fallback
}
func optInt(params map[string]any, key string, fallback int) int {
v, ok := params[key]
if !ok {
return fallback
}
if f, ok := toFloat(v); ok {
return int(f)
}
return fallback
}
func optBool(params map[string]any, key string, fallback bool) bool {
v, ok := params[key]
if !ok {
return fallback
}
if b, ok := v.(bool); ok {
return b
}
return fallback
}
func toFloat(v any) (float64, bool) {
switch n := v.(type) {
case float64:
return n, true
case int:
return float64(n), true
case int64:
return float64(n), true
}
return 0, false
}
func reqString(params map[string]any, key string) (string, error) {
v, ok := params[key]
if !ok {
return "", fmt.Errorf("missing required field %q", key)
}
s, ok := v.(string)
if !ok || strings.TrimSpace(s) == "" {
return "", fmt.Errorf("field %q must be a non-empty string", key)
}
return s, nil
}
func optString(params map[string]any, key, fallback string) string {
v, ok := params[key]
if !ok {
return fallback
}
s, ok := v.(string)
if !ok {
return fallback
}
return s
}
func reqFloatSlice(params map[string]any, key string) ([]float64, error) {
v, ok := params[key]
if !ok {
return nil, fmt.Errorf("missing required field %q (expected an array of numbers)", key)
}
raw, ok := v.([]any)
if !ok {
return nil, fmt.Errorf("field %q must be an array of numbers", key)
}
out := make([]float64, len(raw))
for i, x := range raw {
f, ok := toFloat(x)
if !ok {
return nil, fmt.Errorf("field %q: element %d is not a number", key, i)
}
out[i] = f
}
return out, nil
}
func reqStringSlice(params map[string]any, key string) ([]string, error) {
v, ok := params[key]
if !ok {
return nil, fmt.Errorf("missing required field %q (expected an array of strings)", key)
}
raw, ok := v.([]any)
if !ok {
return nil, fmt.Errorf("field %q must be an array of strings", key)
}
out := make([]string, len(raw))
for i, x := range raw {
s, ok := x.(string)
if !ok {
return nil, fmt.Errorf("field %q: element %d is not a string", key, i)
}
out[i] = s
}
return out, nil
}
func reqMatrix(params map[string]any, key string) (Matrix, error) {
v, ok := params[key]
if !ok {
return nil, fmt.Errorf("missing required field %q (expected a 2D array of numbers)", key)
}
rows, ok := v.([]any)
if !ok {
return nil, fmt.Errorf("field %q must be a 2D array of numbers, e.g. [[1,2],[3,4]]", key)
}
m := make(Matrix, len(rows))
for i, r := range rows {
rowRaw, ok := r.([]any)
if !ok {
return nil, fmt.Errorf("field %q: row %d is not an array", key, i)
}
row := make([]float64, len(rowRaw))
for j, x := range rowRaw {
f, ok := toFloat(x)
if !ok {
return nil, fmt.Errorf("field %q: row %d, column %d is not a number", key, i, j)
}
row[j] = f
}
m[i] = row
}
return m, nil
}
// varMap builds the variable bindings an expression needs from an
// "at" (or "variables") object field — {"x": 1, "y": 2} — falling
// back to an empty map if the field is absent (a single-variable
// expression evaluated via the "variable"/"at" pair still works).
func varMap(params map[string]any, key string) (map[string]float64, error) {
v, ok := params[key]
if !ok {
return map[string]float64{}, nil
}
raw, ok := v.(map[string]any)
if !ok {
return nil, fmt.Errorf("field %q must be an object mapping variable names to numbers, e.g. {\"x\": 1.5}", key)
}
out := make(map[string]float64, len(raw))
for k, x := range raw {
f, ok := toFloat(x)
if !ok {
return nil, fmt.Errorf("field %q: value for %q is not a number", key, k)
}
out[k] = f
}
return out, nil
}
func parseExprField(params map[string]any) (*Expr, error) {
src, err := reqString(params, "expression")
if err != nil {
return nil, err
}
expr, err := ParseExpr(src)
if err != nil {
return nil, err
}
return expr, nil
}
func fmtFloat(f float64) string {
if math.IsNaN(f) {
return "NaN"
}
if math.IsInf(f, 1) {
return "+Inf"
}
if math.IsInf(f, -1) {
return "-Inf"
}
return strings.TrimRight(strings.TrimRight(fmt.Sprintf("%.10f", f), "0"), ".")
}
func fmtFloats(fs []float64) string {
parts := make([]string, len(fs))
for i, f := range fs {
parts[i] = fmtFloat(f)
}
return "[" + strings.Join(parts, ", ") + "]"
}
func fmtMatrix(m Matrix) string {
var sb strings.Builder
for _, row := range m {
sb.WriteString(fmtFloats(row))
sb.WriteByte('\n')
}
return strings.TrimRight(sb.String(), "\n")
}
// ---------- operation implementations ----------
func solveEvaluate(params map[string]any) (string, error) {
expr, err := parseExprField(params)
if err != nil {
return "", err
}
vars, err := varMap(params, "variables")
if err != nil {
return "", err
}
result, err := expr.Eval(vars)
if err != nil {
return "", err
}
return fmt.Sprintf("%s = %s", expr.String(), fmtFloat(result)), nil
}
func solveMatrixBinOp(op string, params map[string]any) (string, error) {
a, err := reqMatrix(params, "a")
if err != nil {
return "", err
}
b, err := reqMatrix(params, "b")
if err != nil {
return "", err
}
var out Matrix
switch op {
case "matrix_add":
out, err = MatAdd(a, b)
case "matrix_sub":
out, err = MatSub(a, b)
case "matrix_mul":
out, err = MatMul(a, b)
}
if err != nil {
return "", err
}
return fmt.Sprintf("%s result:\n%s", op, fmtMatrix(out)), nil
}
func solveMatrixTranspose(params map[string]any) (string, error) {
a, err := reqMatrix(params, "a")
if err != nil {
return "", err
}
out, err := Transpose(a)
if err != nil {
return "", err
}
return fmt.Sprintf("transpose result:\n%s", fmtMatrix(out)), nil
}
func solveMatrixDeterminant(params map[string]any) (string, error) {
a, err := reqMatrix(params, "a")
if err != nil {
return "", err
}
det, err := Determinant(a)
if err != nil {
return "", err
}
return fmt.Sprintf("det(A) = %s", fmtFloat(det)), nil
}
func solveMatrixInverse(params map[string]any) (string, error) {
a, err := reqMatrix(params, "a")
if err != nil {
return "", err
}
inv, err := Inverse(a)
if err != nil {
return "", err
}
return fmt.Sprintf("A^-1:\n%s", fmtMatrix(inv)), nil
}
func solveMatrixRank(params map[string]any) (string, error) {
a, err := reqMatrix(params, "a")
if err != nil {
return "", err
}
r, err := Rank(a)
if err != nil {
return "", err
}
return fmt.Sprintf("rank(A) = %d", r), nil
}
func solveMatrixTrace(params map[string]any) (string, error) {
a, err := reqMatrix(params, "a")
if err != nil {
return "", err
}
tr, err := Trace(a)
if err != nil {
return "", err
}
return fmt.Sprintf("trace(A) = %s", fmtFloat(tr)), nil
}
func solveDotProduct(params map[string]any) (string, error) {
a, err := reqFloatSlice(params, "a")
if err != nil {
return "", err
}
b, err := reqFloatSlice(params, "b")
if err != nil {
return "", err
}
d, err := Dot(a, b)
if err != nil {
return "", err
}
return fmt.Sprintf("a · b = %s", fmtFloat(d)), nil
}
func solveVectorNorm(params map[string]any) (string, error) {
a, err := reqFloatSlice(params, "a")
if err != nil {
return "", err
}
p := optFloat(params, "p", 2)
n := Norm(a, p)
label := fmt.Sprintf("L%g", p)
if math.IsInf(p, 1) {
label = "L∞"
}
return fmt.Sprintf("%s norm = %s", label, fmtFloat(n)), nil
}
func solveDominantEigen(params map[string]any) (string, error) {
a, err := reqMatrix(params, "a")
if err != nil {
return "", err
}
iters := optInt(params, "iterations", 0)
lambda, v, err := DominantEigen(a, iters)
if err != nil {
return "", err
}
return fmt.Sprintf("dominant eigenvalue ≈ %s\neigenvector ≈ %s", fmtFloat(lambda), fmtFloats(v)), nil
}
func solveDerivative(params map[string]any) (string, error) {
expr, err := parseExprField(params)
if err != nil {
return "", err
}
variable, err := reqString(params, "variable")
if err != nil {
return "", err
}
at, err := varMap(params, "at")
if err != nil {
return "", err
}
if _, ok := at[variable]; !ok {
at[variable] = optFloat(params, "value", 0)
}
h := optFloat(params, "h", 0)
d, err := Derivative(expr, variable, at, h)
if err != nil {
return "", err
}
return fmt.Sprintf("d/d%s[%s] at %s=%s ≈ %s", variable, expr.String(), variable, fmtFloat(at[variable]), fmtFloat(d)), nil
}
func solveGradient(params map[string]any) (string, error) {
expr, err := parseExprField(params)
if err != nil {
return "", err
}
variables, err := reqStringSlice(params, "variables")
if err != nil {
return "", err
}
at, err := varMap(params, "at")
if err != nil {
return "", err
}
h := optFloat(params, "h", 0)
grad, err := Gradient(expr, variables, at, h)
if err != nil {
return "", err
}
return fmt.Sprintf("∇[%s] with respect to %s ≈ %s", expr.String(), strings.Join(variables, ","), fmtFloats(grad)), nil
}
func solveIntegral(params map[string]any) (string, error) {
expr, err := parseExprField(params)
if err != nil {
return "", err
}
variable, err := reqString(params, "variable")
if err != nil {
return "", err
}
from, err := reqFloat(params, "from")
if err != nil {
return "", err
}
to, err := reqFloat(params, "to")
if err != nil {
return "", err
}
steps := optInt(params, "steps", 0)
fixed, err := varMap(params, "variables")
if err != nil {
return "", err
}
result, err := Integrate(expr, variable, from, to, steps, fixed)
if err != nil {
return "", err
}
return fmt.Sprintf("∫[%s to %s] %s d%s ≈ %s", fmtFloat(from), fmtFloat(to), expr.String(), variable, fmtFloat(result)), nil
}
func solveGradientDescent(params map[string]any) (string, error) {
expr, err := parseExprField(params)
if err != nil {
return "", err
}
variables, err := reqStringSlice(params, "variables")
if err != nil {
return "", err
}
start, err := reqFloatSlice(params, "start")
if err != nil {
return "", err
}
lr, err := reqFloat(params, "learning_rate")
if err != nil {
return "", err
}
iterations := optInt(params, "iterations", 100)
tol := optFloat(params, "tolerance", 1e-8)
res, err := GradientDescent(expr, variables, start, lr, iterations, tol)
if err != nil {
return "", err
}
last := len(res.Path) - 1
var sb strings.Builder
fmt.Fprintf(&sb, "gradient descent on f(%s) = %s\n", strings.Join(variables, ","), expr.String())
fmt.Fprintf(&sb, "start: %s -> f = %s\n", fmtFloats(res.Path[0]), fmtFloat(res.Values[0]))
fmt.Fprintf(&sb, "ran %d step(s), converged=%v\n", last, res.Converged)
fmt.Fprintf(&sb, "final: %s -> f = %s\n", fmtFloats(res.Path[last]), fmtFloat(res.Values[last]))
// Show a short trajectory sample so the model can see the descent
// shape without a huge wall of numbers for large iteration counts.
sampleN := 5
if last+1 <= sampleN*2 {
for i := 0; i <= last; i++ {
fmt.Fprintf(&sb, " step %d: x=%s f=%s\n", i, fmtFloats(res.Path[i]), fmtFloat(res.Values[i]))
}
} else {
for i := 0; i < sampleN; i++ {
fmt.Fprintf(&sb, " step %d: x=%s f=%s\n", i, fmtFloats(res.Path[i]), fmtFloat(res.Values[i]))
}
sb.WriteString(" ...\n")
for i := last - sampleN + 1; i <= last; i++ {
fmt.Fprintf(&sb, " step %d: x=%s f=%s\n", i, fmtFloats(res.Path[i]), fmtFloat(res.Values[i]))
}
}
return strings.TrimRight(sb.String(), "\n"), nil
}
func solveNewtonRoot(params map[string]any) (string, error) {
expr, err := parseExprField(params)
if err != nil {
return "", err
}
variable, err := reqString(params, "variable")
if err != nil {
return "", err
}
start, err := reqFloat(params, "start")
if err != nil {
return "", err
}
iterations := optInt(params, "iterations", 0)
tol := optFloat(params, "tolerance", 0)
fixed, err := varMap(params, "variables")
if err != nil {
return "", err
}
res, err := NewtonRoot(expr, variable, start, iterations, tol, fixed)
if err != nil {
return "", err
}
last := res.Path[len(res.Path)-1]
return fmt.Sprintf("newton's method on %s(%s)=0: converged=%v after %d step(s), root ≈ %s=%s",
expr.String(), variable, res.Converged, len(res.Path)-1, variable, fmtFloat(last)), nil
}
func solveReduce1(params map[string]any, name string, f func([]float64) (float64, error)) (string, error) {
data, err := reqFloatSlice(params, "data")
if err != nil {
return "", err
}
v, err := f(data)
if err != nil {
return "", err
}
return fmt.Sprintf("%s(%s) = %s", name, fmtFloats(data), fmtFloat(v)), nil
}
func solveReduce1Pop(params map[string]any, name string, f func([]float64, bool) (float64, error)) (string, error) {
data, err := reqFloatSlice(params, "data")
if err != nil {
return "", err
}
population := optBool(params, "population", false)
v, err := f(data, population)
if err != nil {
return "", err
}
kind := "sample"
if population {
kind = "population"
}
return fmt.Sprintf("%s %s(%s) = %s", kind, name, fmtFloats(data), fmtFloat(v)), nil
}
func solveReduce2(params map[string]any, name string, f func([]float64, []float64) (float64, error)) (string, error) {
x, err := reqFloatSlice(params, "x")
if err != nil {
return "", err
}
y, err := reqFloatSlice(params, "y")
if err != nil {
return "", err
}
v, err := f(x, y)
if err != nil {
return "", err
}
return fmt.Sprintf("%s(x, y) = %s", name, fmtFloat(v)), nil
}
func solveReduce2Pop(params map[string]any, name string, f func([]float64, []float64, bool) (float64, error)) (string, error) {
x, err := reqFloatSlice(params, "x")
if err != nil {
return "", err
}
y, err := reqFloatSlice(params, "y")
if err != nil {
return "", err
}
population := optBool(params, "population", false)
v, err := f(x, y, population)
if err != nil {
return "", err
}
kind := "sample"
if population {
kind = "population"
}
return fmt.Sprintf("%s %s(x, y) = %s", kind, name, fmtFloat(v)), nil
}
func solveNormalize(params map[string]any, name string, f func([]float64) ([]float64, error)) (string, error) {
data, err := reqFloatSlice(params, "data")
if err != nil {
return "", err
}
out, err := f(data)
if err != nil {
return "", err
}
return fmt.Sprintf("%s(%s) = %s", name, fmtFloats(data), fmtFloats(out)), nil
}
func solveNormalDist(params map[string]any, name string, f func(float64, float64, float64) (float64, error)) (string, error) {
x, err := reqFloat(params, "x")
if err != nil {
return "", err
}
mean := optFloat(params, "mean", 0)
std := optFloat(params, "std_dev", 1)
v, err := f(x, mean, std)
if err != nil {
return "", err
}
return fmt.Sprintf("%s(x=%s, mean=%s, std_dev=%s) = %s", name, fmtFloat(x), fmtFloat(mean), fmtFloat(std), fmtFloat(v)), nil
}
func solveBinomialPMF(params map[string]any) (string, error) {
k := optInt(params, "k", -1)
if k < 0 {
return "", fmt.Errorf("missing required field %q", "k")
}
n := optInt(params, "n", -1)
if n < 0 {
return "", fmt.Errorf("missing required field %q", "n")
}
p, err := reqFloat(params, "p")
if err != nil {
return "", err
}
v, err := BinomialPMF(k, n, p)
if err != nil {
return "", err
}
return fmt.Sprintf("P(X=%d) for X~Binomial(n=%d, p=%s) = %s", k, n, fmtFloat(p), fmtFloat(v)), nil
}
func solveElementwise(params map[string]any, name string, f func(float64) float64) (string, error) {
if raw, ok := params["x"].([]any); ok {
data := make([]float64, len(raw))
for i, v := range raw {
fv, ok := toFloat(v)
if !ok {
return "", fmt.Errorf("field \"x\": element %d is not a number", i)
}
data[i] = fv
}
out := make([]float64, len(data))
for i, v := range data {
out[i] = f(v)
}
return fmt.Sprintf("%s(%s) = %s", name, fmtFloats(data), fmtFloats(out)), nil
}
x, err := reqFloat(params, "x")
if err != nil {
return "", fmt.Errorf(`field "x" must be a number or an array of numbers: %w`, err)
}
return fmt.Sprintf("%s(%s) = %s", name, fmtFloat(x), fmtFloat(f(x))), nil
}
func solveSoftmax(params map[string]any) (string, error) {
x, err := reqFloatSlice(params, "x")
if err != nil {
return "", err
}
out := Softmax(x)
sum := 0.0
for _, v := range out {
sum += v
}
return fmt.Sprintf("softmax(%s) = %s (sums to %s)", fmtFloats(x), fmtFloats(out), fmtFloat(sum)), nil
}
func solveLoss(params map[string]any, name string, f func([]float64, []float64) (float64, error)) (string, error) {
pred, err := reqFloatSlice(params, "predictions")
if err != nil {
return "", err
}
target, err := reqFloatSlice(params, "targets")
if err != nil {
return "", err
}
v, err := f(pred, target)
if err != nil {
return "", err
}
return fmt.Sprintf("%s(predictions, targets) = %s", name, fmtFloat(v)), nil
}
func solveLinearRegression(params map[string]any) (string, error) {
X, err := reqMatrix(params, "x")
if err != nil {
return "", err
}
y, err := reqFloatSlice(params, "y")
if err != nil {
return "", err
}
res, err := LinearRegression(X, y)
if err != nil {
return "", err
}
var sb strings.Builder
fmt.Fprintf(&sb, "ordinary least squares fit: y ≈ %s", fmtFloat(res.Coefficients[0]))
for i, c := range res.Coefficients[1:] {
fmt.Fprintf(&sb, " + %s*x%d", fmtFloat(c), i+1)
}
fmt.Fprintf(&sb, "\nR² = %s\n", fmtFloat(res.RSquared))
fmt.Fprintf(&sb, "predictions: %s\n", fmtFloats(res.Predictions))
fmt.Fprintf(&sb, "residuals: %s", fmtFloats(res.Residuals))
return sb.String(), nil
}
// Operations returns the sorted list of every operation name Solve
// accepts — used by the agent's tool-catalog description so the
// supported operation list is generated from the real dispatch table
// instead of hand-copied prose that can drift out of sync with it.
func Operations() []string {
ops := []string{
"evaluate",
"matrix_add", "matrix_sub", "matrix_mul", "matrix_transpose", "matrix_determinant",
"matrix_inverse", "matrix_rank", "matrix_trace", "dot_product", "vector_norm", "dominant_eigen",
"derivative", "gradient", "integral",
"gradient_descent", "newton_root",
"mean", "variance", "std_dev", "median", "covariance", "correlation",
"z_score_normalize", "min_max_normalize", "normal_pdf", "normal_cdf", "binomial_pmf",
"sigmoid", "relu", "tanh", "softmax", "mse_loss", "cross_entropy_loss", "linear_regression",
}
sort.Strings(ops)
return ops
}