aaronrene / knowtation public
sampling-rerank.mjs
85 lines 2.9 KB
Raw
sha256:65ccb454656ea5acdea0a10e559b78bcde1eb6ff753ecc2911bc99d1c3d7cadd feat(calendar): enforce agent context tiers in retrieval AP… Human minor ⚠ breaking 1 day ago
1 /**
2 * Issue #1 Phase F4 — sampling-based search result reranking.
3 * After vector search returns candidates, use the client LLM to rerank by semantic relevance.
4 */
5
6 import { trySampling } from '../sampling.mjs';
7
8 const RERANK_SYSTEM = `You are a search result reranker. Given a query and numbered search results, return ONLY a JSON array of the result numbers (1-based) sorted by relevance to the query, most relevant first. Example: [3, 1, 5, 2, 4]
9 Do not include results that are clearly irrelevant. Return at most the number requested.`;
10
11 const MAX_RESULTS_FOR_RERANK = 20;
12
13 /**
14 * Attempt to rerank search results using the client LLM via sampling.
15 * Falls back to the original order when sampling is unavailable or fails.
16 *
17 * @param {import('@modelcontextprotocol/sdk/server/mcp.js').McpServer} mcpServer
18 * @param {string} query
19 * @param {{ path: string, snippet?: string, score?: number }[]} results
20 * @param {number} [limit]
21 * @returns {Promise<{ path: string, snippet?: string, score?: number }[]>}
22 */
23 export async function rerankWithSampling(mcpServer, query, results, limit) {
24 if (!results || results.length <= 1) return results;
25
26 const candidates = results.slice(0, MAX_RESULTS_FOR_RERANK);
27 const numbered = candidates
28 .map((r, i) => `${i + 1}. [${r.path}] ${(r.snippet || '').slice(0, 200)}`)
29 .join('\n');
30
31 const user = `Query: "${query}"\nReturn up to ${limit || candidates.length} results.\n\nResults:\n${numbered}`;
32
33 const raw = await trySampling(mcpServer, {
34 system: RERANK_SYSTEM,
35 user,
36 maxTokens: 256,
37 });
38
39 if (!raw) return results;
40
41 const ranked = parseRerankResponse(raw, candidates.length);
42 if (!ranked || ranked.length === 0) return results;
43
44 const reordered = [];
45 const used = new Set();
46 for (const idx of ranked) {
47 if (idx >= 0 && idx < candidates.length && !used.has(idx)) {
48 reordered.push(candidates[idx]);
49 used.add(idx);
50 }
51 }
52 for (let i = 0; i < candidates.length; i++) {
53 if (!used.has(i)) reordered.push(candidates[i]);
54 }
55 if (results.length > MAX_RESULTS_FOR_RERANK) {
56 reordered.push(...results.slice(MAX_RESULTS_FOR_RERANK));
57 }
58
59 return limit ? reordered.slice(0, limit) : reordered;
60 }
61
62 /**
63 * Parse the LLM rerank response into an array of 0-based indices.
64 * @param {string} raw
65 * @param {number} maxIdx
66 * @returns {number[] | null}
67 */
68 export function parseRerankResponse(raw, maxIdx) {
69 if (!raw) return null;
70 try {
71 const cleaned = raw.replace(/^```(?:json)?\s*/i, '').replace(/```\s*$/, '').trim();
72 const arr = JSON.parse(cleaned);
73 if (!Array.isArray(arr)) return null;
74 return arr
75 .filter((n) => typeof n === 'number' && Number.isFinite(n))
76 .map((n) => Math.round(n) - 1)
77 .filter((i) => i >= 0 && i < maxIdx);
78 } catch (_) {
79 const matches = raw.match(/\d+/g);
80 if (!matches) return null;
81 return matches
82 .map((s) => parseInt(s, 10) - 1)
83 .filter((i) => Number.isFinite(i) && i >= 0 && i < maxIdx);
84 }
85 }
File History 2 commits
sha256:65ccb454656ea5acdea0a10e559b78bcde1eb6ff753ecc2911bc99d1c3d7cadd feat(calendar): enforce agent context tiers in retrieval AP… Human minor ⚠ 1 day ago
sha256:9103f98c89257ed2b01c237cea895dabb3e85ea337dccb1161c175e4422355b6 docs: accept Calendar Events v0 spec with Phase 0 security … Human 1 day ago