mirror of
https://github.com/okieselbach/AutopilotMonitor.git
synced 2026-10-24 04:32:32 +00:00
fix(mcp): put every doc token back into search_docs ranking
Two measured recall gaps in search_docs, both silent. 1. Truncation. MiniLM's context is 256 tokens; 87 of 324 chunks exceeded it and 13% of all tokens never reached the ranking — the tokenizer drops the tail without complaint. Concretely, in built-in-rules.md :: Apps (734 tokens) the strings ANALYZE-APP-013, ANALYZE-OFFICE-001, ANALYZE-CORR-003 and Click-to-Run all sit past the cutoff, so "proxy configuration blocking app downloads" did not retrieve that section at all. Shrinking chunks cannot fix this: the corpus ranges 2.43 to 5.81 characters per token, so no character budget bounds tokens, and a budget safe for the worst case (~620 chars) would shred prose and tables alike. So the chunk stays whole and the embedding is split instead — index() cuts long text into 250-token windows with 40 overlap, a document carries embeddings[][], and search() scores it by its BEST window. Max rather than mean: averaging would blur a long section's several topics into one mediocre direction, the same dilution that makes large context windows worse than chunking. The tokenizer is a build-time dependency only (vocabulary, not ONNX weights, ~93 ms); production boots from precomputed vectors and never loads it. 2. No general fallback. Proper nouns embed poorly — "RealmJoin" occurs in 3 chunks yet scores under the threshold. Fuzzy matching was built, measured and rejected: at a threshold low enough to surface RealmJoin (0.28), Fuse also returns 5 hits for "SCEPman", a word this corpus contains zero times. Literal containment has neither problem, and every provider already exposes it. So a short result set is now topped up via lexicalMatch on distinctive query terms, requiring two matching terms as corroboration — without that rule "how do I bake sourdough bread" matched a page about LLM providers, because "bake" is a substring of "baked into the image". Keyword hits are appended and flagged, never sorted together with cosine scores. Also sharpens the tool description with the measured reason to prefer this over fetching the docs site (~1k tokens for topK=3 against ~9k raw / ~600k rendered, version-matched, and matching on meaning), and replaces the growing tail of positional docs parameters with a DocsSearchBundle. Verified: token coverage 87% -> 100%; "Office Click-to-Run install never finished" 0.318 -> 0.442; "proxy configuration blocking app downloads" absent -> 0.559 at rank 1; the ten reference queries keep byte-identical scores and rank 1 (short chunks are one window, so the same vector as before — no long-chunk bias). The fallback is not a dead path: across 904 rare corpus tokens queried alone it contributes hits for 431, while negative controls stay empty. Precompute output is accepted by the boot validator and a mutated corpus still rejected; index 2.63 -> 3.12 MB for 324 chunks / 424 vectors. The rules corpus gains the same treatment unplanned: 128 rules now carry 172 vectors. 445 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -2,6 +2,8 @@
|
||||
|
||||
## 2026-07-21
|
||||
|
||||
* **Update**: Closed the two measured recall gaps in `search_docs` (`mcp/docs-corpus.md`). (1) **Token windows** — 87 of 324 chunks exceeded the model's 256-token context and 13% of all tokens never reached the ranking; a character budget cannot fix this (the corpus ranges 2.43–5.81 chars/token), so long text is now embedded as 250-token windows with 40 overlap and scored by its BEST window (`embeddings: number[][]`, max not mean). Token coverage 87% → 100%; "proxy configuration blocking app downloads" went from absent to 0.559 at rank 1; the ten reference queries kept byte-identical scores. Costs +31% vectors, index 2.63 → 3.12 MB. The rules corpus gained the same treatment (128 rules → 172 vectors). (2) **Keyword fallback** — when the semantic pass returns fewer than `topK`, results are topped up via literal `lexicalMatch` on distinctive query terms, requiring two matching terms for corroboration (one incidental substring match returned a page about LLM providers for a sourdough query). Fuzzy matching was measured and rejected: at a threshold low enough to surface "RealmJoin" it invents 5 hits for "SCEPman", which the corpus never mentions. Also sharpened the tool description with the measured context argument for preferring it over fetching the docs site.
|
||||
|
||||
* **Creation**: Added `mcp/docs-corpus.md` (new `mcp/` section) — the MCP server gains `search_docs`, a third semantic corpus built from the published customer documentation repository (`autopilotmonitor-docs`, checked out into the Docker build context by the deploy workflow). Documents the h2/h3 + `<details>` chunking with breadcrumb prefixes, the two deliberate size exceptions (tables never split, unbroken changelog lists split at line level), the CRLF precondition (`.` does not match `\r` in JS — headings silently stop matching, which produced 286 wrong-but-plausible chunks before it was found), the optional `docs` section in `search-index.json` and its validation asymmetry, embedding rounding to 6 decimals, and the cross-repo coupling: a docs edit reaches clients only after an MCP redeploy, made checkable via `DOCS_COMMIT` on `/health`. Container resources raised to 0.5 vCPU / 1.0 GiB for the cold-start query-embedder load.
|
||||
|
||||
## 2026-07-20
|
||||
|
||||
+62
-1
@@ -87,6 +87,67 @@ question sections.
|
||||
`normalizeNewlines()` runs before any parsing. Do not remove it, and do not
|
||||
"fix" a future parsing bug here by adding `\r?` to one pattern.
|
||||
|
||||
## Token windows: why a chunk can hold several vectors
|
||||
|
||||
Chunking at heading boundaries keeps sections coherent, but it does not bound the
|
||||
result by *tokens* — and it cannot. Measured on this corpus the ratio ranges from
|
||||
**2.43 to 5.81 characters per token**, a factor of 2.4, so no character budget maps
|
||||
reliably onto the model's 256-token window. A budget safe for the worst case
|
||||
(~620 characters) would shred prose and tables alike.
|
||||
|
||||
Left alone, 87 of 324 chunks exceeded the window and **13% of all tokens never
|
||||
reached the ranking** — silently, since the tokenizer truncates without complaint.
|
||||
The effect was concrete: in `built-in-rules.md :: Apps` (734 tokens), the strings
|
||||
`ANALYZE-APP-013`, `ANALYZE-OFFICE-001`, `ANALYZE-CORR-003` and `Click-to-Run` all
|
||||
sit past the cutoff, so *"proxy configuration blocking app downloads"* did not
|
||||
retrieve that section at all.
|
||||
|
||||
So the *chunk* stays whole and the *embedding* is split instead:
|
||||
`VectorSearchProvider.index()` cuts long text into 250-token windows with 40 tokens
|
||||
of overlap and embeds each. A document therefore carries `embeddings: number[][]`,
|
||||
and `search()` scores it by its **best** window.
|
||||
|
||||
- **Max, not mean.** Averaging window vectors would blur a long section's several
|
||||
topics into one mediocre direction — the same dilution that makes large context
|
||||
windows worse than chunking. Max keeps a document findable by any one thing it
|
||||
discusses.
|
||||
- **The tokenizer is a build-time dependency only.** `index()` runs during the
|
||||
Docker build; production boots from precomputed vectors and never loads it. It is
|
||||
also only the vocabulary, not the ONNX weights (~93 ms).
|
||||
- Cost: 324 chunks → 424 vectors (+31%), index 2.63 → 3.12 MB, and 424 instead of
|
||||
324 dot products per query — still microseconds.
|
||||
- The rules corpus benefits too, unplanned: 128 rules now carry 172 vectors.
|
||||
|
||||
Measured effect: token coverage 87% → **100%**, *"Office Click-to-Run install never
|
||||
finished"* 0.318 → **0.442**, *"proxy configuration blocking app downloads"* absent
|
||||
→ **0.559 at rank 1**. The ten reference queries kept byte-identical scores — their
|
||||
chunks are short, hence one window, hence the same vector as before.
|
||||
|
||||
## Keyword fallback
|
||||
|
||||
Semantic ranking alone leaves a gap: proper nouns embed poorly. `RealmJoin` occurs
|
||||
in 3 chunks yet scores under the threshold.
|
||||
|
||||
Fuzzy matching is the wrong remedy, and the corpus says so. At a threshold low
|
||||
enough to surface `RealmJoin` (0.28), Fuse also returns 5 hits for `SCEPman` — a
|
||||
word this corpus contains **zero** times. Literal containment has neither problem.
|
||||
|
||||
So when the semantic pass returns fewer than `topK`, `search_docs` tops the list up
|
||||
via `lexicalMatch`, the substring scan every provider already exposes:
|
||||
|
||||
- Query terms are filtered to distinctive ones (≥4 characters, non-stopword, max 6).
|
||||
- A hit needs **corroboration**: at least two matching terms, unless the query has
|
||||
only one. Without this rule, "how do I bake sourdough bread" matched a page about
|
||||
LLM providers, because "bake" is a substring of "baked into the image".
|
||||
- Hits are **appended, never merged**: they carry no cosine similarity, so semantic
|
||||
results keep the head of the list and keyword hits fill the tail, flagged
|
||||
`matchType: "keyword"` with `keywordFallback` in the response.
|
||||
|
||||
It is not a rare path: across 904 rare corpus tokens queried on their own, the
|
||||
fallback contributed hits for 431. Multi-word natural-language questions rarely
|
||||
trigger it — those the semantic pass answers well. It earns its keep when someone
|
||||
pastes an identifier.
|
||||
|
||||
## Build-time embedding
|
||||
|
||||
Like the rules corpus, docs are embedded once at Docker build time by
|
||||
@@ -101,7 +162,7 @@ if the server *has* a docs corpus and the file does *not*, that is staleness and
|
||||
is rejected like any other mismatch.
|
||||
|
||||
Embeddings are rounded to 6 decimals on serialize. A float64 component
|
||||
serializes to ~20 characters of JSON, so 384 dims cost ~7.8 KB per document;
|
||||
serializes to ~20 characters of JSON, so 384 dims cost ~7.8 KB per vector;
|
||||
rounding roughly halves the index file, which is `JSON.parse`d on every
|
||||
scale-to-zero cold start. The induced L2 norm error (~1e-5 over 384 dims) sits
|
||||
two orders of magnitude inside `UNIT_NORM_EPSILON`.
|
||||
|
||||
@@ -26,7 +26,7 @@ const unit = (v: number[]): number[] => {
|
||||
};
|
||||
|
||||
const withEmbedding = (d: SearchDocument, embedding: number[] = [0.1, 0.2, 0.3]): PrecomputedDocument =>
|
||||
({ ...d, embedding: unit(embedding) });
|
||||
({ ...d, embeddings: [unit(embedding)] });
|
||||
|
||||
const KB_DOCS = [doc('rule-1', 'TPM not ready', { type: 'analyze-rule' }), doc('rule-2', 'ESP timeout', { type: 'analyze-rule' })];
|
||||
const ET_DOCS = [doc('et-1', 'bitlocker status', { eventType: 'bitlocker_status' })];
|
||||
@@ -110,7 +110,7 @@ describe('validatePrecomputedIndex', () => {
|
||||
|
||||
it('rejects inconsistent embedding dimensions', () => {
|
||||
const f = validFile();
|
||||
f.knowledgeBase.entries[1].embedding = [0.1, 0.2]; // 2 dims vs 3
|
||||
f.knowledgeBase.entries[1].embeddings = [[0.1, 0.2]]; // 2 dims vs 3
|
||||
expect(validatePrecomputedIndex(f, MODEL, KB_DOCS, ET_DOCS)).toMatchObject({
|
||||
ok: false, reason: expect.stringContaining('dimensions'),
|
||||
});
|
||||
@@ -157,14 +157,14 @@ describe('VectorSearchProvider serialize/hydrate roundtrip', () => {
|
||||
it('rounds embeddings on serialize but stays well inside the unit-norm tolerance', () => {
|
||||
const p = new VectorSearchProvider();
|
||||
const raw = unit([0.123456789012345, -0.987654321098765, 0.5555555555555]);
|
||||
p.indexPrecomputed([{ ...doc('d1', 'text'), embedding: raw }]);
|
||||
p.indexPrecomputed([{ ...doc('d1', 'text'), embeddings: [raw] }]);
|
||||
|
||||
const [out] = p.serialize();
|
||||
// Rounded to 6 decimals — the point of the exercise (halves the index JSON).
|
||||
expect(out.embedding.every((v) => v === Number(v.toFixed(6)))).toBe(true);
|
||||
expect(out.embedding).not.toEqual(raw);
|
||||
expect(out.embeddings[0].every((v) => v === Number(v.toFixed(6)))).toBe(true);
|
||||
expect(out.embeddings[0]).not.toEqual(raw);
|
||||
|
||||
const norm = Math.sqrt(out.embedding.reduce((s, x) => s + x * x, 0));
|
||||
const norm = Math.sqrt(out.embeddings[0].reduce((s, x) => s + x * x, 0));
|
||||
expect(Math.abs(norm - 1)).toBeLessThan(1e-5);
|
||||
// Re-hydrating the rounded vectors must not trip the assertion.
|
||||
expect(() => new VectorSearchProvider().indexPrecomputed(p.serialize())).not.toThrow();
|
||||
@@ -172,7 +172,7 @@ describe('VectorSearchProvider serialize/hydrate roundtrip', () => {
|
||||
|
||||
it('rejects a non-unit embedding on hydrate rather than skewing every score silently', () => {
|
||||
const p = new VectorSearchProvider();
|
||||
expect(() => p.indexPrecomputed([{ ...doc('d1', 'text'), embedding: [0.1, 0.2, 0.3] }]))
|
||||
expect(() => p.indexPrecomputed([{ ...doc('d1', 'text'), embeddings: [[0.1, 0.2, 0.3]] }]))
|
||||
.toThrow(/not L2-unit-normalized/);
|
||||
});
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@ function registeredToolNames(ga: boolean, strictGa: boolean = ga, delegated: boo
|
||||
const server = new McpServer({ name: 'test', version: '0.0.0' });
|
||||
// knowledgeBase / eventTypeIndex are optional — pass undefined so this stays a
|
||||
// pure unit test with no search-provider or backend dependency.
|
||||
registerTools(server, undefined, undefined, undefined, [], ga, strictGa, delegated);
|
||||
registerTools(server, undefined, undefined, undefined, ga, strictGa, delegated);
|
||||
const internal = server as unknown as { _registeredTools: Record<string, unknown> };
|
||||
return Object.keys(internal._registeredTools);
|
||||
}
|
||||
@@ -44,7 +44,7 @@ function stubDocsIndex(size = 3): SearchProvider {
|
||||
|
||||
function namesWithDocs(ga: boolean, strictGa: boolean = ga, delegated: boolean = false): string[] {
|
||||
const server = new McpServer({ name: 'test', version: '0.0.0' });
|
||||
registerTools(server, undefined, undefined, stubDocsIndex(), ['concepts', 'trust'], ga, strictGa, delegated);
|
||||
registerTools(server, undefined, undefined, { vector: stubDocsIndex(), sections: ['concepts', 'trust'] }, ga, strictGa, delegated);
|
||||
const internal = server as unknown as { _registeredTools: Record<string, unknown> };
|
||||
return Object.keys(internal._registeredTools);
|
||||
}
|
||||
@@ -65,7 +65,7 @@ describe('search_docs registration', () => {
|
||||
|
||||
it('is skipped for an empty corpus, so a doc-less build advertises no broken tool', () => {
|
||||
const server = new McpServer({ name: 'test', version: '0.0.0' });
|
||||
registerTools(server, undefined, undefined, stubDocsIndex(0), [], true, true, false);
|
||||
registerTools(server, undefined, undefined, { vector: stubDocsIndex(0), sections: [] }, true, true, false);
|
||||
const internal = server as unknown as { _registeredTools: Record<string, unknown> };
|
||||
expect(Object.keys(internal._registeredTools)).not.toContain('search_docs');
|
||||
});
|
||||
|
||||
@@ -27,7 +27,7 @@ const MANAGED = 'aaaa-1111-managed';
|
||||
/** Build a server for the given caller shape and return a tool's raw handler by name. */
|
||||
function handlerFor(name: string, opts: { ga?: boolean; strictGa?: boolean; delegated?: boolean }): ToolHandler {
|
||||
const server = new McpServer({ name: 'test', version: '0.0.0' });
|
||||
registerTools(server, undefined, undefined, undefined, [], opts.ga ?? false, opts.strictGa ?? (opts.ga ?? false), opts.delegated ?? false);
|
||||
registerTools(server, undefined, undefined, undefined, opts.ga ?? false, opts.strictGa ?? (opts.ga ?? false), opts.delegated ?? false);
|
||||
const registry = (server as unknown as { _registeredTools: Record<string, { handler: ToolHandler }> })._registeredTools;
|
||||
const tool = registry[name];
|
||||
if (!tool) throw new Error(`tool ${name} not registered for this caller shape`);
|
||||
|
||||
@@ -29,7 +29,7 @@ type ToolHandler = (args: Record<string, unknown>, extra: unknown) => Promise<{
|
||||
/** Build a strict-GA server and return a tool's raw handler by name. */
|
||||
function gaHandler(name: string): ToolHandler {
|
||||
const server = new McpServer({ name: 'test', version: '0.0.0' });
|
||||
registerTools(server, undefined, undefined, undefined, [], /*ga*/ true, /*strictGa*/ true, /*delegated*/ false);
|
||||
registerTools(server, undefined, undefined, undefined, /*ga*/ true, /*strictGa*/ true, /*delegated*/ false);
|
||||
const registry = (server as unknown as { _registeredTools: Record<string, { handler: ToolHandler }> })._registeredTools;
|
||||
const tool = registry[name];
|
||||
if (!tool) throw new Error(`tool ${name} not registered for strict-GA caller`);
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
/**
|
||||
* Unit tests for the search_docs handler — specifically the keyword fallback.
|
||||
*
|
||||
* Semantic search misses exact proper nouns that embed poorly. Measured on the real
|
||||
* corpus, "RealmJoin" occurs in 3 chunks yet scores below the semantic threshold,
|
||||
* while fuzzy matching at a threshold low enough to surface it also returns 5 hits
|
||||
* for "SCEPman", a word the corpus does not contain at all. The fallback therefore
|
||||
* uses literal containment, and the contract under test is:
|
||||
* appended after semantic hits, flagged, never re-sorted together, never invented.
|
||||
*
|
||||
* The provider is a stub: no model, no backend, no corpus on disk.
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
|
||||
import { registerTools } from '../tools.js';
|
||||
import { extractDistinctiveTerms } from '../tools/search.js';
|
||||
import type { SearchDocument, SearchProvider, SearchResult } from '../search-provider.js';
|
||||
|
||||
type ToolHandler = (args: Record<string, unknown>) => Promise<{ content: Array<{ text: string }> }>;
|
||||
|
||||
const chunk = (id: string, section: string, text = 'body'): SearchResult => ({
|
||||
id,
|
||||
text,
|
||||
metadata: { title: id, heading: null, section, path: `${section}/x.md`, url: null, tags: [] },
|
||||
score: 0.5,
|
||||
});
|
||||
|
||||
/**
|
||||
* Provider returning fixed lists, ignoring the query. `lexical` omitted entirely
|
||||
* models a provider that cannot expose document text (lexicalMatch is optional).
|
||||
*/
|
||||
function stubProvider(semantic: SearchResult[], lexical?: SearchResult[]): SearchProvider {
|
||||
return {
|
||||
name: 'stub',
|
||||
semanticCapable: true,
|
||||
size: Math.max(semantic.length, 1),
|
||||
index: async (_docs: SearchDocument[]) => {},
|
||||
search: async () => semantic,
|
||||
...(lexical ? { lexicalMatch: () => lexical } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
function searchDocsHandler(provider: SearchProvider): ToolHandler {
|
||||
const server = new McpServer({ name: 'test', version: '0.0.0' });
|
||||
registerTools(server, undefined, undefined, { vector: provider, sections: ['trust', 'reference'] }, true, true, false);
|
||||
const registry = (server as unknown as { _registeredTools: Record<string, { handler: ToolHandler }> })._registeredTools;
|
||||
return registry.search_docs.handler;
|
||||
}
|
||||
|
||||
async function run(handler: ToolHandler, args: Record<string, unknown>) {
|
||||
const res = await handler(args);
|
||||
return JSON.parse(res.content[0].text) as {
|
||||
resultCount: number;
|
||||
keywordFallback?: { used: boolean; added: number };
|
||||
results: Array<{ id: string; matchType?: string; section: string }>;
|
||||
};
|
||||
}
|
||||
|
||||
describe('extractDistinctiveTerms', () => {
|
||||
it('keeps product and setting names, drops filler', () => {
|
||||
expect(extractDistinctiveTerms('how does the RealmJoin watcher work'))
|
||||
.toEqual(['realmjoin', 'watcher', 'work']);
|
||||
});
|
||||
|
||||
it('drops words too short to be distinctive', () => {
|
||||
expect(extractDistinctiveTerms('who can do it')).toEqual([]);
|
||||
});
|
||||
|
||||
it('keeps identifier-shaped tokens intact', () => {
|
||||
expect(extractDistinctiveTerms('what is CollectorIdleTimeoutMinutes'))
|
||||
.toContain('collectoridletimeoutminutes');
|
||||
});
|
||||
|
||||
it('deduplicates and caps the term list', () => {
|
||||
const terms = extractDistinctiveTerms('alpha alpha beta gamma delta epsilon zeta eta theta');
|
||||
expect(new Set(terms).size).toBe(terms.length);
|
||||
expect(terms.length).toBeLessThanOrEqual(6);
|
||||
});
|
||||
});
|
||||
|
||||
describe('search_docs keyword fallback', () => {
|
||||
it('does not fire when semantic search already fills topK', async () => {
|
||||
const handler = searchDocsHandler(
|
||||
stubProvider([chunk('sem-1', 'trust'), chunk('sem-2', 'trust')], [chunk('kw-1', 'trust')]),
|
||||
);
|
||||
|
||||
const body = await run(handler, { query: 'notifications setup', topK: 2 });
|
||||
|
||||
expect(body.resultCount).toBe(2);
|
||||
expect(body.keywordFallback).toBeUndefined();
|
||||
expect(body.results.every((r) => r.matchType === undefined)).toBe(true);
|
||||
});
|
||||
|
||||
it('tops up a short result set and flags the added hits', async () => {
|
||||
const both = 'the realmjoin watcher toggle';
|
||||
const handler = searchDocsHandler(
|
||||
stubProvider([chunk('sem-1', 'trust')], [chunk('kw-1', 'trust', both), chunk('kw-2', 'trust', both)]),
|
||||
);
|
||||
|
||||
const body = await run(handler, { query: 'realmjoin watcher', topK: 3 });
|
||||
|
||||
expect(body.resultCount).toBe(3);
|
||||
expect(body.keywordFallback).toEqual({ used: true, added: 2 });
|
||||
// Semantic keeps the head; keyword fills the tail. Never interleaved — literal
|
||||
// hits carry no cosine similarity to sort against.
|
||||
expect(body.results.map((r) => r.matchType)).toEqual([undefined, 'keyword', 'keyword']);
|
||||
});
|
||||
|
||||
it('ranks fallback hits by how many query terms they contain', async () => {
|
||||
const handler = searchDocsHandler(
|
||||
stubProvider([], [
|
||||
chunk('two-terms', 'trust', 'mentions realmjoin and the watcher'),
|
||||
chunk('three-terms', 'trust', 'mentions realmjoin, the watcher and its toggle'),
|
||||
]),
|
||||
);
|
||||
|
||||
const body = await run(handler, { query: 'realmjoin watcher toggle', topK: 2 });
|
||||
|
||||
expect(body.results.map((r) => r.id)).toEqual(['three-terms', 'two-terms']);
|
||||
});
|
||||
|
||||
it('requires corroboration: one incidental term match is not a hit', async () => {
|
||||
// "bake" is a substring of "baked into the image" — matching ANY single term
|
||||
// would return a page about LLM providers for a sourdough question. Measured
|
||||
// on the real corpus before the corroboration rule was added.
|
||||
const handler = searchDocsHandler(
|
||||
stubProvider([], [chunk('incidental', 'trust', 'the corpus is baked into the image')]),
|
||||
);
|
||||
|
||||
const body = await run(handler, { query: 'how do I bake sourdough bread', topK: 3 });
|
||||
|
||||
expect(body.resultCount).toBe(0);
|
||||
});
|
||||
|
||||
it('never duplicates a document already returned semantically', async () => {
|
||||
const shared = chunk('same-doc', 'trust', 'realmjoin');
|
||||
const handler = searchDocsHandler(stubProvider([shared], [shared, chunk('kw-1', 'trust', 'realmjoin')]));
|
||||
|
||||
const body = await run(handler, { query: 'realmjoin', topK: 3 });
|
||||
|
||||
expect(body.results.map((r) => r.id)).toEqual(['same-doc', 'kw-1']);
|
||||
expect(body.keywordFallback).toEqual({ used: true, added: 1 });
|
||||
});
|
||||
|
||||
it('honours the section filter for fallback hits too', async () => {
|
||||
const handler = searchDocsHandler(
|
||||
stubProvider([], [chunk('kw-other', 'reference', 'realmjoin'), chunk('kw-match', 'trust', 'realmjoin')]),
|
||||
);
|
||||
|
||||
const body = await run(handler, { query: 'realmjoin', topK: 3, section: 'trust' });
|
||||
|
||||
expect(body.results.map((r) => r.id)).toEqual(['kw-match']);
|
||||
expect(body.results.every((r) => r.section === 'trust')).toBe(true);
|
||||
});
|
||||
|
||||
it('returns nothing when neither pass matches — no invented results', async () => {
|
||||
const handler = searchDocsHandler(stubProvider([], []));
|
||||
|
||||
const body = await run(handler, { query: 'how do I bake sourdough bread', topK: 5 });
|
||||
|
||||
expect(body.resultCount).toBe(0);
|
||||
expect(body.results).toEqual([]);
|
||||
expect(body.keywordFallback).toBeUndefined();
|
||||
});
|
||||
|
||||
it('skips the fallback when the query has no distinctive terms', async () => {
|
||||
const handler = searchDocsHandler(stubProvider([], [chunk('kw-1', 'trust')]));
|
||||
|
||||
const body = await run(handler, { query: 'who can do it', topK: 3 });
|
||||
|
||||
expect(body.resultCount).toBe(0);
|
||||
});
|
||||
|
||||
it('works with a provider that exposes no lexicalMatch at all', async () => {
|
||||
const handler = searchDocsHandler(stubProvider([chunk('sem-1', 'trust')]));
|
||||
|
||||
const body = await run(handler, { query: 'realmjoin watcher', topK: 3 });
|
||||
|
||||
expect(body.resultCount).toBe(1);
|
||||
expect(body.keywordFallback).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -5,7 +5,7 @@
|
||||
* assertUnitNorm directly with canned vectors — no ML model load required.
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { assertUnitNorm } from '../vector-search-provider.js';
|
||||
import { assertUnitNorm, splitIntoWindows, VectorSearchProvider, embed } from '../vector-search-provider.js';
|
||||
|
||||
/** Build an L2-unit vector of the given dimension (all components equal). */
|
||||
function unitVector(dim: number): number[] {
|
||||
@@ -46,3 +46,62 @@ describe('assertUnitNorm', () => {
|
||||
expect(() => assertUnitNorm([5, 0], 'index:doc-42')).toThrow(/index:doc-42/);
|
||||
});
|
||||
});
|
||||
|
||||
describe('splitIntoWindows', () => {
|
||||
it('returns short text unchanged as a single window', async () => {
|
||||
const text = 'Retention is per tenant and configurable, default 90 days.';
|
||||
await expect(splitIntoWindows(text)).resolves.toEqual([text]);
|
||||
});
|
||||
|
||||
it('splits text beyond the model context into several overlapping windows', async () => {
|
||||
// ~1200 distinct words — comfortably past the 250-token window.
|
||||
const long = Array.from({ length: 1200 }, (_, i) => `token${i}`).join(' ');
|
||||
const windows = await splitIntoWindows(long);
|
||||
|
||||
expect(windows.length).toBeGreaterThan(1);
|
||||
// Overlap: consecutive windows must share content, or a sentence sitting on a
|
||||
// boundary would be split across two vectors and match neither well.
|
||||
const tailOfFirst = windows[0].split(/\s+/).slice(-10);
|
||||
expect(tailOfFirst.some((w) => windows[1].includes(w))).toBe(true);
|
||||
// Coverage is the whole point: the last token must survive somewhere.
|
||||
expect(windows.some((w) => w.includes('token1199'))).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('VectorSearchProvider scoring across windows', () => {
|
||||
const unit3 = (v: number[]) => {
|
||||
const n = Math.sqrt(v.reduce((s, x) => s + x * x, 0));
|
||||
return v.map((x) => x / n);
|
||||
};
|
||||
|
||||
it('scores a document by its BEST window, not by its first', async () => {
|
||||
// Uses the real embedder to obtain a query vector, then plants that exact
|
||||
// direction as a document's SECOND window. Before windowing, content past the
|
||||
// first 250 tokens could not be matched at all — this pins that it now can.
|
||||
const query = 'delivery optimization download timeout';
|
||||
const target = await embed(query);
|
||||
const decoy = target.map((_, i) => (i === 0 ? 1 : 0)); // unit, but unrelated direction
|
||||
|
||||
const p = new VectorSearchProvider();
|
||||
p.indexPrecomputed([
|
||||
{ id: 'match-in-second-window', text: 'long doc', metadata: {}, embeddings: [decoy, target] },
|
||||
{ id: 'first-window-only', text: 'short doc', metadata: {}, embeddings: [decoy] },
|
||||
]);
|
||||
|
||||
const results = await p.search(query, { topK: 2, minScore: 0.5 });
|
||||
expect(results[0].id).toBe('match-in-second-window');
|
||||
// Cosine with itself is 1 — proof the second window, not the first, set the score.
|
||||
expect(results[0].score).toBeCloseTo(1, 5);
|
||||
expect(p.vectorCount).toBe(3);
|
||||
expect(p.size).toBe(2);
|
||||
});
|
||||
|
||||
it('counts vectors separately from documents', () => {
|
||||
const p = new VectorSearchProvider();
|
||||
p.indexPrecomputed([
|
||||
{ id: 'a', text: 'a', metadata: {}, embeddings: [unit3([1, 0, 0]), unit3([0, 1, 0]), unit3([0, 0, 1])] },
|
||||
]);
|
||||
expect(p.size).toBe(1);
|
||||
expect(p.vectorCount).toBe(3);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -165,8 +165,12 @@ if (precomputed) {
|
||||
console.error(`Search provider ready: ${knowledgeBase.name} — ${knowledgeBase.size} documents indexed${precomputed ? ' (precomputed)' : ''}.`);
|
||||
console.error(`Event-type index ready: ${eventTypeIndex.name} — ${eventTypeIndex.size} types indexed${precomputed ? ' (precomputed)' : ''}.`);
|
||||
if (docsIndex) {
|
||||
// Vectors exceed chunks whenever a chunk spans several token windows — surfacing
|
||||
// both makes a regression in the windowing visible in the boot log.
|
||||
const vectors = (docsIndex as { vectorCount?: number }).vectorCount;
|
||||
console.error(
|
||||
`Docs index ready: ${docsIndex.name} — ${docsIndex.size} chunks indexed${precomputed ? ' (precomputed)' : ''} ` +
|
||||
`Docs index ready: ${docsIndex.name} — ${docsIndex.size} chunks` +
|
||||
`${vectors ? ` / ${vectors} vectors` : ''} indexed${precomputed ? ' (precomputed)' : ''} ` +
|
||||
`(docs commit ${DOCS_COMMIT}).`,
|
||||
);
|
||||
} else {
|
||||
@@ -175,7 +179,7 @@ if (docsIndex) {
|
||||
|
||||
// Computed once: the tool catalog is rebuilt per request, and this is a pure
|
||||
// function of the (immutable) corpus.
|
||||
const DOCS_SECTIONS = docSections(docsDocs);
|
||||
const DOCS = docsIndex ? { vector: docsIndex, sections: docSections(docsDocs) } : undefined;
|
||||
|
||||
// Server-level guidance. The host surfaces this once per connection, so it is
|
||||
// the right home for cross-cutting strategy that would otherwise be duplicated
|
||||
@@ -227,7 +231,7 @@ function createMcpServer(ga: boolean, strictGa: boolean, delegated: boolean, man
|
||||
{ name: 'Autopilot-Monitor', version: SERVER_VERSION },
|
||||
{ instructions: buildInstructions(ga, delegated, managedTenants, homeTenantId) },
|
||||
);
|
||||
registerTools(s, knowledgeBase, eventTypeIndex, docsIndex, DOCS_SECTIONS, ga, strictGa, delegated);
|
||||
registerTools(s, knowledgeBase, eventTypeIndex, DOCS, ga, strictGa, delegated);
|
||||
registerResources(s);
|
||||
// A delegated caller has no platform scope, so prompts get the tenant-user surface (ga=false) —
|
||||
// the cross-tenant prompt wording would be misleading for a tenant-bounded MSP user.
|
||||
|
||||
@@ -78,12 +78,20 @@ function validateSection(
|
||||
if (section.entries.length !== sourceDocs.length) {
|
||||
return `${name}: entry count ${section.entries.length} != corpus size ${sourceDocs.length}`;
|
||||
}
|
||||
const dims = section.entries[0]?.embedding?.length ?? 0;
|
||||
// Each entry carries one vector per token window (see splitIntoWindows) — usually
|
||||
// one, more for long documents. Dimensionality must be uniform across every window
|
||||
// of every entry: a ragged file would otherwise produce silently wrong dot products.
|
||||
const dims = section.entries[0]?.embeddings?.[0]?.length ?? 0;
|
||||
if (dims === 0) return `${name}: first entry has no embedding`;
|
||||
for (const e of section.entries) {
|
||||
if (typeof e.id !== 'string' || typeof e.text !== 'string') return `${name}: entry missing id/text`;
|
||||
if (!Array.isArray(e.embedding) || e.embedding.length !== dims) {
|
||||
return `${name}: inconsistent embedding dimensions`;
|
||||
if (!Array.isArray(e.embeddings) || e.embeddings.length === 0) {
|
||||
return `${name}: entry ${e.id} has no embeddings`;
|
||||
}
|
||||
for (const vec of e.embeddings) {
|
||||
if (!Array.isArray(vec) || vec.length !== dims) {
|
||||
return `${name}: inconsistent embedding dimensions`;
|
||||
}
|
||||
}
|
||||
}
|
||||
return null;
|
||||
|
||||
@@ -90,3 +90,16 @@ export function scanLexical(documents: SearchDocument[], needles: string[]): Sea
|
||||
// ── Provider identifiers ─────────────────────────────────────
|
||||
|
||||
export type SearchBackend = 'vector' | 'fuse';
|
||||
|
||||
/**
|
||||
* What `search_docs` runs on, as one value rather than a growing tail of positional
|
||||
* parameters through registerTools → registerSearchTools.
|
||||
*
|
||||
* The keyword fallback needs no second provider: it uses `vector.lexicalMatch`,
|
||||
* the literal-containment scan every provider already exposes.
|
||||
*/
|
||||
export interface DocsSearchBundle {
|
||||
vector: SearchProvider;
|
||||
/** Top-level documentation areas present in the corpus — the `section` filter hint. */
|
||||
sections: string[];
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ import type { SearchProvider } from './search-provider.js';
|
||||
import { registerSessionTools } from './tools/sessions.js';
|
||||
import { registerSearchTools } from './tools/search.js';
|
||||
import { registerAdminTools } from './tools/admin.js';
|
||||
import type { DocsSearchBundle } from './search-provider.js';
|
||||
|
||||
/**
|
||||
* Registers the tool catalog for a single request, tailored to the caller's role.
|
||||
@@ -26,14 +27,13 @@ export function registerTools(
|
||||
server: McpServer,
|
||||
knowledgeBase: SearchProvider | undefined,
|
||||
eventTypeIndex: SearchProvider | undefined,
|
||||
docsIndex: SearchProvider | undefined,
|
||||
docsSections: string[],
|
||||
docs: DocsSearchBundle | undefined,
|
||||
ga: boolean,
|
||||
strictGa: boolean = ga,
|
||||
delegated: boolean = false,
|
||||
): void {
|
||||
registerSessionTools(server, ga, delegated);
|
||||
registerSearchTools(server, knowledgeBase, eventTypeIndex, docsIndex, docsSections, ga, delegated);
|
||||
registerSearchTools(server, knowledgeBase, eventTypeIndex, docs, ga, delegated);
|
||||
registerAdminTools(server, ga, strictGa, delegated);
|
||||
sortToolCatalog(server);
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
|
||||
import { z } from 'zod';
|
||||
import { apiFetch, buildQuery, enforceDelegatedTenant, followNextLink, pickGlobalOrTenantPath } from '../client.js';
|
||||
import { withToolTelemetry } from '../telemetry.js';
|
||||
import type { SearchProvider } from '../search-provider.js';
|
||||
import type { SearchProvider, DocsSearchBundle } from '../search-provider.js';
|
||||
import { READ_ONLY, MAX_RESULT_SIZE_CHARS, toolResultText, SessionIdSchema, isBenignHealthDetectionReport, tenantIdDescription } from './shared.js';
|
||||
import { toolError } from './error-handler.js';
|
||||
import { ALL_EVENT_TYPES } from '../resource-catalog.js';
|
||||
@@ -959,15 +959,42 @@ export function diversifyBySession<T extends { event: EventEntry }>(
|
||||
return picked;
|
||||
}
|
||||
|
||||
/**
|
||||
* Terms from a natural-language query that are worth matching literally.
|
||||
*
|
||||
* Only distinctive words: short and common ones ("how", "the", "does") appear in
|
||||
* nearly every chunk, so matching on them would return the whole corpus in
|
||||
* arbitrary order. What is left is the vocabulary literal matching is actually good
|
||||
* at — product names, setting names, identifiers — which is exactly what embeddings
|
||||
* are bad at.
|
||||
*/
|
||||
const KEYWORD_STOPWORDS = new Set([
|
||||
'about', 'after', 'again', 'also', 'been', 'before', 'being', 'between', 'both', 'can',
|
||||
'could', 'does', 'doing', 'during', 'each', 'from', 'have', 'having', 'here', 'how',
|
||||
'into', 'just', 'more', 'most', 'much', 'must', 'need', 'only', 'other', 'over',
|
||||
'same', 'should', 'some', 'such', 'than', 'that', 'their', 'them', 'then', 'there',
|
||||
'these', 'they', 'this', 'those', 'through', 'under', 'until', 'very', 'were', 'what',
|
||||
'when', 'where', 'which', 'while', 'will', 'with', 'would', 'your',
|
||||
]);
|
||||
|
||||
export function extractDistinctiveTerms(query: string): string[] {
|
||||
const seen = new Set<string>();
|
||||
for (const raw of query.toLowerCase().split(/[^a-z0-9_.-]+/)) {
|
||||
const term = raw.replace(/^[.-]+|[.-]+$/g, '');
|
||||
if (term.length < 4 || KEYWORD_STOPWORDS.has(term)) continue;
|
||||
seen.add(term);
|
||||
if (seen.size >= 6) break;
|
||||
}
|
||||
return [...seen];
|
||||
}
|
||||
|
||||
// ── Registration ────────────────────────────────────────────────────────
|
||||
|
||||
export function registerSearchTools(
|
||||
server: McpServer,
|
||||
knowledgeBase: SearchProvider | undefined,
|
||||
eventTypeIndex: SearchProvider | undefined,
|
||||
docsIndex: SearchProvider | undefined,
|
||||
/** Top-level documentation areas present in the corpus — surfaced as the `section` filter hint. */
|
||||
docsSections: string[],
|
||||
docs: DocsSearchBundle | undefined,
|
||||
ga: boolean,
|
||||
delegated: boolean = false,
|
||||
): void {
|
||||
@@ -1186,7 +1213,8 @@ export function registerSearchTools(
|
||||
// normal tenant user has exactly as much business asking "how do I deploy the
|
||||
// agent" as a Global Admin. It is skipped entirely when no corpus was baked
|
||||
// into the image, so a doc-less build advertises no broken tool.
|
||||
if (docsIndex && docsIndex.size > 0) {
|
||||
if (docs && docs.vector.size > 0) {
|
||||
const docsIndex = docs.vector;
|
||||
server.registerTool(
|
||||
'search_docs',
|
||||
{
|
||||
@@ -1197,8 +1225,16 @@ export function registerSearchTools(
|
||||
'Use this for questions about how the PRODUCT works rather than what a specific enrollment did: ' +
|
||||
'setup and onboarding, agent deployment, portal features, roles and permissions, settings, ' +
|
||||
'notifications, network endpoints, security/privacy/data-residency, plans, and troubleshooting guidance. ' +
|
||||
'PREFER THIS OVER FETCHING OR WEB-SEARCHING THE DOCS SITE: it returns the matching SECTIONS ' +
|
||||
'(~1k tokens for topK=3) instead of whole pages (~9k tokens raw, ~600k rendered), it is the exact ' +
|
||||
'version this server ships rather than a cached or third-party copy, and it matches on meaning — ' +
|
||||
'"where is my data stored" finds the page that says "Germany West Central" with no shared words, ' +
|
||||
'which keyword or web search cannot do. ' +
|
||||
'Each result carries the page `title`, its `heading` breadcrumb and a citable `url` — quote the url when ' +
|
||||
'answering so the user can verify. ' +
|
||||
'If semantic matching returns too few hits, exact keyword matches are APPENDED (flagged ' +
|
||||
'`matchType: "keyword"`, and their scores are on a different scale — rank order across the two is ' +
|
||||
'not comparable). ' +
|
||||
'NOT the same corpus as search_knowledge: that one holds analysis/gather rules and IME log patterns and ' +
|
||||
'explains why an enrollment failed. Rule of thumb — "how does X work?" → search_docs; ' +
|
||||
'"why did this session fail?" → search_knowledge or get_session_summary.',
|
||||
@@ -1207,7 +1243,7 @@ export function registerSearchTools(
|
||||
topK: z.coerce.number().min(1).max(20).optional().default(5).describe('Number of results to return (1-20, default 5)'),
|
||||
section: z.string().optional()
|
||||
.describe('Restrict to one documentation area, matching the bundle\'s top-level folder ' +
|
||||
`(one of: ${docsSections.join(', ')}). Omit to search everything.`),
|
||||
`(one of: ${docs.sections.join(', ')}). Omit to search everything.`),
|
||||
minScore: z.coerce.number().min(0).max(1).optional().default(0.25)
|
||||
.describe('Minimum similarity score (0-1, default 0.25). Short keyword queries score low on ' +
|
||||
'all-MiniLM embeddings, so the default is tuned to keep the relevant-but-marginal band.'),
|
||||
@@ -1243,6 +1279,48 @@ export function registerSearchTools(
|
||||
}
|
||||
results = results.slice(0, topK);
|
||||
|
||||
// Keyword fallback: exact, literal term matching to top up a short result set.
|
||||
//
|
||||
// Proper nouns embed poorly — measured on this corpus, "RealmJoin" appears in
|
||||
// 3 chunks yet scores below the semantic threshold. Fuzzy matching is the wrong
|
||||
// remedy: at a threshold low enough to surface "RealmJoin" (0.28), Fuse also
|
||||
// returns 5 hits for "SCEPman", a word that appears in the corpus zero times.
|
||||
// Literal containment has neither problem — it finds the first and returns
|
||||
// nothing for the second.
|
||||
//
|
||||
// APPENDED, never merged: these carry no cosine similarity, so semantic hits
|
||||
// keep the head of the list and keyword hits fill the tail, flagged for the
|
||||
// caller. Ranked among themselves by how many query terms they contain.
|
||||
const keywordHitIds = new Set<string>();
|
||||
const terms = extractDistinctiveTerms(query);
|
||||
if (results.length < topK && terms.length > 0 && docsIndex.lexicalMatch) {
|
||||
// lexicalMatch is ANY-of-needles, which alone is far too generous: for
|
||||
// "how do I bake sourdough bread" the single word "bake" matches "baked
|
||||
// into the image" and a page about LLM providers comes back as a hit.
|
||||
// A multi-term query must therefore corroborate — at least two of its
|
||||
// terms in the same chunk. A single-term query ("RealmJoin") still needs
|
||||
// just that one, which is the case this fallback exists for.
|
||||
const required = Math.min(2, terms.length);
|
||||
const seen = new Set(results.map((r) => r.id));
|
||||
const ranked = docsIndex
|
||||
.lexicalMatch(terms)
|
||||
.map((hit) => {
|
||||
const haystack = hit.text.toLowerCase();
|
||||
return { hit, matched: terms.filter((t) => haystack.includes(t)).length };
|
||||
})
|
||||
.filter(({ matched }) => matched >= required)
|
||||
.sort((a, b) => b.matched - a.matched);
|
||||
|
||||
for (const { hit } of ranked) {
|
||||
if (results.length >= topK) break;
|
||||
if (seen.has(hit.id)) continue;
|
||||
if (section && hit.metadata.section !== section) continue;
|
||||
seen.add(hit.id);
|
||||
keywordHitIds.add(hit.id);
|
||||
results.push(hit);
|
||||
}
|
||||
}
|
||||
|
||||
const formatted = results.map((r) => ({
|
||||
id: r.id,
|
||||
score: Math.round(r.score * 1000) / 1000,
|
||||
@@ -1253,7 +1331,11 @@ export function registerSearchTools(
|
||||
url: r.metadata.url,
|
||||
tags: r.metadata.tags,
|
||||
content: r.text,
|
||||
matchType: errorCodeHitIds.has(r.id) ? ('error-code' as const) : undefined,
|
||||
matchType: errorCodeHitIds.has(r.id)
|
||||
? ('error-code' as const)
|
||||
: keywordHitIds.has(r.id)
|
||||
? ('keyword' as const)
|
||||
: undefined,
|
||||
}));
|
||||
|
||||
return toolResultText({
|
||||
@@ -1262,6 +1344,7 @@ export function registerSearchTools(
|
||||
...(section ? { section } : {}),
|
||||
resultCount: formatted.length,
|
||||
...(needles.length > 0 ? { errorCodeFallback: { codes: needles, matchedCount: errorCodeHitIds.size } } : {}),
|
||||
...(keywordHitIds.size > 0 ? { keywordFallback: { used: true, added: keywordHitIds.size } } : {}),
|
||||
results: formatted,
|
||||
}, MAX_RESULT_SIZE_CHARS.small);
|
||||
} catch (error: unknown) {
|
||||
|
||||
@@ -45,6 +45,68 @@ function getEmbedder(): Promise<FeatureExtractionPipeline> {
|
||||
return embedderPromise;
|
||||
}
|
||||
|
||||
// ── Token windowing ──────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* The model's context window is 256 tokens INCLUDING the [CLS]/[SEP] specials, so
|
||||
* windows carry at most 250 content tokens.
|
||||
*
|
||||
* Anything past the window is silently dropped by the tokenizer — not an error, just
|
||||
* a vector that represents the opening of a document and nothing else. Measured on
|
||||
* the docs corpus, 87 of 324 chunks exceeded it and 13% of all tokens never reached
|
||||
* the ranking. Splitting long text into several windows and scoring by the best one
|
||||
* (see search()) puts every token back in play without shrinking chunks, which would
|
||||
* break up the tables and breadcrumbs the chunker deliberately keeps whole.
|
||||
*/
|
||||
const WINDOW_TOKENS = 250;
|
||||
|
||||
/**
|
||||
* Overlap between consecutive windows. A sentence that straddles a boundary would
|
||||
* otherwise be split across two vectors and match neither well.
|
||||
*/
|
||||
const WINDOW_OVERLAP_TOKENS = 40;
|
||||
|
||||
// Tokenizer only — NOT the ONNX weights. Needed exclusively on the indexing path
|
||||
// (build time); the production boot hydrates precomputed vectors and never calls
|
||||
// index(), so this stays unloaded there. Measured cold load: ~93 ms.
|
||||
let tokenizerPromise: Promise<{
|
||||
(text: string, opts?: object): Promise<{ input_ids: { data: BigInt64Array | Int32Array } }>;
|
||||
decode(ids: number[], opts?: object): string;
|
||||
}> | null = null;
|
||||
|
||||
function getTokenizer() {
|
||||
if (!tokenizerPromise) {
|
||||
tokenizerPromise = (async () => {
|
||||
const { AutoTokenizer, env } = await import('@huggingface/transformers');
|
||||
if (process.env.HF_CACHE_DIR) env.cacheDir = process.env.HF_CACHE_DIR;
|
||||
return (await AutoTokenizer.from_pretrained(MODEL_NAME)) as never;
|
||||
})().catch((err) => {
|
||||
tokenizerPromise = null;
|
||||
throw err;
|
||||
});
|
||||
}
|
||||
return tokenizerPromise;
|
||||
}
|
||||
|
||||
/**
|
||||
* Split `text` into overlapping token windows, returned as decoded strings.
|
||||
* Short text yields exactly one window (the original string, untouched).
|
||||
*/
|
||||
export async function splitIntoWindows(text: string): Promise<string[]> {
|
||||
const tokenizer = await getTokenizer();
|
||||
const encoded = await tokenizer(text, { add_special_tokens: false });
|
||||
const ids = Array.from(encoded.input_ids.data as ArrayLike<bigint | number>, Number);
|
||||
if (ids.length <= WINDOW_TOKENS) return [text];
|
||||
|
||||
const stride = WINDOW_TOKENS - WINDOW_OVERLAP_TOKENS;
|
||||
const windows: string[] = [];
|
||||
for (let start = 0; start < ids.length; start += stride) {
|
||||
windows.push(tokenizer.decode(ids.slice(start, start + WINDOW_TOKENS), { skip_special_tokens: true }));
|
||||
if (start + WINDOW_TOKENS >= ids.length) break;
|
||||
}
|
||||
return windows;
|
||||
}
|
||||
|
||||
/**
|
||||
* Every embedding this module produces MUST be L2-unit-normalized: cosineSimilarity is a
|
||||
* bare dot product (see below) that silently returns wrong scores for non-unit inputs. We
|
||||
@@ -105,9 +167,16 @@ function cosineSimilarity(a: number[], b: number[]): number {
|
||||
|
||||
// ── Provider implementation ──────────────────────────────────
|
||||
|
||||
/** An indexed document with its embedding — the unit of build-time serialization. */
|
||||
/**
|
||||
* An indexed document with its embeddings — the unit of build-time serialization.
|
||||
*
|
||||
* One vector per token window (see splitIntoWindows). Short documents — which is
|
||||
* every rule and every event type — have exactly one, so the array is not an
|
||||
* optimization for the common case but the single uniform shape: no branch, no
|
||||
* "primary vs extra" special case to get wrong.
|
||||
*/
|
||||
export interface PrecomputedDocument extends SearchDocument {
|
||||
embedding: number[];
|
||||
embeddings: number[][];
|
||||
}
|
||||
|
||||
export class VectorSearchProvider implements SearchProvider {
|
||||
@@ -131,10 +200,17 @@ export class VectorSearchProvider implements SearchProvider {
|
||||
// a bare dot product, so a non-unit vector here would silently skew every
|
||||
// score rather than fail. One check is enough — the whole section is written
|
||||
// by a single serialize() call, so they stand or fall together.
|
||||
if (docs.length > 0) assertUnitNorm(docs[0].embedding, 'indexPrecomputed');
|
||||
if (docs.length > 0 && docs[0].embeddings.length > 0) {
|
||||
assertUnitNorm(docs[0].embeddings[0], 'indexPrecomputed');
|
||||
}
|
||||
this.documents.push(...docs);
|
||||
}
|
||||
|
||||
/** Total vectors held — exceeds `size` whenever documents span several windows. */
|
||||
get vectorCount(): number {
|
||||
return this.documents.reduce((n, d) => n + d.embeddings.length, 0);
|
||||
}
|
||||
|
||||
/**
|
||||
* Export the indexed documents with embeddings for build-time serialization.
|
||||
*
|
||||
@@ -149,7 +225,7 @@ export class VectorSearchProvider implements SearchProvider {
|
||||
serialize(): PrecomputedDocument[] {
|
||||
return this.documents.map((doc) => ({
|
||||
...doc,
|
||||
embedding: doc.embedding.map((v) => Number(v.toFixed(SERIALIZED_PRECISION))),
|
||||
embeddings: doc.embeddings.map((vec) => vec.map((v) => Number(v.toFixed(SERIALIZED_PRECISION)))),
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -160,14 +236,21 @@ export class VectorSearchProvider implements SearchProvider {
|
||||
const batch = docs.slice(i, i + batchSize);
|
||||
const embeddings = await Promise.all(
|
||||
batch.map(async (d) => {
|
||||
const out = await model(d.text, { pooling: 'mean', normalize: true });
|
||||
const vec = Array.from(out.data as Float32Array);
|
||||
assertUnitNorm(vec, `index:${d.id}`);
|
||||
return vec;
|
||||
// One vector per token window: text longer than the model's context would
|
||||
// otherwise be truncated to its opening and the rest never ranked at all.
|
||||
const windows = await splitIntoWindows(d.text);
|
||||
return Promise.all(
|
||||
windows.map(async (window, w) => {
|
||||
const out = await model(window, { pooling: 'mean', normalize: true });
|
||||
const vec = Array.from(out.data as Float32Array);
|
||||
assertUnitNorm(vec, `index:${d.id}#w${w}`);
|
||||
return vec;
|
||||
})
|
||||
);
|
||||
})
|
||||
);
|
||||
for (let j = 0; j < batch.length; j++) {
|
||||
this.documents.push({ ...batch[j], embedding: embeddings[j] });
|
||||
this.documents.push({ ...batch[j], embeddings: embeddings[j] });
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -176,12 +259,18 @@ export class VectorSearchProvider implements SearchProvider {
|
||||
const { topK = 5, minScore = 0.3 } = options;
|
||||
const queryEmbedding = await embed(query);
|
||||
|
||||
const scored = this.documents.map((doc) => ({
|
||||
id: doc.id,
|
||||
text: doc.text,
|
||||
metadata: doc.metadata,
|
||||
score: cosineSimilarity(queryEmbedding, doc.embedding),
|
||||
}));
|
||||
// Score a document by its BEST window, not by an average. Averaging would blur a
|
||||
// long document's several topics into one mediocre vector — the exact dilution
|
||||
// that makes big context windows worse than chunking. Max keeps a document
|
||||
// findable by any one of the things it actually talks about.
|
||||
const scored = this.documents.map((doc) => {
|
||||
let best = -Infinity;
|
||||
for (const vec of doc.embeddings) {
|
||||
const score = cosineSimilarity(queryEmbedding, vec);
|
||||
if (score > best) best = score;
|
||||
}
|
||||
return { id: doc.id, text: doc.text, metadata: doc.metadata, score: best };
|
||||
});
|
||||
|
||||
return scored
|
||||
.filter((r) => r.score >= minScore)
|
||||
|
||||
Reference in New Issue
Block a user