Share lexical search ranker across tool, skill, and agent search

Author: Gustav-Simonsson

src/agent/agent-search.ts

@@ -3,12 +3,9 @@ import type { AgentTool } from "@intx/agent";
 import type { ToolDefinition } from "@intx/types/runtime";
 import { type } from "arktype";
 import { scrubSecretShapedContent } from "../plugins/tool-result-secret-scrub.js";
+import { rankLexicalMatches, tokenize } from "./lexical-search.js";
 import type { AgentProfile } from "./profiles.js";
 
-function tokenize(text: string): string[] {
-  return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
-}
-
 function profileSearchText(profile: AgentProfile): string {
   const parts = [profile.id, profile.description ?? "", profile.systemPromptRole ?? ""];
   return parts.join(" ");
@@ -18,36 +15,22 @@ export interface AgentIndex {
   search(query: string, limit?: number): AgentProfile[];
 }
 
-// Lexical ranker over id, description, and role text — same spirit as tool_search.
+// Lexical ranker over id, description, and role text — same weights as
+// tool_search / skill_search, shared via lexical-search.ts.
 export function createAgentIndex(getProfiles: () => readonly AgentProfile[]): AgentIndex {
-  const score = (profile: AgentProfile, queryTokens: string[], rawQuery: string): number => {
-    const idTokens = tokenize(profile.id);
-    const blob = profileSearchText(profile).toLowerCase();
-    const blobTokens = new Set(tokenize(blob));
-    let total = 0;
-    for (const token of queryTokens) {
-      if (idTokens.includes(token)) total += 3;
-      else if (blobTokens.has(token)) total += 1;
-      else if (profile.id.toLowerCase().includes(token)) total += 0.75;
-      else if (blob.includes(token)) total += 0.25;
-    }
-    if (profile.id.toLowerCase().includes(rawQuery)) total += 1;
-    if ((profile.description ?? "").toLowerCase().includes(rawQuery)) total += 0.5;
-    return total;
-  };
-
   return {
     search(query: string, limit = 12): AgentProfile[] {
-      const rawQuery = query.toLowerCase().trim();
-      const queryTokens = tokenize(query);
       const profiles = getProfiles();
-      if (queryTokens.length === 0) return profiles.slice(0, limit);
-      return profiles
-        .map((p) => ({ profile: p, score: score(p, queryTokens, rawQuery) }))
-        .filter((entry) => entry.score > 0)
-        .sort((a, b) => b.score - a.score)
-        .slice(0, limit)
-        .map((entry) => entry.profile);
+      if (tokenize(query).length === 0) return profiles.slice(0, limit);
+      return rankLexicalMatches(
+        profiles,
+        (profile) => ({ name: profile.id, text: profileSearchText(profile) }),
+        query,
+        limit,
+        // A raw-query hit in the description outranks a role-text-only hit.
+        (profile, rawQuery) =>
+          (profile.description ?? "").toLowerCase().includes(rawQuery) ? 0.5 : 0,
+      );
     },
   };
 }

src/agent/lexical-search.ts

@@ -0,0 +1,80 @@
+// Shared lexical ranker for the agent-side search tools (tool_search,
+// skill_search, search_agents). Each used to carry its own copy of the
+// tokenizer and scoring weights; keeping them here means a ranking change
+// (weight tuning, tokenizer tweaks) lands in one place instead of drifting
+// across three copies.
+
+/** Lowercase alphanumeric tokens; any other character splits tokens. */
+export function tokenize(text: string): string[] {
+  return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
+}
+
+/** The two scored surfaces of a searchable document. */
+export interface LexicalSearchFields {
+  /** Short identifier — exact token hits weigh most (tool name, skill name, agent id). */
+  name: string;
+  /** Body text — token hits weigh less (description, role text). */
+  text: string;
+}
+
+/**
+ * Score one document against a query. Exact name-token hits weigh most (3),
+ * then text-token hits (1), then raw-substring matches (0.75 in the name,
+ * 0.25 in the text — so "linear" finds mcp__linear__* even though it is not
+ * a whole token there), plus a raw-query hit in the name (+1). `extra` lets a
+ * caller add a document-specific term without forking the weights.
+ */
+export function scoreLexicalMatch(
+  fields: LexicalSearchFields,
+  queryTokens: readonly string[],
+  rawQuery: string,
+  extra = 0,
+): number {
+  const nameTokens = tokenize(fields.name);
+  const textTokens = new Set(tokenize(fields.text));
+  const nameLower = fields.name.toLowerCase();
+  const textLower = fields.text.toLowerCase();
+  let total = 0;
+  for (const token of queryTokens) {
+    if (nameTokens.includes(token)) total += 3;
+    else if (textTokens.has(token)) total += 1;
+    else if (nameLower.includes(token)) total += 0.75;
+    else if (textLower.includes(token)) total += 0.25;
+  }
+  if (nameLower.includes(rawQuery)) total += 1;
+  return total + extra;
+}
+
+/**
+ * Rank `items` against a query: score each document, drop zero-score hits,
+ * order by descending score (stable — ties keep input order), and cap at
+ * `limit`. Returns [] for an empty or whitespace-only query (no tokens to
+ * match), so callers that treat an empty query as "return everything" (agent
+ * search) keep that branch on their side. `extra` lets a caller add a
+ * per-item scoring term without forking the weights.
+ */
+export function rankLexicalMatches<T>(
+  items: readonly T[],
+  fieldsFor: (item: T) => LexicalSearchFields,
+  query: string,
+  limit: number,
+  extra?: (item: T, rawQuery: string) => number,
+): T[] {
+  const rawQuery = query.toLowerCase().trim();
+  const queryTokens = tokenize(query);
+  if (queryTokens.length === 0) return [];
+  return items
+    .map((item) => ({
+      item,
+      score: scoreLexicalMatch(
+        fieldsFor(item),
+        queryTokens,
+        rawQuery,
+        extra?.(item, rawQuery) ?? 0,
+      ),
+    }))
+    .filter((entry) => entry.score > 0)
+    .sort((a, b) => b.score - a.score)
+    .slice(0, limit)
+    .map((entry) => entry.item);
+}

src/agent/skill-search.ts

@@ -4,6 +4,7 @@ import type { ToolDefinition } from "@intx/types/runtime";
 import { type } from "arktype";
 
 import type { SkillSummary } from "../extensions/skills.js";
+import { rankLexicalMatches } from "./lexical-search.js";
 
 // Catalog lookup for skills. Names live in the system prompt; this tool returns
 // matching name + description so the model can choose. Bodies load via use_skill.
@@ -33,10 +34,6 @@ export interface CreateSkillSearchToolArgs {
   allowedNames?: readonly string[];
 }
 
-function tokenize(text: string): string[] {
-  return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
-}
-
 function visibleSkills(
   skills: readonly SkillSummary[],
   allowedNames: readonly string[] | undefined,
@@ -46,20 +43,6 @@ function visibleSkills(
   return skills.filter((skill) => allowed.has(skill.name));
 }
 
-function scoreSkill(skill: SkillSummary, queryTokens: string[], rawQuery: string): number {
-  const nameTokens = tokenize(skill.name);
-  const descTokens = new Set(tokenize(skill.description));
-  let total = 0;
-  for (const token of queryTokens) {
-    if (nameTokens.includes(token)) total += 3;
-    else if (descTokens.has(token)) total += 1;
-    else if (skill.name.toLowerCase().includes(token)) total += 0.75;
-    else if (skill.description.toLowerCase().includes(token)) total += 0.25;
-  }
-  if (skill.name.toLowerCase().includes(rawQuery)) total += 1;
-  return total;
-}
-
 const SkillSearchArgs = type({ query: "string" });
 
 const DEFAULT_LIMIT = 8;
@@ -75,17 +58,12 @@ export function createSkillSearchTool(args: CreateSkillSearchToolArgs): AgentToo
       }
       const query = parsed.query.trim();
       if (query.length === 0) return "Error: skill_search requires a non-empty query.";
-      const rawQuery = query.toLowerCase();
-      const queryTokens = tokenize(query);
-      if (queryTokens.length === 0) {
-        return `No skills matched "${query}". Try different keywords describing the capability.`;
-      }
-      const matches = catalog
-        .map((skill) => ({ skill, score: scoreSkill(skill, queryTokens, rawQuery) }))
-        .filter((entry) => entry.score > 0)
-        .sort((a, b) => b.score - a.score)
-        .slice(0, DEFAULT_LIMIT)
-        .map((entry) => entry.skill);
+      const matches = rankLexicalMatches(
+        catalog,
+        (skill) => ({ name: skill.name, text: skill.description }),
+        query,
+        DEFAULT_LIMIT,
+      );
       if (matches.length === 0) {
         return `No skills matched "${query}". Try different keywords describing the capability.`;
       }

src/agent/tool-search.ts

@@ -5,6 +5,7 @@ import { type } from "arktype";
 
 import type { SessionMode } from "../config/session-mode.js";
 import { sessionModeEnablesSubAgents } from "../config/session-mode.js";
+import { rankLexicalMatches } from "./lexical-search.js";
 
 // Tools whose full schema is always advertised to the model. Everything else is
 // registered and dispatchable but discovered on demand via tool_search, keeping
@@ -194,43 +195,22 @@ export interface ToolIndex {
   search(query: string, limit?: number): string[];
 }
 
-function tokenize(text: string): string[] {
-  return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
-}
-
-// A dependency-free lexical ranker over each tool's name + description. Exact name
-// token hits weigh most, then description token hits, then raw-substring matches
-// (so "linear" finds mcp__linear__* even though it is not a whole token there).
+// Lexical ranker over each tool's name + description — weights shared with
+// skill_search and search_agents (see lexical-search.ts). Exact name token hits
+// weigh most, then description token hits, then raw-substring matches (so
+// "linear" finds mcp__linear__* even though it is not a whole token there).
 export function createToolIndex(
   getDefs: () => readonly ToolDefinition[],
   advertisedNames: readonly string[] = ADVERTISED_TOOL_NAMES,
 ): ToolIndex {
-  const score = (def: ToolDefinition, queryTokens: string[], rawQuery: string): number => {
-    const nameTokens = tokenize(def.name);
-    const descTokens = new Set(tokenize(def.description ?? ""));
-    let total = 0;
-    for (const token of queryTokens) {
-      if (nameTokens.includes(token)) total += 3;
-      else if (descTokens.has(token)) total += 1;
-      else if (def.name.toLowerCase().includes(token)) total += 0.75;
-      else if ((def.description ?? "").toLowerCase().includes(token)) total += 0.25;
-    }
-    if (def.name.toLowerCase().includes(rawQuery)) total += 1;
-    return total;
-  };
-
   return {
     search(query: string, limit = 8): string[] {
-      const rawQuery = query.toLowerCase().trim();
-      const queryTokens = tokenize(query);
-      if (queryTokens.length === 0) return [];
-      return getDefs()
-        .filter((def) => !advertisedNames.includes(def.name))
-        .map((def) => ({ name: def.name, score: score(def, queryTokens, rawQuery) }))
-        .filter((entry) => entry.score > 0)
-        .sort((a, b) => b.score - a.score)
-        .slice(0, limit)
-        .map((entry) => entry.name);
+      return rankLexicalMatches(
+        getDefs().filter((def) => !advertisedNames.includes(def.name)),
+        (def) => ({ name: def.name, text: def.description ?? "" }),
+        query,
+        limit,
+      ).map((def) => def.name);
     },
   };
 }