Repository navigation
feat(prompt): bound the skill listing and add an opt-in smaller tool list #1428
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
anandgupta42
wants to merge
10
commits into
main
Choose a base branch
from
feat/smaller-tool-context
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
10 commits
Select commit
Hold shift + click to select a range
b740981
feat(skills): keep the skill listing inside a token budget
4d618d9
feat(tools): offer a smaller default tool list, with a fixed path to …
f2bc930
fix(prompt): address review feedback on the skill listing and the too…
9be11e0
fix(skills): search the current skill list when a lookup misses
b49bcdc
fix(skills): mark the agent captured for the skill lookup
8d0dec3
fix(prompt): second review round on the tool router and the skill search
anandgupta42 4365a0c
fix(prompt): third review round, smaller edge cases
e8d1db0
fix(prompt): last review round on router rules, path casing and tests
3e89847
fix(prompt): enforce permission denies when a governed tool runs
7c308a2
fix(prompt): refuse a denied tool inside its own wrapper, not before it
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,214 @@ | ||
| /** | ||
| * Bounded skill listing. | ||
| * | ||
| * The `skill` tool description and the system prompt both enumerate the installed skills. Without a | ||
| * bound, the system-prompt copy grows with every skill (about 190 tokens each: name, description and a | ||
| * file URL) and the tool-description copy stops at 50 entries, leaving skill 51 onward unreachable by | ||
| * the model unless it already knows the name. Both sit in the prompt prefix that is re-sent on every | ||
| * request. | ||
| * | ||
| * This renders a listing that stays inside a token budget and never hides a skill: | ||
| * 1. Skills are ordered deterministically (embedded skills first, then by name), so the text is the | ||
| * same for the same set of skills and the cached prefix is stable. | ||
| * 2. Every skill gets at least its name. Descriptions (single line, truncated) are added in order | ||
| * while the budget lasts. | ||
| * 3. If even the names do not fit, the tail is replaced by a count and a pointer to the search below. | ||
| * 4. Passing a keyword as the skill `name` searches all installed skills (see `findSkills`), so a | ||
| * skill that is not shown can still be found, then loaded by its exact name. | ||
| * | ||
| * Budgets are in estimated tokens (4 characters each); the estimate is only used to bound size. | ||
| */ | ||
| import { Skill } from "../skill" | ||
|
|
||
| export const TOOL_LISTING_BUDGET_TOKENS = 1_000 | ||
| export const SYSTEM_LISTING_BUDGET_TOKENS = 2_500 | ||
| /** | ||
| * Description lengths to try, longest first. The listing uses the longest one at which every skill | ||
| * still fits its budget, so a short list keeps its trigger conditions and a long one degrades evenly. | ||
| */ | ||
| const TOOL_DESCRIPTION_CAPS = [160, 70] | ||
| const SYSTEM_DESCRIPTION_CAPS = [400, 160] | ||
| const SYSTEM_DESCRIPTION_CHARS = 160 | ||
| const CHARS_PER_TOKEN = 4 | ||
| const MATCH_LIMIT = 15 | ||
| const NAME_LIMIT = 40 | ||
| /** A name longer than this is shortened for display; it is still found by search and loaded by its full name. */ | ||
| const DISPLAY_NAME_CHARS = 100 | ||
| /** The longest name a search result shows whole. */ | ||
| const RESULT_NAME_CHARS = 256 | ||
| /** Room kept for the wrapper tags and the "not listed" footer, which the entries must not crowd out. */ | ||
| const OVERHEAD_CHARS = 400 | ||
|
|
||
| export type ListingKind = "tool" | "system" | ||
|
|
||
| type Entry = Pick<Skill.Info, "name" | "description" | "location"> | ||
|
|
||
| const compare = (a: string, b: string) => (a < b ? -1 : a > b ? 1 : 0) | ||
|
|
||
| /** Embedded (shipped) skills first, then everything else; each group by name. */ | ||
| export function orderSkills<T extends Entry>(skills: readonly T[]): T[] { | ||
| const rank = (skill: Entry) => (Skill.hasNoSkillDirectory(skill.location) ? 0 : 1) | ||
| return [...skills].sort((a, b) => rank(a) - rank(b) || compare(a.name, b.name) || compare(a.location, b.location)) | ||
| } | ||
|
|
||
| /** Collapse to one line and cut to `max` characters by code point, so a surrogate pair is never split. */ | ||
| function oneLine(text: string, max: number): string { | ||
| const flat = text.replace(/\s+/g, " ").trim() | ||
| const chars = Array.from(flat) | ||
| return chars.length <= max ? flat : `${chars.slice(0, max - 1).join("").trimEnd()}…` | ||
| } | ||
|
|
||
| /** | ||
| * How a name is shown. A plain name is shown as it is, so it can be copied back. A name with line breaks, | ||
| * repeated spaces or tabs is shown JSON-quoted, which keeps it exact on one line. A name over the display | ||
| * length is shortened with an ellipsis; search shows the longer form. | ||
| */ | ||
| function display(name: string, max: number): string { | ||
| const neutral = Skill.neutralizeListingWrapper(name) | ||
| const plain = neutral === neutral.replace(/\s+/g, " ").trim() | ||
| // JSON leaves U+2028 and U+2029 as they are; they would still break a line, so escape them too. | ||
| const shown = plain ? neutral : JSON.stringify(neutral).replace(/\u2028/g, "\\u2028").replace(/\u2029/g, "\\u2029") | ||
| const chars = Array.from(shown) | ||
| return chars.length <= max ? shown : `${chars.slice(0, max - 1).join("")}…` | ||
| } | ||
|
|
||
| function label(skill: Entry) { | ||
| return display(skill.name, DISPLAY_NAME_CHARS) | ||
| } | ||
|
|
||
| function line(skill: Entry, max: number, name = label(skill)): string { | ||
| const description = oneLine(Skill.neutralizeListingWrapper(skill.description ?? ""), max) | ||
| return description ? `${name}: ${description}` : name | ||
| } | ||
|
|
||
| export function renderBoundedListing(skills: readonly Entry[], kind: ListingKind): string { | ||
| const ordered = orderSkills(skills) | ||
| const budget = | ||
| (kind === "tool" ? TOOL_LISTING_BUDGET_TOKENS : SYSTEM_LISTING_BUDGET_TOKENS) * CHARS_PER_TOKEN - OVERHEAD_CHARS | ||
| const caps = kind === "tool" ? TOOL_DESCRIPTION_CAPS : SYSTEM_DESCRIPTION_CAPS | ||
| const names = ordered.map(label) | ||
| const render = (max: number) => ordered.map((skill) => line(skill, max)) | ||
| // Same per-entry cost as the spending loop below: the line, a separator and a newline. | ||
| const total = (lines: string[]) => lines.reduce((sum, l) => sum + l.length + 2, 0) | ||
| let full = render(caps[caps.length - 1]!) | ||
| for (const max of caps) { | ||
| const lines = render(max) | ||
| if (total(lines) <= budget) { | ||
| full = lines | ||
| break | ||
| } | ||
| } | ||
|
|
||
| // Names are the floor: show as many as fit, in order. | ||
| let used = 0 | ||
| let listed = 0 | ||
| for (const name of names) { | ||
| const next = used + name.length + 2 | ||
| if (next > budget) break | ||
| used = next | ||
| listed++ | ||
| } | ||
|
|
||
| // Spend what is left on descriptions, front to back, and stop at the first one that does not fit. | ||
| let described = 0 | ||
| while (described < listed) { | ||
| const extra = full[described]!.length - names[described]!.length | ||
| if (used + extra > budget) break | ||
| used += extra | ||
| described++ | ||
| } | ||
|
|
||
| const hidden = ordered.length - listed | ||
| const out = ["<available_skills>"] | ||
| for (let i = 0; i < described; i++) out.push(full[i]!) | ||
| if (listed > described) out.push(names.slice(described, listed).join(", ")) | ||
| if (hidden > 0) { | ||
| out.push( | ||
| `... and ${hidden} more installed skills (${ordered.length} in total) are not listed here. ` + | ||
| `Call this tool with a keyword as the name (for example "snowflake") to search all of them.`, | ||
| ) | ||
| } else if (ordered.length > 0 && listed > described) { | ||
| out.push(`Call this tool with a keyword as the name to search skills by description.`) | ||
| } | ||
| out.push("</available_skills>") | ||
| return out.join("\n") | ||
| } | ||
|
|
||
| function words(text: string): string[] { | ||
| return text.toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? [] | ||
| } | ||
|
|
||
| /** | ||
| * Rank every installed skill against a query. Deterministic: score, then the listing order. | ||
| * Used when the model passes something that is not an exact skill name. | ||
| */ | ||
| export function findSkills<T extends Entry>(skills: readonly T[], query: string, limit = MATCH_LIMIT): T[] { | ||
| const terms = [...new Set(words(query))] | ||
| if (terms.length === 0) return [] | ||
| const exact = query.trim().toLowerCase() | ||
| return orderSkills(skills) | ||
| .map((skill, index) => { | ||
| const name = skill.name.toLowerCase() | ||
| const description = (skill.description ?? "").toLowerCase() | ||
| // A skill is always found by its own name: an exact match outranks any keyword score. | ||
| let score = name === exact ? 1_000_000 : 0 | ||
| let covered = 0 | ||
| for (const term of terms) { | ||
| const inName = name.includes(term) | ||
| const inDescription = description.includes(term) | ||
| if (inName) score += 3 | ||
| if (inDescription) score += 1 | ||
| if (inName || inDescription) covered++ | ||
| } | ||
| // Matching more of the query beats matching one term strongly. | ||
| score += covered * 10_000 | ||
| return { skill, score, index } | ||
| }) | ||
| .filter((hit) => hit.score > 0) | ||
| .sort((a, b) => b.score - a.score || a.index - b.index) | ||
| .slice(0, limit) | ||
| .map((hit) => hit.skill) | ||
| } | ||
|
|
||
| /** The text of a failed lookup: matches with descriptions when the keyword hits, else a bounded name list. */ | ||
| export function notFoundMessage(skills: readonly Entry[], requested: string): string { | ||
| const matches = findSkills(skills, requested) | ||
| const head = `Skill "${oneLine(Skill.neutralizeListingWrapper(requested), 80)}" not found. ${skills.length} skills are installed.` | ||
| if (matches.length > 0) { | ||
| const total = findSkills(skills, requested, Number.MAX_SAFE_INTEGER).length | ||
| return [ | ||
| head, | ||
| `Closest matches (call this tool again with the exact name to load one):`, | ||
| // Search results carry the full name: it is what the skill is loaded by. | ||
| ...matches.map((skill) => `- ${line(skill, SYSTEM_DESCRIPTION_CHARS, display(skill.name, RESULT_NAME_CHARS))}`), | ||
| ...(total > matches.length ? [`... ${total - matches.length} more match; use a more specific keyword.`] : []), | ||
| ].join("\n") | ||
| } | ||
| const ordered = orderSkills(skills) | ||
| const shown = ordered.slice(0, NAME_LIMIT).map(label) | ||
| const rest = ordered.length - shown.length | ||
| return [ | ||
| head, | ||
| `No skill matches that keyword. ${rest > 0 ? "First" : "Available"} skills: ${shown.join(", ") || "none"}${rest > 0 ? `, and ${rest} more` : ""}.`, | ||
| ].join("\n") | ||
| } | ||
|
|
||
| /** | ||
| * Configuration switch. The environment variable wins over the config file; `fallback` applies when | ||
| * neither is set. | ||
| */ | ||
| export function switchEnabled(envKey: string, configured: boolean | undefined, fallback: boolean): boolean { | ||
| const raw = process.env[envKey]?.toLowerCase() | ||
| if (raw === "1" || raw === "true") return true | ||
| if (raw === "0" || raw === "false") return false | ||
| return configured ?? fallback | ||
| } | ||
|
|
||
| /** Default for the bounded listing. */ | ||
| export const BOUNDED_SKILL_LISTING_DEFAULT = true | ||
|
|
||
| export function boundedSkillListingEnabled(configured: boolean | undefined): boolean { | ||
| return switchEnabled("ALTIMATE_BOUNDED_SKILL_LISTING", configured, BOUNDED_SKILL_LISTING_DEFAULT) | ||
| } | ||
|
|
||
| export * as SkillListing from "./skill-listing" | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| /** | ||
| * `tool_run`: the one fixed tool through which a tool outside the offered list is called. | ||
| * | ||
| * Its definition (name, schema, and the index of reachable tools) is fixed for the whole session, so | ||
| * reaching a hidden tool never changes the tool block of the prompt (see tool-selection.ts). | ||
| */ | ||
| import { tool, jsonSchema, asSchema, type Tool as AITool } from "ai" | ||
| import { Log } from "@/util/log" | ||
| import { ToolLookup } from "./tools/tool-lookup" | ||
| import { ToolSelection } from "./tool-selection" | ||
|
|
||
| const log = Log.create({ service: "tool.run" }) | ||
|
|
||
| /** Calls the model addressed to a hidden tool directly, which were rerouted here instead of failing. */ | ||
| const rerouted = new Set<string>() | ||
| export function markRerouted(toolCallId: string) { | ||
| // A rerouted call that never executes (an abort) would leave its id behind; keep the set small. | ||
| if (rerouted.size > 1000) rerouted.clear() | ||
| rerouted.add(toolCallId) | ||
| } | ||
|
|
||
| function contract(target: AITool): string { | ||
| try { | ||
| const schema = (asSchema(target.inputSchema) as { jsonSchema: any }).jsonSchema | ||
| const params = ToolLookup.describeJsonSchema(schema) | ||
| if (params.length === 0) return "No parameters." | ||
| return params | ||
| .map((p) => ` ${p.name} (${p.type}, ${p.required ? "required" : "optional"})${p.description ? ` ${p.description}` : ""}`) | ||
| .join("\n") | ||
| } catch { | ||
| return "Use tool_lookup to see the parameters." | ||
| } | ||
| } | ||
|
|
||
| export function createRunTool(input: { | ||
| hidden: Record<string, AITool> | ||
| }): AITool { | ||
| const names = Object.keys(input.hidden) | ||
| const run = tool({ | ||
| description: ToolSelection.runDescription(names), | ||
| inputSchema: jsonSchema<{ name?: unknown; arguments?: unknown }>({ | ||
| type: "object", | ||
| properties: { | ||
| name: { type: "string", description: "Exact name of the tool to run" }, | ||
| arguments: { type: "object", description: "The tool's own parameters", additionalProperties: true }, | ||
| }, | ||
|
anandgupta42 marked this conversation as resolved.
|
||
| required: ["name"], | ||
| }), | ||
| async execute(args: { name?: unknown; arguments?: unknown }, options) { | ||
| const name = typeof args.name === "string" ? args.name.trim() : "" | ||
| const params: unknown = args.arguments === undefined ? {} : args.arguments | ||
| if (typeof params !== "object" || params === null || Array.isArray(params)) { | ||
| throw new Error(`tool_run: "arguments" must be an object holding the tool's parameters.`) | ||
| } | ||
| const target = Object.hasOwn(input.hidden, name) ? input.hidden[name] : undefined | ||
| if (!target?.execute) { | ||
| throw new Error( | ||
| name === "" | ||
| ? `tool_run needs a tool name. Available through tool_run: ${names.join(", ")}` | ||
| : `No tool named "${name}" is available through tool_run. Available: ${names.join(", ")}. ` + | ||
| `Tools in your own tool list are called directly.`, | ||
| ) | ||
| } | ||
| // The permission deny for a governed tool is enforced inside the target's own wrapper (session/prompt.ts), | ||
| // after the call's execution has been registered, so a refusal pairs with its call even when a provider | ||
| // repeats call ids. | ||
| const wasRerouted = options.toolCallId ? rerouted.delete(options.toolCallId) : false | ||
| log.info("run", { tool: name, rerouted: wasRerouted }) | ||
| try { | ||
| const result = (await target.execute(params as never, options)) as { metadata?: Record<string, unknown> } | ||
| if (result && typeof result === "object") { | ||
| return { ...result, metadata: { ...result.metadata, tool_run: name, ...(wasRerouted ? { rerouted: true } : {}) } } | ||
| } | ||
| return result | ||
| } catch (e) { | ||
| const message = e instanceof Error ? e.message : String(e) | ||
| if (/invalid arguments/i.test(message)) throw new Error(`${message}\n\nParameters of ${name}:\n${contract(target)}`) | ||
| throw e | ||
| } | ||
| }, | ||
| }) | ||
| ToolSelection.attachHidden(run, input.hidden) | ||
| return run | ||
| } | ||
|
|
||
| export * as ToolRun from "./tool-run" | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.