Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/agent/agent-search.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { type } from "arktype";
import { scrubSecretShapedContent } from "../plugins/tool-result-secret-scrub.js";
import {
lexicalFields,
NO_LEXICAL_PARAMS,
rankAndCut,
scoreLexical,
tokenizeLexical,
Expand Down Expand Up @@ -35,7 +36,7 @@ export function createAgentIndex(
rawQuery: string,
): number =>
scoreLexical(
lexicalFields(profile.id, profileSearchText(profile)),
lexicalFields(profile.id, profileSearchText(profile), NO_LEXICAL_PARAMS),
queryTokens,
rawQuery,
) +
Expand Down
101 changes: 101 additions & 0 deletions src/agent/lexical-rank.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
import { describe, expect, test } from "bun:test";

import {
lexicalFields,
NO_LEXICAL_PARAMS,
scoreLexical,
tokenizeLexical,
} from "./lexical-rank.js";

describe("scoreLexical", () => {
test("counts a word repeated in the query once", () => {
const fields = lexicalFields(
"mcp__linear__list_release_notes",
"List release notes in the workspace",
NO_LEXICAL_PARAMS,
);
// Same raw query on both sides so only the token list differs; the raw
// query feeds a separate whole-query substring bonus.
const rawQuery = "release notes release";
expect(
scoreLexical(fields, tokenizeLexical("release notes release"), rawQuery),
).toBe(scoreLexical(fields, tokenizeLexical("release notes"), rawQuery));
});

// Two-word raw query that no fixture contains whole, so the whole-query
// substring bonus never applies and only the per-token tiers are compared.
const rawQuery = "release filter";
const queryTokens = tokenizeLexical(rawQuery);

test("ranks a name substring above a parameter name above a parameter description", () => {
// Every fixture matches "filter" in its blob, so the parameter tiers have
// a direct match to boost.
const nameSubstring = lexicalFields(
"mcp__acme__releaselog",
"Filter logs",
NO_LEXICAL_PARAMS,
);
const paramName = lexicalFields("mcp__acme__records", "Filter records", {
names: ["release"],
blob: "",
});
const paramDescription = lexicalFields(
"mcp__acme__records",
"Filter records",
{
names: ["id"],
blob: "Release ID or slug",
},
);
const scores = [nameSubstring, paramName, paramDescription].map((fields) =>
scoreLexical(fields, queryTokens, rawQuery),
);
expect(scores[0]).toBeGreaterThan(scores[1] ?? 0);
expect(scores[1]).toBeGreaterThan(scores[2] ?? 0);
expect(scores[2]).toBeGreaterThan(0);
});

test("ranks a parameter name above a description substring of the same word", () => {
// "release" sits only inside "prereleases", so the description offers
// just a raw-substring hit.
const without = lexicalFields(
"mcp__acme__records",
"Read prereleases",
NO_LEXICAL_PARAMS,
);
const withParam = lexicalFields("mcp__acme__records", "Read prereleases", {
names: ["release"],
blob: "",
});
expect(scoreLexical(withParam, queryTokens, rawQuery)).toBeGreaterThan(
scoreLexical(without, queryTokens, rawQuery),
);
});

test("a parameter match alone scores nothing", () => {
const fields = lexicalFields("mcp__acme__records", "Read records", {
names: ["release"],
blob: "Release ID or slug",
});
expect(scoreLexical(fields, queryTokens, rawQuery)).toBe(0);
});

test("a parameter match adds nothing to a word already matched by name or description", () => {
const params = { names: ["release"], blob: "Release ID or slug" };
for (const [name, blob] of [
["mcp__acme__release", "Read logs"],
["mcp__acme__records", "Read a release"],
["mcp__acme__releaselog", "Read logs"],
] as const) {
expect(
scoreLexical(lexicalFields(name, blob, params), queryTokens, rawQuery),
).toBe(
scoreLexical(
lexicalFields(name, blob, NO_LEXICAL_PARAMS),
queryTokens,
rawQuery,
),
);
}
});
});
73 changes: 62 additions & 11 deletions src/agent/lexical-rank.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
// One shared dependency-free lexical ranker behind tool_search,
// search_agents, and skill_search. Exact name-token hits weigh most, then
// blob-token hits, then raw-substring matches (so "linear" finds
// mcp__linear__* even though it is not a whole token there).
// blob-token hits, then a raw name substring (so "linear" finds
// mcp__linear__* even though it is not a whole token there), then a raw blob
// substring. Parameter name and description tokens (tool_search only) weigh
// less than a name substring and only boost an item that already matches by
// name or blob: a filter makes the tool that carries it rank higher, but a
// query that matches nothing except parameter text stays a miss, so
// tool_search still waits for servers that are connecting.
export function tokenizeLexical(text: string): string[] {
return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
}
Expand All @@ -11,33 +16,79 @@ export interface LexicalFields {
readonly blob: string;
readonly nameTokens: readonly string[];
readonly blobTokens: ReadonlySet<string>;
readonly paramNameTokens: ReadonlySet<string>;
readonly paramBlobTokens: ReadonlySet<string>;
}

export function lexicalFields(name: string, blob: string): LexicalFields {
interface LexicalParams {
readonly names: readonly string[];
readonly blob: string;
}

export const NO_LEXICAL_PARAMS: LexicalParams = { names: [], blob: "" };

// JSON Schema keys are usually camelCase, and an undescribed flag such as
// `includeArchived` carries its word only inside the key.
function splitCamelCase(key: string): string {
return key
.replace(/([a-z0-9])([A-Z])/g, "$1 $2")
.replace(/([A-Z])([A-Z][a-z])/g, "$1 $2");
}

export function lexicalFields(
name: string,
blob: string,
params: LexicalParams,
): LexicalFields {
return {
name,
blob,
nameTokens: tokenizeLexical(name),
blobTokens: new Set(tokenizeLexical(blob)),
paramNameTokens: new Set(
tokenizeLexical(params.names.map(splitCamelCase).join(" ")),
),
paramBlobTokens: new Set(tokenizeLexical(params.blob)),
};
}

function directTokenScore(
fields: LexicalFields,
name: string,
blob: string,
token: string,
): number {
if (fields.nameTokens.includes(token)) return 3;
if (fields.blobTokens.has(token)) return 1;
if (name.includes(token)) return 0.75;
if (blob.includes(token)) return 0.25;
return 0;
}

function paramTokenScore(fields: LexicalFields, token: string): number {
if (fields.paramNameTokens.has(token)) return 0.5;
if (fields.paramBlobTokens.has(token)) return 0.25;
return 0;
}

export function scoreLexical(
fields: LexicalFields,
queryTokens: readonly string[],
rawQuery: string,
): number {
const name = fields.name.toLowerCase();
const blob = fields.blob.toLowerCase();
let total = 0;
for (const token of queryTokens) {
if (fields.nameTokens.includes(token)) total += 3;
else if (fields.blobTokens.has(token)) total += 1;
else if (name.includes(token)) total += 0.75;
else if (blob.includes(token)) total += 0.25;
let matched = name.includes(rawQuery) ? 1 : 0;
let total = matched;
// A model repeats a word because its sentence needs it ("release ID or
// release slug"), not to weight it; counting each copy multiplies the
// bonus for every tool that matches that one word.
for (const token of new Set(queryTokens)) {
const direct = directTokenScore(fields, name, blob, token);
matched += direct;
total += Math.max(direct, paramTokenScore(fields, token));
}
if (name.includes(rawQuery)) total += 1;
return total;
return matched > 0 ? total : 0;
}

// One shared rank→filter→sort→slice→map cut behind tool_search,
Expand Down
3 changes: 2 additions & 1 deletion src/agent/skill-search.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { type } from "arktype";
import type { SkillSummary } from "../extensions/skills.js";
import {
lexicalFields,
NO_LEXICAL_PARAMS,
rankAndCut,
scoreLexical,
tokenizeLexical,
Expand Down Expand Up @@ -67,7 +68,7 @@ function scoreSkill(
rawQuery: string,
): number {
return scoreLexical(
lexicalFields(skill.name, skill.description),
lexicalFields(skill.name, skill.description, NO_LEXICAL_PARAMS),
queryTokens,
rawQuery,
);
Expand Down
Loading
Loading