Tool RAG: lexical projection of the bound tool catalog (off by default)

When ARTEX_TOOL_RAG is enabled, AugmentTools keeps the core tools plus the top-k (ARTEX_TOOL_RAG_TOPK, default 12) BM25-ranked bound tools for the worker intent / planner goal. Runs after ToolResolve, so it only removes; falls back to the full list on empty query, catalog under 20, or scorer error.
2026-11-13 11:40:53 +08:00 · 2026-11-12 15:42:31 +01:00
parent 34ece5937a
commit f3a20ef2c6
5 changed files with 356 additions and 2 deletions
@@ -60,6 +60,8 @@ func AugmentTools(ctx context.Context, agentKey string, base []actool.CoreTool)
 	if ToolResolve != nil {
 		out = ToolResolve(ctx, agentKey, out)
 	}
+	// Tool RAG (off by default): lexical projection of the bound list; only removes.
+	out = applyToolRAG(ctx, agentKey, out, def.Deferred, bm25Score)
 	out, def.FindingGuidance = findingWorkflowTools(agentKey, out)
 	for i, t := range out {
 		out[i] = guardPanic(t)
@@ -299,6 +299,18 @@ func factIDsYielded(ts *db.ExplorationStore, id int64) string {
 	return strings.Join(parts, ", ")
 }

+// plannerToolQuery is the Tool RAG query for one planning round: the goal plus the
+// short texts the round's triggers already carry (no store lookups).
+func plannerToolQuery(goal string, triggers []TriggerEvent) string {
+	parts := []string{goal}
+	for _, ev := range triggers {
+		parts = append(parts, ev.Summary, ev.NewGoal)
+		parts = append(parts, ev.Goals...)
+		parts = append(parts, ev.Hints...)
+	}
+	return strings.TrimSpace(strings.Join(parts, " "))
+}
+
 // intentSummary reads an intent node's one-line summary (best-effort, "?" on miss).
 func intentSummary(ts *db.ExplorationStore, id int64) string {
 	n, err := ts.GetNode(id)
@@ -470,7 +482,7 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts
 	// When the asset-coverage feature is off, drop add_task_scope/list_untested_assets (not in the prompt).
 	base := append(tsx.DropCoverageTools(tsx.PlannerTools()), actool.DefaultTools()...)
 	ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts)})
-	tools, def, cleanup := AugmentTools(ctx, "planner", base)
+	tools, def, cleanup := AugmentTools(WithToolQuery(ctx, plannerToolQuery(goal, triggers)), "planner", base)
 	defer cleanup()
 	// The key situation (just-finished intents + the pre-fetched full graph) goes into
 	// [this round's user input] (see input below); the system prompt keeps only the
