@maxgfr/codeindex 2.9.0 → 2.10.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +25 -0
- package/docs/MIGRATION.md +29 -0
- package/package.json +1 -1
- package/scripts/engine.d.mts +57 -2
- package/scripts/engine.mjs +474 -23
- package/src/embed/encode.ts +133 -0
- package/src/embed/endpoint.ts +55 -0
- package/src/embed/index.ts +125 -0
- package/src/embed/model.ts +117 -0
- package/src/embed/search.ts +90 -0
- package/src/engine-cli.ts +107 -5
- package/src/engine.ts +21 -0
- package/src/mcp.ts +43 -6
- package/src/types.ts +1 -1
|
@@ -0,0 +1,133 @@
|
|
|
1
|
+
import { foldText } from "../util.js";
|
|
2
|
+
import type { StaticEmbedModel } from "./model.js";
|
|
3
|
+
|
|
4
|
+
// The pure-JS deterministic encoder — the reason this whole tier can produce
|
|
5
|
+
// goldens. Every step is chosen to be byte-identical on any platform:
|
|
6
|
+
// 1. tokenize : foldText (NFKD, strip combining marks) + camelCase/ACRONYM
|
|
7
|
+
// split + lowercase + split on non-alphanumeric runs.
|
|
8
|
+
// 2. wordpiece : greedy longest-match against the model vocab, with BERT-style
|
|
9
|
+
// "##" continuation pieces; an unsplittable word → the UNK row
|
|
10
|
+
// (or nothing when the model has no UNK).
|
|
11
|
+
// 3. gather : look up each token's row (double precision).
|
|
12
|
+
// 4. mean-pool : sum rows in FIXED token order, divide by token count. IEEE-754
|
|
13
|
+
// double add/div is deterministic for a fixed evaluation order.
|
|
14
|
+
// 5. L2-norm : divide by sqrt(Σ v²) — sqrt is correctly-rounded per IEEE-754,
|
|
15
|
+
// so identical across engines.
|
|
16
|
+
// 6. quantize : multiply the unit vector by 127, round-half-to-EVEN (NOT
|
|
17
|
+
// Math.round, which is half-up), clamp to [-127, 127] → int8.
|
|
18
|
+
// The fixed 1/127 scale means ranking is a PURE INTEGER dot product (no per-vector
|
|
19
|
+
// float scale to store or multiply): for a fixed query every corpus vector shares
|
|
20
|
+
// the same scale, so the integer dot Σ q·c is monotonic with the cosine estimate.
|
|
21
|
+
|
|
22
|
+
const QUANT = 127;
|
|
23
|
+
|
|
24
|
+
// Split raw text into lowercase, diacritic-folded alphanumeric words. camelCase
|
|
25
|
+
// and ACRONYMWord boundaries become spaces FIRST (so "verifyAuthToken" exposes
|
|
26
|
+
// "verify","auth","token" as separate words to wordpiece), then any
|
|
27
|
+
// non-alphanumeric run splits. Mirrors the boundary rules bm25 `subtokens` uses,
|
|
28
|
+
// keeping the two search tiers tokenizing code identifiers alike.
|
|
29
|
+
export function basicTokenize(text: string): string[] {
|
|
30
|
+
const spaced = foldText(text)
|
|
31
|
+
.replace(/([a-z0-9])([A-Z])/g, "$1 $2")
|
|
32
|
+
.replace(/([A-Z]+)([A-Z][a-z])/g, "$1 $2");
|
|
33
|
+
const out: string[] = [];
|
|
34
|
+
for (const part of spaced.toLowerCase().split(/[^a-z0-9]+/)) {
|
|
35
|
+
if (part) out.push(part);
|
|
36
|
+
}
|
|
37
|
+
return out;
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
// Greedy longest-match wordpiece for ONE word. The first piece is looked up as
|
|
41
|
+
// written; every subsequent piece is looked up with a "##" prefix. If no piece
|
|
42
|
+
// matches at some position the whole word is unresolvable → [unkId] (or [] when
|
|
43
|
+
// the model has no UNK row). Deterministic: longest-match is a fixed scan.
|
|
44
|
+
export function wordpiece(word: string, model: StaticEmbedModel): number[] {
|
|
45
|
+
if (!word) return [];
|
|
46
|
+
const ids: number[] = [];
|
|
47
|
+
let start = 0;
|
|
48
|
+
const n = word.length;
|
|
49
|
+
while (start < n) {
|
|
50
|
+
let end = n;
|
|
51
|
+
let match = -1;
|
|
52
|
+
while (end > start) {
|
|
53
|
+
const piece = start === 0 ? word.slice(start, end) : "##" + word.slice(start, end);
|
|
54
|
+
const id = model.vocab.get(piece);
|
|
55
|
+
if (id !== undefined) {
|
|
56
|
+
match = id;
|
|
57
|
+
break;
|
|
58
|
+
}
|
|
59
|
+
end--;
|
|
60
|
+
}
|
|
61
|
+
if (match === -1) return model.unkId >= 0 ? [model.unkId] : []; // whole word is OOV
|
|
62
|
+
ids.push(match);
|
|
63
|
+
start = end;
|
|
64
|
+
}
|
|
65
|
+
return ids;
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
// Full token-id sequence for a text: tokenize → wordpiece each word, concatenated
|
|
69
|
+
// in order. The order is load-bearing for determinism, but mean-pooling is
|
|
70
|
+
// order-independent anyway; the fixed order simply makes reasoning trivial.
|
|
71
|
+
export function tokenize(text: string, model: StaticEmbedModel): number[] {
|
|
72
|
+
const ids: number[] = [];
|
|
73
|
+
for (const word of basicTokenize(text)) {
|
|
74
|
+
for (const id of wordpiece(word, model)) ids.push(id);
|
|
75
|
+
}
|
|
76
|
+
return ids;
|
|
77
|
+
}
|
|
78
|
+
|
|
79
|
+
// Round half to EVEN (banker's rounding). JS `Math.round` rounds half toward
|
|
80
|
+
// +∞, which is not symmetric and would bias quantization; half-to-even is the
|
|
81
|
+
// IEEE-754 default and keeps quantize(x) === -quantize(-x).
|
|
82
|
+
export function roundHalfToEven(x: number): number {
|
|
83
|
+
const f = Math.floor(x);
|
|
84
|
+
const diff = x - f;
|
|
85
|
+
if (diff < 0.5) return f;
|
|
86
|
+
if (diff > 0.5) return f + 1;
|
|
87
|
+
return f % 2 === 0 ? f : f + 1; // exactly .5 → nearest even
|
|
88
|
+
}
|
|
89
|
+
|
|
90
|
+
// Encode a text to an int8 unit-ish vector. Empty / all-OOV / zero-norm inputs
|
|
91
|
+
// yield an all-zero vector (dot-products against it are 0 — it simply ranks
|
|
92
|
+
// last), never NaN. Pure and deterministic: same model + text → same bytes.
|
|
93
|
+
export function encode(model: StaticEmbedModel, text: string): Int8Array {
|
|
94
|
+
const { dim, weights } = model;
|
|
95
|
+
const out = new Int8Array(dim);
|
|
96
|
+
const ids = tokenize(text, model);
|
|
97
|
+
if (ids.length === 0) return out;
|
|
98
|
+
|
|
99
|
+
// mean-pool in double precision, fixed order.
|
|
100
|
+
const pooled = new Float64Array(dim);
|
|
101
|
+
for (const id of ids) {
|
|
102
|
+
const base = id * dim;
|
|
103
|
+
for (let d = 0; d < dim; d++) pooled[d]! += weights[base + d]!;
|
|
104
|
+
}
|
|
105
|
+
const inv = 1 / ids.length;
|
|
106
|
+
for (let d = 0; d < dim; d++) pooled[d]! *= inv;
|
|
107
|
+
|
|
108
|
+
// L2-normalize.
|
|
109
|
+
let sumsq = 0;
|
|
110
|
+
for (let d = 0; d < dim; d++) sumsq += pooled[d]! * pooled[d]!;
|
|
111
|
+
const norm = Math.sqrt(sumsq);
|
|
112
|
+
if (norm === 0) return out; // zero vector — nothing to rank on
|
|
113
|
+
|
|
114
|
+
// quantize to int8 at the fixed 1/127 scale, round-half-to-even, clamp.
|
|
115
|
+
for (let d = 0; d < dim; d++) {
|
|
116
|
+
let q = roundHalfToEven((pooled[d]! / norm) * QUANT);
|
|
117
|
+
if (q > QUANT) q = QUANT;
|
|
118
|
+
else if (q < -QUANT) q = -QUANT;
|
|
119
|
+
out[d] = q;
|
|
120
|
+
}
|
|
121
|
+
return out;
|
|
122
|
+
}
|
|
123
|
+
|
|
124
|
+
// Integer dot product of two int8 vectors (the ranking primitive). Widened to a
|
|
125
|
+
// JS number (double), but every operand and partial sum is an exact integer
|
|
126
|
+
// (max |Σ| for dim 512 ≈ 512·127² ≈ 8.3M, far inside 2⁵³) — so this is an EXACT
|
|
127
|
+
// integer computation, deterministic everywhere.
|
|
128
|
+
export function intDot(a: Int8Array, b: Int8Array): number {
|
|
129
|
+
const n = Math.min(a.length, b.length);
|
|
130
|
+
let dot = 0;
|
|
131
|
+
for (let i = 0; i < n; i++) dot += a[i]! * b[i]!;
|
|
132
|
+
return dot;
|
|
133
|
+
}
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
// HTTP embedding endpoint client — the "rich" tier SKELETON for v2.11.0. It is
|
|
2
|
+
// intentionally NOT wired into the CLI or MCP in v2.10.0: the deterministic
|
|
3
|
+
// static tier is the default, and an endpoint's vectors are float and provider-
|
|
4
|
+
// dependent (NOT byte-deterministic), so this stays an explicit, opt-in escape
|
|
5
|
+
// hatch a future release finalizes. Shipped now only so the module surface is
|
|
6
|
+
// stable and the contract is documented.
|
|
7
|
+
//
|
|
8
|
+
// The contract (v2.11): POST { texts: string[] } → { vectors: number[][] } to
|
|
9
|
+
// CODEINDEX_EMBED_ENDPOINT (or opts.url). Node >=18 has global fetch, so this
|
|
10
|
+
// stays zero-dependency. The library never orchestrates docker; running such a
|
|
11
|
+
// server (e.g. ghcr.io/maxgfr/codeindex-embed) is a user/CI concern.
|
|
12
|
+
|
|
13
|
+
export interface EmbedEndpointOptions {
|
|
14
|
+
// Endpoint URL. Falls back to CODEINDEX_EMBED_ENDPOINT when omitted.
|
|
15
|
+
url?: string;
|
|
16
|
+
// Abort the request after this many ms (default 30_000).
|
|
17
|
+
timeoutMs?: number;
|
|
18
|
+
// Extra request headers (e.g. an auth token for a private endpoint).
|
|
19
|
+
headers?: Record<string, string>;
|
|
20
|
+
}
|
|
21
|
+
|
|
22
|
+
// Resolve the configured endpoint URL, or undefined when neither opts.url nor
|
|
23
|
+
// CODEINDEX_EMBED_ENDPOINT is set — the silent "endpoint tier unavailable" signal.
|
|
24
|
+
export function resolveEmbedEndpoint(opts: EmbedEndpointOptions = {}): string | undefined {
|
|
25
|
+
const url = opts.url ?? process.env.CODEINDEX_EMBED_ENDPOINT;
|
|
26
|
+
return url && url.trim() ? url.trim() : undefined;
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
// Request float embeddings for a batch of texts from the endpoint. Throws a
|
|
30
|
+
// clear error when no endpoint is configured or the response is malformed; the
|
|
31
|
+
// caller decides whether to degrade. NOTE: endpoint vectors are float and are
|
|
32
|
+
// NOT part of any determinism guarantee — this path is for the rich tier only.
|
|
33
|
+
export async function embedViaEndpoint(texts: string[], opts: EmbedEndpointOptions = {}): Promise<number[][]> {
|
|
34
|
+
const url = resolveEmbedEndpoint(opts);
|
|
35
|
+
if (!url) throw new Error("no embedding endpoint configured (set CODEINDEX_EMBED_ENDPOINT or pass opts.url)");
|
|
36
|
+
const controller = new AbortController();
|
|
37
|
+
const timer = setTimeout(() => controller.abort(), opts.timeoutMs ?? 30_000);
|
|
38
|
+
try {
|
|
39
|
+
const res = await fetch(url, {
|
|
40
|
+
method: "POST",
|
|
41
|
+
headers: { "content-type": "application/json", ...(opts.headers ?? {}) },
|
|
42
|
+
body: JSON.stringify({ texts }),
|
|
43
|
+
signal: controller.signal,
|
|
44
|
+
});
|
|
45
|
+
if (!res.ok) throw new Error(`embedding endpoint ${url} returned HTTP ${res.status}`);
|
|
46
|
+
const data = (await res.json()) as { vectors?: unknown };
|
|
47
|
+
const vectors = data.vectors;
|
|
48
|
+
if (!Array.isArray(vectors) || !vectors.every((v) => Array.isArray(v) && v.every((x) => typeof x === "number"))) {
|
|
49
|
+
throw new Error(`embedding endpoint ${url} returned a malformed { vectors } payload`);
|
|
50
|
+
}
|
|
51
|
+
return vectors as number[][];
|
|
52
|
+
} finally {
|
|
53
|
+
clearTimeout(timer);
|
|
54
|
+
}
|
|
55
|
+
}
|
|
@@ -0,0 +1,125 @@
|
|
|
1
|
+
import type { RepoScan } from "../scan.js";
|
|
2
|
+
import { encode } from "./encode.js";
|
|
3
|
+
import { EMBED_VERSION, type StaticEmbedModel } from "./model.js";
|
|
4
|
+
|
|
5
|
+
// The corpus embedding artifact. Granularity is PER-SYMBOL (each symbol gets its
|
|
6
|
+
// own record, so search can surface the exact symbol that matched), with a
|
|
7
|
+
// PER-FILE fallback record for files that declare no symbols (docs, config) so
|
|
8
|
+
// every file with content is still represented. Deterministic: records follow
|
|
9
|
+
// scan order (files sorted by rel) then declaration order within a file.
|
|
10
|
+
export interface EmbeddingRecord {
|
|
11
|
+
file: string; // repo-relative path
|
|
12
|
+
symbol?: string; // the declared symbol this vector represents (absent for a file-level record)
|
|
13
|
+
line?: number; // 1-based symbol line (absent for a file-level record)
|
|
14
|
+
vec: Int8Array; // length === dim
|
|
15
|
+
}
|
|
16
|
+
|
|
17
|
+
export interface EmbeddingIndex {
|
|
18
|
+
embedVersion: number;
|
|
19
|
+
modelId: string;
|
|
20
|
+
dim: number;
|
|
21
|
+
records: EmbeddingRecord[];
|
|
22
|
+
}
|
|
23
|
+
|
|
24
|
+
// The text encoded for one symbol: its name (camelCase-split by the tokenizer),
|
|
25
|
+
// its signature, the owning file's one-line summary, and the file's path segments
|
|
26
|
+
// — the same signal bm25 indexes, but pooled into a single vector. Newline-joined
|
|
27
|
+
// so the tokenizer's non-alphanumeric split cleanly separates the parts.
|
|
28
|
+
function symbolText(rel: string, name: string, signature: string | undefined, summary: string | undefined): string {
|
|
29
|
+
return [name, signature ?? "", summary ?? "", rel.replace(/\//g, " ")].join("\n");
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
// A file-level record's text (symbol-less files): title, summary, headings, path.
|
|
33
|
+
function fileText(rel: string, title: string | undefined, summary: string | undefined, headings: string[]): string {
|
|
34
|
+
return [title ?? "", summary ?? "", ...headings, rel.replace(/\//g, " ")].join("\n");
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
// Build the corpus embedding index from a scan + a loaded model. Pure and
|
|
38
|
+
// deterministic (encode is byte-stable, scan order is fixed) → two builds of an
|
|
39
|
+
// unchanged repo produce byte-identical serialized bytes.
|
|
40
|
+
export function buildEmbeddingIndex(scan: RepoScan, model: StaticEmbedModel): EmbeddingIndex {
|
|
41
|
+
const records: EmbeddingRecord[] = [];
|
|
42
|
+
for (const f of scan.files) {
|
|
43
|
+
const seen = new Set<string>();
|
|
44
|
+
let hadSymbol = false;
|
|
45
|
+
for (const s of f.symbols) {
|
|
46
|
+
if (seen.has(s.name)) continue; // dedupe by name within a file (bm25 parity)
|
|
47
|
+
seen.add(s.name);
|
|
48
|
+
hadSymbol = true;
|
|
49
|
+
records.push({
|
|
50
|
+
file: f.rel,
|
|
51
|
+
symbol: s.name,
|
|
52
|
+
line: s.line,
|
|
53
|
+
vec: encode(model, symbolText(f.rel, s.name, s.signature, f.summary)),
|
|
54
|
+
});
|
|
55
|
+
}
|
|
56
|
+
if (!hadSymbol) {
|
|
57
|
+
const text = fileText(f.rel, f.title, f.summary, f.headings);
|
|
58
|
+
if (text.replace(/\s+/g, "")) {
|
|
59
|
+
records.push({ file: f.rel, vec: encode(model, text) });
|
|
60
|
+
}
|
|
61
|
+
}
|
|
62
|
+
}
|
|
63
|
+
return { embedVersion: EMBED_VERSION, modelId: model.modelId, dim: model.dim, records };
|
|
64
|
+
}
|
|
65
|
+
|
|
66
|
+
const MAGIC = "CIE1"; // codeindex embeddings, format 1
|
|
67
|
+
|
|
68
|
+
// Serialize to embeddings.bin: a fixed ASCII magic, a uint32-LE header length, a
|
|
69
|
+
// UTF-8 JSON header (per-record metadata in build order — carries NO absolute
|
|
70
|
+
// path or timestamp), then the packed int8 body (count × dim signed bytes). The
|
|
71
|
+
// header JSON key order is fixed by construction, and every body byte is written
|
|
72
|
+
// explicitly, so the bytes are fully deterministic. Returns a Uint8Array (a
|
|
73
|
+
// Buffer at runtime) so the public type surface stays free of the Node `Buffer`
|
|
74
|
+
// global — a consumer type-checking without @types/node still resolves it.
|
|
75
|
+
export function serializeEmbeddings(index: EmbeddingIndex): Uint8Array {
|
|
76
|
+
const header = JSON.stringify({
|
|
77
|
+
embedVersion: index.embedVersion,
|
|
78
|
+
modelId: index.modelId,
|
|
79
|
+
dim: index.dim,
|
|
80
|
+
count: index.records.length,
|
|
81
|
+
records: index.records.map((r) => ({ file: r.file, symbol: r.symbol ?? "", line: r.line ?? 0 })),
|
|
82
|
+
});
|
|
83
|
+
const headerBuf = Buffer.from(header, "utf8");
|
|
84
|
+
const body = Buffer.alloc(index.records.length * index.dim);
|
|
85
|
+
let off = 0;
|
|
86
|
+
for (const r of index.records) {
|
|
87
|
+
for (let d = 0; d < index.dim; d++) body.writeInt8(r.vec[d] ?? 0, off++);
|
|
88
|
+
}
|
|
89
|
+
const out = Buffer.alloc(8 + headerBuf.length + body.length);
|
|
90
|
+
out.write(MAGIC, 0, "ascii");
|
|
91
|
+
out.writeUInt32LE(headerBuf.length, 4);
|
|
92
|
+
headerBuf.copy(out, 8);
|
|
93
|
+
body.copy(out, 8 + headerBuf.length);
|
|
94
|
+
return out;
|
|
95
|
+
}
|
|
96
|
+
|
|
97
|
+
// Inverse of serializeEmbeddings. Accepts any Uint8Array (e.g. a fs.readFileSync
|
|
98
|
+
// Buffer) and wraps it as a Buffer VIEW (no copy) for the numeric reads. Throws
|
|
99
|
+
// on a bad magic (a corrupt or foreign file) so a caller fails loudly rather
|
|
100
|
+
// than misreading arbitrary bytes.
|
|
101
|
+
export function deserializeEmbeddings(bytes: Uint8Array): EmbeddingIndex {
|
|
102
|
+
const buf = Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength);
|
|
103
|
+
if (buf.length < 8 || buf.toString("ascii", 0, 4) !== MAGIC) {
|
|
104
|
+
throw new Error("embeddings.bin: bad magic (not a codeindex embeddings artifact)");
|
|
105
|
+
}
|
|
106
|
+
const headerLen = buf.readUInt32LE(4);
|
|
107
|
+
const header = JSON.parse(buf.toString("utf8", 8, 8 + headerLen)) as {
|
|
108
|
+
embedVersion: number;
|
|
109
|
+
modelId: string;
|
|
110
|
+
dim: number;
|
|
111
|
+
count: number;
|
|
112
|
+
records: { file: string; symbol: string; line: number }[];
|
|
113
|
+
};
|
|
114
|
+
const bodyOff = 8 + headerLen;
|
|
115
|
+
const { dim } = header;
|
|
116
|
+
const records: EmbeddingRecord[] = header.records.map((m, i) => {
|
|
117
|
+
const vec = new Int8Array(dim);
|
|
118
|
+
for (let d = 0; d < dim; d++) vec[d] = buf.readInt8(bodyOff + i * dim + d);
|
|
119
|
+
const rec: EmbeddingRecord = { file: m.file, vec };
|
|
120
|
+
if (m.symbol) rec.symbol = m.symbol;
|
|
121
|
+
if (m.line) rec.line = m.line;
|
|
122
|
+
return rec;
|
|
123
|
+
});
|
|
124
|
+
return { embedVersion: header.embedVersion, modelId: header.modelId, dim, records };
|
|
125
|
+
}
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
import { existsSync, readFileSync } from "node:fs";
|
|
2
|
+
import { join } from "node:path";
|
|
3
|
+
|
|
4
|
+
// The deterministic static-embedding tier (v2.10.0). A *static* embedding model
|
|
5
|
+
// is a plain lookup table: token → row vector. Inference is tokenize → gather
|
|
6
|
+
// rows → mean-pool → L2-normalize → int8-quantize (see encode.ts). There is NO
|
|
7
|
+
// neural forward pass, no wasm, no float matmul — so the encode is pure-JS and
|
|
8
|
+
// byte-identical across platforms, which is the whole reason this tier exists.
|
|
9
|
+
//
|
|
10
|
+
// Activation is OPT-IN by asset presence, exactly like the tree-sitter grammar
|
|
11
|
+
// tier (src/ast/loader.ts): no model on disk → the engine silently stays
|
|
12
|
+
// lexical. Models are NEVER shipped in the npm tarball; the asset is fetched by
|
|
13
|
+
// `codeindex embed pull` into CODEINDEX_EMBED_DIR (or <repo>/.codeindex/models/).
|
|
14
|
+
|
|
15
|
+
// Bumped independently of SCHEMA_VERSION / EXTRACTOR_VERSION whenever the
|
|
16
|
+
// embeddings.bin artifact shape or the encode algorithm changes, so a consumer
|
|
17
|
+
// rejects an embeddings.bin written by an incompatible engine (or a different
|
|
18
|
+
// model) instead of misranking with it. SCHEMA_VERSION is deliberately left
|
|
19
|
+
// UNTOUCHED — embeddings are a purely additive sidecar with zero impact on the
|
|
20
|
+
// graph.json / symbols.json consumers.
|
|
21
|
+
export const EMBED_VERSION = 1;
|
|
22
|
+
|
|
23
|
+
// The default model-name a `pull` writes and a `status` reports. The real asset
|
|
24
|
+
// URL is intentionally NOT hard-coded yet (see resolveEmbedPullUrl): the asset
|
|
25
|
+
// is unpublished, so `pull` fails cleanly asking for CODEINDEX_EMBED_URL rather
|
|
26
|
+
// than fetching a phantom.
|
|
27
|
+
export const DEFAULT_EMBED_DIRNAME = "models";
|
|
28
|
+
|
|
29
|
+
// A loaded static-embedding model. `vocab` maps a wordpiece token to its row
|
|
30
|
+
// index; `weights` is the flat row-major (vocabSize × dim) matrix as IEEE-754
|
|
31
|
+
// doubles (float rows widen to double exactly, so mean-pool stays deterministic).
|
|
32
|
+
export interface StaticEmbedModel {
|
|
33
|
+
modelId: string;
|
|
34
|
+
dim: number;
|
|
35
|
+
unk: string; // the OOV token (e.g. "[UNK]"); may be absent from vocab
|
|
36
|
+
unkId: number; // row index of `unk`, or -1 when the model has no UNK row
|
|
37
|
+
vocabSize: number;
|
|
38
|
+
vocab: Map<string, number>;
|
|
39
|
+
weights: Float64Array; // length === vocabSize * dim, row-major
|
|
40
|
+
}
|
|
41
|
+
|
|
42
|
+
// The on-disk model.json header shape. `weights` is an array of `vocabSize`
|
|
43
|
+
// rows, each of length `dim`. Kept human-authorable (plain JSON numbers) so a
|
|
44
|
+
// tiny fixture model is committable and diffable; a real `pull`ed model uses the
|
|
45
|
+
// identical shape at scale.
|
|
46
|
+
interface ModelFile {
|
|
47
|
+
modelId: string;
|
|
48
|
+
dim: number;
|
|
49
|
+
unk?: string;
|
|
50
|
+
vocab: string[];
|
|
51
|
+
weights: number[][];
|
|
52
|
+
}
|
|
53
|
+
|
|
54
|
+
// Where the model asset lives. CODEINDEX_EMBED_DIR wins outright (explicit
|
|
55
|
+
// override, mirrors CODEINDEX_GRAMMAR_DIR); otherwise <repo>/.codeindex/models/
|
|
56
|
+
// then <cwd>/.codeindex/models/. Returns undefined when no model.json is found
|
|
57
|
+
// anywhere — the silent-degradation signal. When CODEINDEX_EMBED_DIR is set but
|
|
58
|
+
// empty of a model, that is still "no model" (undefined), never an error.
|
|
59
|
+
export function resolveEmbedModelDir(repo?: string): string | undefined {
|
|
60
|
+
const env = process.env.CODEINDEX_EMBED_DIR;
|
|
61
|
+
const candidates: string[] = [];
|
|
62
|
+
if (env) candidates.push(env);
|
|
63
|
+
if (repo) candidates.push(join(repo, ".codeindex", DEFAULT_EMBED_DIRNAME));
|
|
64
|
+
candidates.push(join(process.cwd(), ".codeindex", DEFAULT_EMBED_DIRNAME));
|
|
65
|
+
for (const c of candidates) {
|
|
66
|
+
if (existsSync(join(c, "model.json"))) return c;
|
|
67
|
+
}
|
|
68
|
+
return undefined;
|
|
69
|
+
}
|
|
70
|
+
|
|
71
|
+
// True when an activatable model is present (used by `embed status` and the MCP
|
|
72
|
+
// gate before attempting a semantic search).
|
|
73
|
+
export function hasEmbedModel(repo?: string): boolean {
|
|
74
|
+
return resolveEmbedModelDir(repo) !== undefined;
|
|
75
|
+
}
|
|
76
|
+
|
|
77
|
+
// Load and validate a static model from a directory containing model.json.
|
|
78
|
+
// Throws on a malformed file (bad dim, ragged weights) so a corrupt asset fails
|
|
79
|
+
// loudly at load rather than silently misranking. Returns undefined only when
|
|
80
|
+
// the directory has no model.json (the not-present case).
|
|
81
|
+
export function loadEmbedModel(dir?: string): StaticEmbedModel | undefined {
|
|
82
|
+
if (!dir) return undefined;
|
|
83
|
+
const path = join(dir, "model.json");
|
|
84
|
+
if (!existsSync(path)) return undefined;
|
|
85
|
+
const raw = JSON.parse(readFileSync(path, "utf8")) as ModelFile;
|
|
86
|
+
const { modelId, dim, vocab, weights } = raw;
|
|
87
|
+
if (typeof modelId !== "string" || !modelId) throw new Error(`embed model: missing modelId in ${path}`);
|
|
88
|
+
if (!Number.isInteger(dim) || dim <= 0) throw new Error(`embed model: bad dim ${dim} in ${path}`);
|
|
89
|
+
if (!Array.isArray(vocab) || !Array.isArray(weights) || vocab.length !== weights.length) {
|
|
90
|
+
throw new Error(`embed model: vocab/weights length mismatch in ${path}`);
|
|
91
|
+
}
|
|
92
|
+
const vocabSize = vocab.length;
|
|
93
|
+
const flat = new Float64Array(vocabSize * dim);
|
|
94
|
+
const vmap = new Map<string, number>();
|
|
95
|
+
for (let i = 0; i < vocabSize; i++) {
|
|
96
|
+
const tok = vocab[i]!;
|
|
97
|
+
if (typeof tok !== "string") throw new Error(`embed model: non-string vocab entry at ${i}`);
|
|
98
|
+
if (!vmap.has(tok)) vmap.set(tok, i); // first occurrence wins (deterministic)
|
|
99
|
+
const row = weights[i]!;
|
|
100
|
+
if (!Array.isArray(row) || row.length !== dim) {
|
|
101
|
+
throw new Error(`embed model: row ${i} has length ${row?.length}, expected ${dim}`);
|
|
102
|
+
}
|
|
103
|
+
for (let d = 0; d < dim; d++) flat[i * dim + d] = Number(row[d]);
|
|
104
|
+
}
|
|
105
|
+
const unk = typeof raw.unk === "string" ? raw.unk : "[UNK]";
|
|
106
|
+
const unkId = vmap.has(unk) ? vmap.get(unk)! : -1;
|
|
107
|
+
return { modelId, dim, unk, unkId, vocabSize, vocab: vmap, weights: flat };
|
|
108
|
+
}
|
|
109
|
+
|
|
110
|
+
// Resolve the URL `embed pull` fetches from. The official asset is not published
|
|
111
|
+
// yet, so there is no built-in default: the user must point CODEINDEX_EMBED_URL
|
|
112
|
+
// at a model.json (or a directory serving one). Returns undefined when unset —
|
|
113
|
+
// the caller turns that into a clean, actionable failure, never a crash.
|
|
114
|
+
export function resolveEmbedPullUrl(): string | undefined {
|
|
115
|
+
const url = process.env.CODEINDEX_EMBED_URL;
|
|
116
|
+
return url && url.trim() ? url.trim() : undefined;
|
|
117
|
+
}
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
import type { RepoScan } from "../scan.js";
|
|
2
|
+
import { rrf } from "../util.js";
|
|
3
|
+
import { byStr } from "../sort.js";
|
|
4
|
+
import { searchIndex, type SearchOptions, type SearchResult } from "../bm25.js";
|
|
5
|
+
import { encode, intDot } from "./encode.js";
|
|
6
|
+
import type { EmbeddingIndex } from "./index.js";
|
|
7
|
+
import type { StaticEmbedModel } from "./model.js";
|
|
8
|
+
|
|
9
|
+
const DEFAULT_LIMIT = 20;
|
|
10
|
+
const RRF_K = 60; // reuse the engine-wide RRF damping constant
|
|
11
|
+
|
|
12
|
+
export interface SemanticSearchOptions extends SearchOptions {
|
|
13
|
+
// The loaded static model, needed to encode the QUERY into the same int8 space
|
|
14
|
+
// as the corpus. Absent → the search degrades to pure lexical (no throw).
|
|
15
|
+
model?: StaticEmbedModel;
|
|
16
|
+
// RRF damping (default 60), exposed for parity with the shared rrf helper.
|
|
17
|
+
rrfK?: number;
|
|
18
|
+
}
|
|
19
|
+
|
|
20
|
+
// A fused result. Extends the lexical SearchResult additively: `semanticSymbol`
|
|
21
|
+
// is the corpus symbol whose embedding was closest to the query for this file
|
|
22
|
+
// (absent when the file was contributed only by the lexical side, or when the
|
|
23
|
+
// search degraded to lexical).
|
|
24
|
+
export interface SemanticSearchResult extends SearchResult {
|
|
25
|
+
semanticSymbol?: string;
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
// RRF-fused semantic + lexical search. The two rankings live on incomparable
|
|
29
|
+
// scales (BM25 score vs integer dot product), so we fuse by RANK via the shared
|
|
30
|
+
// `rrf` helper (k=60) rather than any linear score blend. Deterministic end to
|
|
31
|
+
// end: lexical is deterministic, the semantic ranking sorts by exact integer dot
|
|
32
|
+
// (ties broken by path), and the fused list sorts by RRF score (ties by path).
|
|
33
|
+
//
|
|
34
|
+
// DEGRADATION: with no model, no index, or an empty index, this returns the pure
|
|
35
|
+
// lexical ranking unchanged (as SemanticSearchResult[] with no `semanticSymbol`)
|
|
36
|
+
// — the caller stays on exit 0 and simply prints a note.
|
|
37
|
+
export function searchSemantic(
|
|
38
|
+
scan: RepoScan,
|
|
39
|
+
query: string,
|
|
40
|
+
index: EmbeddingIndex | undefined,
|
|
41
|
+
opts: SemanticSearchOptions = {},
|
|
42
|
+
): SemanticSearchResult[] {
|
|
43
|
+
const limit = opts.limit ?? DEFAULT_LIMIT;
|
|
44
|
+
// Pull a deeper lexical list than the final limit so RRF has enough overlap to
|
|
45
|
+
// fuse meaningfully, then trim after fusion.
|
|
46
|
+
const lexical = searchIndex(scan, query, { limit: Math.max(limit, 50), fuzzy: opts.fuzzy });
|
|
47
|
+
|
|
48
|
+
if (!opts.model || !index || index.records.length === 0) {
|
|
49
|
+
return lexical.slice(0, limit); // pure-lexical degradation
|
|
50
|
+
}
|
|
51
|
+
|
|
52
|
+
const q = encode(opts.model, query);
|
|
53
|
+
// Best (highest integer dot) record per file, remembering its symbol.
|
|
54
|
+
const bestByFile = new Map<string, { score: number; symbol?: string }>();
|
|
55
|
+
for (const r of index.records) {
|
|
56
|
+
const dot = intDot(q, r.vec);
|
|
57
|
+
const prev = bestByFile.get(r.file);
|
|
58
|
+
if (!prev || dot > prev.score) bestByFile.set(r.file, { score: dot, symbol: r.symbol });
|
|
59
|
+
}
|
|
60
|
+
|
|
61
|
+
// Semantic ranked file list: positive similarity only, best-first, ties by path.
|
|
62
|
+
const semList = [...bestByFile.entries()]
|
|
63
|
+
.filter(([, v]) => v.score > 0)
|
|
64
|
+
.sort((a, b) => b[1].score - a[1].score || byStr(a[0], b[0]))
|
|
65
|
+
.map(([file]) => file);
|
|
66
|
+
const lexList = lexical.map((r) => r.file);
|
|
67
|
+
|
|
68
|
+
// Fuse the two ranked lists (identity keyOf — items are already file paths).
|
|
69
|
+
const fused = rrf<string>([lexList, semList], (f) => f, opts.rrfK ?? RRF_K);
|
|
70
|
+
|
|
71
|
+
const lexByFile = new Map(lexical.map((r) => [r.file, r] as const));
|
|
72
|
+
const results: SemanticSearchResult[] = [...fused.entries()]
|
|
73
|
+
.sort((a, b) => b[1] - a[1] || byStr(a[0], b[0]))
|
|
74
|
+
.map(([file, score]) => {
|
|
75
|
+
const lex = lexByFile.get(file);
|
|
76
|
+
const res: SemanticSearchResult = {
|
|
77
|
+
file,
|
|
78
|
+
score: Number(score.toFixed(4)),
|
|
79
|
+
matchedTerms: lex?.matchedTerms ?? [],
|
|
80
|
+
topSymbols: lex?.topSymbols ?? [],
|
|
81
|
+
};
|
|
82
|
+
const sem = bestByFile.get(file);
|
|
83
|
+
if (sem?.symbol) res.semanticSymbol = sem.symbol;
|
|
84
|
+
if (lex?.fuzzyTerms) res.fuzzyTerms = lex.fuzzyTerms;
|
|
85
|
+
return res;
|
|
86
|
+
});
|
|
87
|
+
return results.slice(0, limit);
|
|
88
|
+
}
|
|
89
|
+
|
|
90
|
+
export { DEFAULT_LIMIT as SEMANTIC_DEFAULT_LIMIT };
|