diff --git a/agent/assembly.go b/agent/assembly.go index 0b803a9..f4f8e9a 100644 --- a/agent/assembly.go +++ b/agent/assembly.go @@ -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) diff --git a/agent/planner.go b/agent/planner.go index e422fad..8139ffa 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -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 diff --git a/agent/toolrag.go b/agent/toolrag.go new file mode 100644 index 0000000..f4e6c83 --- /dev/null +++ b/agent/toolrag.go @@ -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 +} diff --git a/agent/toolrag_test.go b/agent/toolrag_test.go new file mode 100644 index 0000000..4148e8a --- /dev/null +++ b/agent/toolrag_test.go @@ -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) + } +} diff --git a/agent/worker.go b/agent/worker.go index 66f13de..582911b 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -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