@@ -0,0 +1,197 @@
+package agent
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"math"
+	"os"
+	"sort"
+	"strconv"
+	"strings"
+	"unicode"
+
+	actool "github.com/Autumn-27/norma/tool"
+)
+
+// Tool RAG: a lexical PROJECTION of an agent's already-bound tool list. It runs in
+// AugmentTools after ToolResolve (binding + enabled filters), so it only ever
+// removes tools from that list: it cannot add an unbound or disabled tool, never
+// rewrites a description or schema, and never reads the database. Off by default.
+
+const (
+	toolRAGEnv        = "ARTEX_TOOL_RAG"      // "1"/"true"/"on"/"yes" enables; anything else (incl. unset) = off
+	toolRAGTopKEnv    = "ARTEX_TOOL_RAG_TOPK" // non-core tools kept; default toolRAGDefaultK
+	toolRAGDefaultK   = 12
+	toolRAGMinCatalog = 20 // smaller catalogs are passed through whole
+)
+
+// toolRAGCore is always kept when bound and never counts against k.
+var toolRAGCore = map[string]bool{
+	"record_fact": true, "insert_assets": true, "report_finding": true, "list_findings": true,
+	"graph_overview": true, "node_detail": true, "add_intent": true, "steer_work": true,
+	"ask_operator": true, "prove_goal": true, "goal_met": true, "list_assets": true,
+}
+
+type toolQueryKey struct{}
+
+// WithToolQuery attaches the retrieval query (goal + current intent/round summary)
+// read by AugmentTools. No query = no projection.
+func WithToolQuery(ctx context.Context, q string) context.Context {
+	return context.WithValue(ctx, toolQueryKey{}, q)
+}
+
+// ToolScorer returns one relevance score per tool (same order as tools); higher is
+// more relevant. It must be pure: no network, no database.
+type ToolScorer func(query string, tools []actool.CoreTool) ([]float64, error)
+
+// projectTools keeps the tools in keep (core / harness / deferred names) plus the
+// k best-scoring other tools, preserving the input order. It falls back to the full
+// list (projected=false) on an empty query, a catalog under toolRAGMinCatalog, or a
+// scorer error.
+func projectTools(query string, tools []actool.CoreTool, keep map[string]bool, k int, score ToolScorer) (out []actool.CoreTool, projected bool, err error) {
+	if strings.TrimSpace(query) == "" || len(tools) < toolRAGMinCatalog {
+		return tools, false, nil
+	}
+	var cands []actool.CoreTool
+	for _, t := range tools {
+		if !keep[t.Name()] {
+			cands = append(cands, t)
+		}
+	}
+	ranked, err := rankTools(query, cands, score)
+	if err != nil {
+		return tools, false, err
+	}
+	sel := map[string]bool{}
+	for i := 0; i < len(ranked) && i < k; i++ {
+		sel[ranked[i].Name()] = true
+	}
+	for _, t := range tools {
+		if keep[t.Name()] || sel[t.Name()] {
+			out = append(out, t)
+		}
+	}
+	return out, true, nil
+}
+
+// rankTools orders tools by score descending, ties broken by name ascending.
+func rankTools(query string, tools []actool.CoreTool, score ToolScorer) ([]actool.CoreTool, error) {
+	scores, err := score(query, tools)
+	if err == nil && len(scores) != len(tools) {
+		err = fmt.Errorf("scorer returned %d scores for %d tools", len(scores), len(tools))
+	}
+	if err != nil {
+		return nil, err
+	}
+	idx := make([]int, len(tools))
+	for i := range idx {
+		idx[i] = i
+	}
+	sort.SliceStable(idx, func(a, b int) bool {
+		if sa, sb := scores[idx[a]], scores[idx[b]]; sa != sb {
+			return sa > sb
+		}
+		return tools[idx[a]].Name() < tools[idx[b]].Name()
+	})
+	out := make([]actool.CoreTool, len(idx))
+	for i, j := range idx {
+		out[i] = tools[j]
+	}
+	return out, nil
+}
+
+// bm25Score is the default ToolScorer: Okapi BM25 (k1=1.2, b=0.75) over each tool's
+// name + description + parameter names, IDF computed across the candidate set.
+func bm25Score(query string, tools []actool.CoreTool) ([]float64, error) {
+	const k1, b = 1.2, 0.75
+	docs := make([]map[string]int, len(tools))
+	lens := make([]float64, len(tools))
+	df := map[string]int{}
+	var total float64
+	for i, t := range tools {
+		text := t.Name() + " " + t.Description()
+		if props, ok := t.InputSchema()["properties"].(map[string]any); ok {
+			for p := range props {
+				text += " " + p
+			}
+		}
+		tf := map[string]int{}
+		for _, w := range tokenize(text) {
+			tf[w]++
+			lens[i]++
+		}
+		for w := range tf {
+			df[w]++
+		}
+		docs[i], total = tf, total+lens[i]
+	}
+	scores := make([]float64, len(tools))
+	n := float64(len(tools))
+	avg, seen := math.Max(total/math.Max(n, 1), 1), map[string]bool{}
+	for _, q := range tokenize(query) {
+		if seen[q] || df[q] == 0 {
+			continue
+		}
+		seen[q] = true
+		idf := math.Log(1 + (n-float64(df[q])+0.5)/(float64(df[q])+0.5))
+		for i, tf := range docs {
+			if f := float64(tf[q]); f > 0 {
+				scores[i] += idf * f * (k1 + 1) / (f + k1*(1-b+b*lens[i]/avg))
+			}
+		}
+	}
+	return scores, nil
+}
+
+// tokenize lowercases and splits on non-alphanumerics ("sql_injection" → sql, injection).
+func tokenize(s string) []string {
+	return strings.FieldsFunc(strings.ToLower(s), func(r rune) bool {
+		return !unicode.IsLetter(r) && !unicode.IsDigit(r)
+	})
+}
+
+// toolRAGKeep: names never hidden by the projection — the core set, the harness
+// tools every agent relies on (file/shell tools, interactive shell, Skill
+// meta-tool), and deferred tools (schemas already withheld; dropping them would
+// desync the deferred block and the skill unlock gate).
+func toolRAGKeep(deferred []string) map[string]bool {
+	keep := map[string]bool{"Skill": true}
+	for n := range toolRAGCore {
+		keep[n] = true
+	}
+	for _, t := range append(actool.DefaultTools(), actool.ShellSessionTools()...) {
+		keep[t.Name()] = true
+	}
+	for _, n := range deferred {
+		keep[n] = true
+	}
+	return keep
+}
+
+// applyToolRAG is the AugmentTools hook. Off unless ARTEX_TOOL_RAG is enabled.
+func applyToolRAG(ctx context.Context, agentKey string, tools []actool.CoreTool, deferred []string, score ToolScorer) []actool.CoreTool {
+	switch strings.ToLower(strings.TrimSpace(os.Getenv(toolRAGEnv))) {
+	case "1", "true", "on", "yes":
+	default:
+		return tools
+	}
+	k, kerr := strconv.Atoi(strings.TrimSpace(os.Getenv(toolRAGTopKEnv)))
+	if kerr != nil || k <= 0 {
+		k = toolRAGDefaultK
+	}
+	query, _ := ctx.Value(toolQueryKey{}).(string)
+	out, projected, err := projectTools(query, tools, toolRAGKeep(deferred), k, score)
+	if err != nil {
+		log.Printf("[tool-rag] %s: scoring failed, using full catalog (%d tools): %v", agentKey, len(tools), err)
+	} else if !projected {
+		log.Printf("[tool-rag] %s: full catalog (%d tools; empty query or under %d)", agentKey, len(tools), toolRAGMinCatalog)
+	} else {
+		names := make([]string, len(out))
+		for i, t := range out {
+			names[i] = t.Name()
+		}
+		log.Printf("[tool-rag] %s: selected %d/%d tools: %s", agentKey, len(out), len(tools), strings.Join(names, ", "))
+	}
+	return out
+}
@@ -0,0 +1,143 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"slices"
+	"testing"
+
+	actool "github.com/Autumn-27/norma/tool"
+)
+
+func ragTool(name, desc string, params ...string) actool.CoreTool {
+	props := map[string]any{}
+	for _, p := range params {
+		props[p] = map[string]any{"type": "string"}
+	}
+	return actool.Build(actool.Spec{Name: name, Description: desc, Schema: map[string]any{"type": "object", "properties": props}})
+}
+
+// ragCatalog: two core tools, three meaningful fixtures, and filler up to n tools.
+func ragCatalog(n int) []actool.CoreTool {
+	tools := []actool.CoreTool{
+		ragTool("record_fact", "Record a fact"),
+		ragTool("port_scan", "Scan TCP ports on a host, nmap-like service discovery", "host", "ports"),
+		ragTool("sqli_probe", "Probe web parameters for SQL injection", "url", "param"),
+		ragTool("http_request", "Send an HTTP request to a web endpoint", "url", "method"),
+		ragTool("list_assets", "List assets"),
+	}
+	for i := len(tools); i < n; i++ {
+		tools = append(tools, ragTool(fmt.Sprintf("filler_%02d", i), "Unrelated helper"))
+	}
+	return tools
+}
+
+func ragNames(tools []actool.CoreTool) (out []string) {
+	for _, t := range tools {
+		out = append(out, t.Name())
+	}
+	return out
+}
+
+const ragQuery = "Find SQL injection in the web login form"
+
+func TestToolRAGRanksInjectionAboveScan(t *testing.T) {
+	ranked, err := rankTools(ragQuery, ragCatalog(5), bm25Score)
+	if err != nil {
+		t.Fatal(err)
+	}
+	n := ragNames(ranked)
+	if scan := slices.Index(n, "port_scan"); slices.Index(n, "sqli_probe") > scan || slices.Index(n, "http_request") > scan {
+		t.Fatalf("ranking = %v, want sqli_probe and http_request above port_scan", n)
+	}
+	if n[0] != "sqli_probe" {
+		t.Fatalf("top = %s, want sqli_probe", n[0])
+	}
+}
+
+func TestToolRAGCoreKeptAndKRespected(t *testing.T) {
+	keep := toolRAGKeep(nil)
+	for _, k := range []int{1, 3} {
+		out, projected, err := projectTools(ragQuery, ragCatalog(25), keep, k, bm25Score)
+		if err != nil || !projected {
+			t.Fatalf("k=%d: projected=%v err=%v", k, projected, err)
+		}
+		n := ragNames(out)
+		if !slices.Contains(n, "record_fact") || !slices.Contains(n, "list_assets") {
+			t.Fatalf("k=%d: core tool hidden: %v", k, n)
+		}
+		if slices.Contains(n, "goal_met") {
+			t.Fatalf("k=%d: unbound core tool added: %v", k, n)
+		}
+		if nonCore := len(out) - 2; nonCore > k {
+			t.Fatalf("k=%d: %d non-core tools selected: %v", k, nonCore, n)
+		}
+		if !slices.Contains(n, "sqli_probe") {
+			t.Fatalf("k=%d: best match missing: %v", k, n)
+		}
+	}
+}
+
+// The projection runs after ToolResolve, so a tool the binding/enabled filter
+// dropped is never selected even when it is the best lexical match.
+func TestToolRAGNeverSelectsUnboundOrDisabled(t *testing.T) {
+	oldResolve, oldAugment, oldPolicy := ToolResolve, ToolAugment, FindingTrafficBindingEnabled
+	t.Cleanup(func() { ToolResolve, ToolAugment, FindingTrafficBindingEnabled = oldResolve, oldAugment, oldPolicy })
+	ToolAugment, FindingTrafficBindingEnabled = nil, nil
+	t.Setenv(toolRAGEnv, "on")
+	t.Setenv(toolRAGTopKEnv, "2")
+	ToolResolve = func(_ context.Context, _ string, in []actool.CoreTool) []actool.CoreTool {
+		var out []actool.CoreTool
+		for _, tool := range in {
+			if tool.Name() != "sql_injection_exploit" {
+				out = append(out, tool)
+			}
+		}
+		return out
+	}
+	base := append(ragCatalog(25), ragTool("sql_injection_exploit", "Exploit SQL injection in web login form", "url"))
+	out, _, cleanup := AugmentTools(WithToolQuery(t.Context(), ragQuery), "worker", base)
+	defer cleanup()
+	n := ragNames(out)
+	if slices.Contains(n, "sql_injection_exploit") {
+		t.Fatalf("unbound tool selected: %v", n)
+	}
+	if len(out) != 4 || !slices.Contains(n, "sqli_probe") {
+		t.Fatalf("got %v, want 2 core + top 2", n)
+	}
+}
+
+func TestToolRAGFallbacks(t *testing.T) {
+	keep := toolRAGKeep(nil)
+	failing := func(string, []actool.CoreTool) ([]float64, error) { return nil, errors.New("boom") }
+	for name, c := range map[string]struct {
+		query string
+		n     int
+		score ToolScorer
+	}{"empty query": {"  ", 25, bm25Score}, "catalog under 20": {ragQuery, 19, bm25Score}, "scoring error": {ragQuery, 25, failing}} {
+		in := ragCatalog(c.n)
+		out, projected, err := projectTools(c.query, in, keep, 1, c.score)
+		if projected || (err != nil) != (name == "scoring error") || !slices.Equal(ragNames(out), ragNames(in)) {
+			t.Errorf("%s: projected=%v err=%v out=%v", name, projected, err, ragNames(out))
+		}
+	}
+}
+
+func TestToolRAGDeterministic(t *testing.T) {
+	keep := toolRAGKeep(nil)
+	// "zzz" matches nothing: every candidate ties at 0 and the name tie-break decides.
+	for _, q := range []string{ragQuery, "zzz"} {
+		first, _, _ := projectTools(q, ragCatalog(30), keep, 4, bm25Score)
+		for i := 0; i < 20; i++ {
+			again, _, _ := projectTools(q, ragCatalog(30), keep, 4, bm25Score)
+			if !slices.Equal(ragNames(first), ragNames(again)) {
+				t.Fatalf("%q: run %d = %v, want %v", q, i, ragNames(again), ragNames(first))
+			}
+		}
+	}
+	ranked, _ := rankTools("zzz", ragCatalog(8), bm25Score)
+	if n := ragNames(ranked); !slices.IsSorted(n) {
+		t.Fatalf("ties not ordered by name: %v", n)
+	}
+}
@@ -415,7 +415,7 @@ func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.
 	// are kept as usual.
 	base = append(base, defaultToolsExcept("MultiEdit", "Glob", "Grep")...)
 	ctx = WithRunInfo(ctx, RunInfo{TaskID: taskID, ExplorationID: explorationID(ts), IntentID: intent.ID})
-	tools, def, cleanup := AugmentTools(ctx, "worker", base)
+	tools, def, cleanup := AugmentTools(WithToolQuery(ctx, nodeSummary(intent)), "worker", base)
 	defer cleanup()

 	// The intent is the worker's [sole duty, an invariant across the whole run] → it