@theihtisham/agent-shadow-brain 5.2.0 → 6.0.2

package/dist/brain/brain-quantize.js

@@ -0,0 +1,346 @@
+// src/brain/brain-quantize.ts — int8 vector quantization for the brain
+// v6.0.2 — Hive Mind Edition
+//
+// Float32 vectors are 4x bigger on disk and in memory than they need to be for
+// approximate-nearest-neighbor recall. Symmetric int8 quantization preserves
+// recall to within ~1% on unit-norm vectors while shrinking the working set
+// by 4x. That means the same RAM budget covers four times as many memories
+// and the cold-start read of a sharded JSON index is dramatically faster.
+//
+// Design:
+//   - Symmetric per-vector int8: scale = max(|v_i|) / 127, q_i = round(v_i/scale).
+//     Per-vector scale (not per-tensor) handles the case where some memories
+//     have a much wider dynamic range than others (e.g. Ollama dense vs hash).
+//   - Compressed format ".q8" is a tight binary:
+//       magic     : 4 bytes  ascii  "Q8I0"
+//       version   : uint32   1
+//       dim       : uint32   embedding dimension
+//       count     : uint32   number of entries
+//       entries[] : { float64 scale; int8[dim] q; uint16 idLen; utf8 id[idLen] }
+//   - Direct cosine on int8: integer dot then divide by scaleA*scaleB and
+//     (dim * 127^2). Skips a dequantize-into-float32 round trip in the hot path.
+//
+// Exposed: BrainQuantize, getBrainQuantize(), resetBrainQuantizeForTests().
+//   .quantize(vec)
+//   .dequantize(q, scale)
+//   .cosineQuantized(a, scaleA, b, scaleB)
+//   .compressIndex(srcPath, outPath)
+//   .loadCompressedIndex(path)
+//   .quantizeRoundtripError(vectors)
+import * as fs from 'fs';
+import * as path from 'path';
+const MAGIC = Buffer.from('Q8I0', 'ascii');
+const VERSION = 1;
+const HEADER_BYTES = 4 + 4 + 4 + 4; // magic + version + dim + count
+export class BrainQuantize {
+    /** Symmetric int8 quantization with per-vector scale. */
+    quantize(vector) {
+        if (!vector || vector.length === 0)
+            return new Int8Array(0);
+        let maxAbs = 0;
+        for (let i = 0; i < vector.length; i++) {
+            const a = Math.abs(vector[i]);
+            if (a > maxAbs)
+                maxAbs = a;
+        }
+        const out = new Int8Array(vector.length);
+        if (maxAbs === 0)
+            return out; // all zeros — quantizes to all zeros
+        const scale = maxAbs / 127;
+        const inv = 1 / scale;
+        for (let i = 0; i < vector.length; i++) {
+            const q = Math.round(vector[i] * inv);
+            out[i] = q > 127 ? 127 : q < -127 ? -127 : q;
+        }
+        return out;
+    }
+    /** Recover an approximation of the original float vector. */
+    dequantize(q, scale) {
+        if (!q || q.length === 0)
+            return [];
+        const out = new Array(q.length);
+        for (let i = 0; i < q.length; i++)
+            out[i] = q[i] * scale;
+        return out;
+    }
+    /**
+     * Cosine similarity on quantized vectors WITHOUT dequantizing first.
+     *
+     * If `a` and `b` were derived from unit-norm float vectors, the cosine is
+     * just the dot product after rescaling. We compute the integer dot then
+     * apply `scaleA * scaleB` to convert back to the float dot, and clamp to
+     * [-1, 1] for numeric safety. Note: the dim/127^2 factor in the spec is the
+     * theoretical worst-case scale (when every element is ±127); the actual
+     * normalization is captured entirely by scaleA*scaleB, so we use that and
+     * do not double-divide. The clamp guards against quantization rounding
+     * pushing the result slightly out of range.
+     */
+    cosineQuantized(a, scaleA, b, scaleB) {
+        if (!a || !b || a.length === 0 || b.length === 0)
+            return 0;
+        const n = Math.min(a.length, b.length);
+        let dot = 0;
+        for (let i = 0; i < n; i++)
+            dot += a[i] * b[i];
+        const sim = dot * scaleA * scaleB;
+        if (sim > 1)
+            return 1;
+        if (sim < -1)
+            return -1;
+        return sim;
+    }
+    /** Convert a v2 sharded JSON index into a single compressed .q8 binary. */
+    async compressIndex(srcPath, outPath) {
+        const start = Date.now();
+        const docs = this.readShardDocs(srcPath);
+        if (docs.length === 0) {
+            // Empty index → still write a valid empty file.
+            const header = Buffer.alloc(HEADER_BYTES);
+            MAGIC.copy(header, 0);
+            header.writeUInt32LE(VERSION, 4);
+            header.writeUInt32LE(0, 8);
+            header.writeUInt32LE(0, 12);
+            this.atomicWrite(outPath, header);
+            return {
+                originalBytes: this.measureBytes(srcPath),
+                compressedBytes: header.length,
+                compressionRatio: 0,
+                vectorCount: 0,
+                dim: 0,
+                encodingError: { meanAbs: 0, max: 0 },
+                durationMs: Date.now() - start,
+            };
+        }
+        const dim = docs[0].vector.length;
+        const quantized = [];
+        for (const d of docs) {
+            if (!d.vector || d.vector.length !== dim)
+                continue; // skip ragged
+            const q = this.quantize(d.vector);
+            const scale = this.scaleOf(d.vector);
+            quantized.push({ id: d.id, q, scale, original: d.vector });
+        }
+        // Estimate output size up front so we allocate one Buffer.
+        let bodyBytes = 0;
+        const idBufs = [];
+        for (const e of quantized) {
+            const idBuf = Buffer.from(e.id, 'utf-8');
+            idBufs.push(idBuf);
+            bodyBytes += 8 + dim + 2 + idBuf.length;
+        }
+        const totalBytes = HEADER_BYTES + bodyBytes;
+        const buf = Buffer.alloc(totalBytes);
+        MAGIC.copy(buf, 0);
+        buf.writeUInt32LE(VERSION, 4);
+        buf.writeUInt32LE(dim, 8);
+        buf.writeUInt32LE(quantized.length, 12);
+        let offset = HEADER_BYTES;
+        for (let i = 0; i < quantized.length; i++) {
+            const e = quantized[i];
+            buf.writeDoubleLE(e.scale, offset);
+            offset += 8;
+            // int8 array copy — Buffer views the Int8Array directly.
+            Buffer.from(e.q.buffer, e.q.byteOffset, e.q.byteLength).copy(buf, offset);
+            offset += dim;
+            const idBuf = idBufs[i];
+            buf.writeUInt16LE(idBuf.length, offset);
+            offset += 2;
+            idBuf.copy(buf, offset);
+            offset += idBuf.length;
+        }
+        this.atomicWrite(outPath, buf);
+        // Encoding error sample (cap at 500 vectors so this stays cheap).
+        const sampleN = Math.min(500, quantized.length);
+        let sumAbs = 0, maxAbs = 0, count = 0;
+        for (let i = 0; i < sampleN; i++) {
+            const e = quantized[i];
+            const recon = this.dequantize(e.q, e.scale);
+            for (let j = 0; j < dim; j++) {
+                const err = Math.abs(recon[j] - e.original[j]);
+                sumAbs += err;
+                if (err > maxAbs)
+                    maxAbs = err;
+                count++;
+            }
+        }
+        const originalBytes = this.measureBytes(srcPath);
+        return {
+            originalBytes,
+            compressedBytes: totalBytes,
+            compressionRatio: originalBytes > 0 ? +(originalBytes / totalBytes).toFixed(3) : 0,
+            vectorCount: quantized.length,
+            dim,
+            encodingError: {
+                meanAbs: count > 0 ? +(sumAbs / count).toFixed(6) : 0,
+                max: +maxAbs.toFixed(6),
+            },
+            durationMs: Date.now() - start,
+        };
+    }
+    /** Load a .q8 file and return vectors plus a ready-to-use recall function. */
+    async loadCompressedIndex(filePath) {
+        const buf = await fs.promises.readFile(filePath);
+        if (buf.length < HEADER_BYTES)
+            throw new Error(`q8 index too short: ${buf.length} bytes`);
+        if (buf.subarray(0, 4).compare(MAGIC) !== 0)
+            throw new Error('q8 index: bad magic');
+        const version = buf.readUInt32LE(4);
+        if (version !== VERSION)
+            throw new Error(`q8 index: unsupported version ${version}`);
+        const dim = buf.readUInt32LE(8);
+        const count = buf.readUInt32LE(12);
+        const vectors = new Map();
+        let offset = HEADER_BYTES;
+        for (let i = 0; i < count; i++) {
+            if (offset + 8 + dim + 2 > buf.length)
+                throw new Error(`q8 index truncated at entry ${i}`);
+            const scale = buf.readDoubleLE(offset);
+            offset += 8;
+            // Copy out of the parent buffer so the underlying ArrayBuffer can be GC'd
+            // and the Int8Array view stays tight to its dim bytes.
+            const q = new Int8Array(dim);
+            buf.copy(Buffer.from(q.buffer, q.byteOffset, q.byteLength), 0, offset, offset + dim);
+            offset += dim;
+            const idLen = buf.readUInt16LE(offset);
+            offset += 2;
+            if (offset + idLen > buf.length)
+                throw new Error(`q8 index truncated id at entry ${i}`);
+            const id = buf.subarray(offset, offset + idLen).toString('utf-8');
+            offset += idLen;
+            vectors.set(id, { q, scale });
+        }
+        const recallFn = (query, topK) => {
+            if (vectors.size === 0 || !query || query.length === 0 || topK <= 0)
+                return [];
+            const qVec = query.length === dim ? query : padOrTruncate(query, dim);
+            const qInt = this.quantize(qVec);
+            const qScale = this.scaleOf(qVec);
+            const hits = [];
+            for (const [id, entry] of vectors) {
+                const score = this.cosineQuantized(qInt, qScale, entry.q, entry.scale);
+                hits.push({ id, score });
+            }
+            hits.sort((a, b) => b.score - a.score);
+            return hits.slice(0, topK);
+        };
+        return { vectors, dim, recallFn };
+    }
+    /** Self-test: report mean / max absolute reconstruction error on `vectors`. */
+    quantizeRoundtripError(vectors) {
+        if (!vectors || vectors.length === 0)
+            return { meanAbs: 0, max: 0, samples: 0 };
+        let sum = 0, maxAbs = 0, n = 0;
+        for (const v of vectors) {
+            if (!v || v.length === 0)
+                continue;
+            const q = this.quantize(v);
+            const scale = this.scaleOf(v);
+            for (let i = 0; i < v.length; i++) {
+                const recon = q[i] * scale;
+                const err = Math.abs(recon - v[i]);
+                sum += err;
+                if (err > maxAbs)
+                    maxAbs = err;
+                n++;
+            }
+        }
+        return {
+            meanAbs: n > 0 ? +(sum / n).toFixed(6) : 0,
+            max: +maxAbs.toFixed(6),
+            samples: n,
+        };
+    }
+    // ── Internals ──────────────────────────────────────────────────────────
+    scaleOf(vector) {
+        let maxAbs = 0;
+        for (let i = 0; i < vector.length; i++) {
+            const a = Math.abs(vector[i]);
+            if (a > maxAbs)
+                maxAbs = a;
+        }
+        return maxAbs === 0 ? 0 : maxAbs / 127;
+    }
+    readShardDocs(srcPath) {
+        // Accept either a single JSON file with {docs: [...]} or a directory of shards.
+        const docs = [];
+        try {
+            const stat = fs.statSync(srcPath);
+            if (stat.isDirectory()) {
+                for (const f of fs.readdirSync(srcPath)) {
+                    if (!f.endsWith('.json'))
+                        continue;
+                    const raw = fs.readFileSync(path.join(srcPath, f), 'utf-8');
+                    this.collectDocs(raw, docs);
+                }
+            }
+            else {
+                const raw = fs.readFileSync(srcPath, 'utf-8');
+                this.collectDocs(raw, docs);
+            }
+        }
+        catch { /* missing or unreadable → empty result */ }
+        return docs;
+    }
+    collectDocs(raw, out) {
+        try {
+            const parsed = JSON.parse(raw);
+            if (Array.isArray(parsed.docs)) {
+                for (const d of parsed.docs) {
+                    if (d && typeof d.id === 'string' && Array.isArray(d.vector))
+                        out.push({ id: d.id, vector: d.vector });
+                }
+            }
+            else if (parsed.vectors && typeof parsed.vectors === 'object') {
+                // Legacy v1 cache shape: { vectors: { id: number[] } }
+                for (const [id, vec] of Object.entries(parsed.vectors)) {
+                    if (Array.isArray(vec))
+                        out.push({ id, vector: vec });
+                }
+            }
+        }
+        catch { /* skip corrupt shard */ }
+    }
+    measureBytes(srcPath) {
+        try {
+            const stat = fs.statSync(srcPath);
+            if (stat.isFile())
+                return stat.size;
+            if (stat.isDirectory()) {
+                let total = 0;
+                for (const f of fs.readdirSync(srcPath)) {
+                    try {
+                        total += fs.statSync(path.join(srcPath, f)).size;
+                    }
+                    catch { /* skip */ }
+                }
+                return total;
+            }
+        }
+        catch { /* missing */ }
+        return 0;
+    }
+    atomicWrite(filePath, data) {
+        fs.mkdirSync(path.dirname(filePath), { recursive: true });
+        const tmp = filePath + '.tmp';
+        fs.writeFileSync(tmp, data);
+        fs.renameSync(tmp, filePath);
+    }
+}
+function padOrTruncate(v, dim) {
+    if (v.length === dim)
+        return v;
+    if (v.length > dim)
+        return v.slice(0, dim);
+    const out = v.slice();
+    while (out.length < dim)
+        out.push(0);
+    return out;
+}
+let _instance = null;
+export function getBrainQuantize() {
+    if (!_instance)
+        _instance = new BrainQuantize();
+    return _instance;
+}
+export function resetBrainQuantizeForTests() { _instance = null; }
+//# sourceMappingURL=brain-quantize.js.map

package/dist/brain/brain-quantize.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"brain-quantize.js","sourceRoot":"","sources":["../../src/brain/brain-quantize.ts"],"names":[],"mappings":"AAAA,uEAAuE;AACvE,6BAA6B;AAC7B,EAAE;AACF,+EAA+E;AAC/E,6EAA6E;AAC7E,4EAA4E;AAC5E,2EAA2E;AAC3E,0EAA0E;AAC1E,EAAE;AACF,UAAU;AACV,mFAAmF;AACnF,6EAA6E;AAC7E,+EAA+E;AAC/E,iDAAiD;AACjD,2CAA2C;AAC3C,+BAA+B;AAC/B,iDAAiD;AACjD,+CAA+C;AAC/C,iFAAiF;AACjF,0EAA0E;AAC1E,iFAAiF;AACjF,EAAE;AACF,4EAA4E;AAC5E,mBAAmB;AACnB,0BAA0B;AAC1B,2CAA2C;AAC3C,qCAAqC;AACrC,+BAA+B;AAC/B,qCAAqC;AAErC,OAAO,KAAK,EAAE,MAAM,IAAI,CAAC;AACzB,OAAO,KAAK,IAAI,MAAM,MAAM,CAAC;AAE7B,MAAM,KAAK,GAAG,MAAM,CAAC,IAAI,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;AAC3C,MAAM,OAAO,GAAG,CAAC,CAAC;AAClB,MAAM,YAAY,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC,gCAAgC;AA6BpE,MAAM,OAAO,aAAa;IACxB,yDAAyD;IACzD,QAAQ,CAAC,MAAgB;QACvB,IAAI,CAAC,MAAM,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC;YAAE,OAAO,IAAI,SAAS,CAAC,CAAC,CAAC,CAAC;QAC5D,IAAI,MAAM,GAAG,CAAC,CAAC;QACf,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACvC,MAAM,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;YAC9B,IAAI,CAAC,GAAG,MAAM;gBAAE,MAAM,GAAG,CAAC,CAAC;QAC7B,CAAC;QACD,MAAM,GAAG,GAAG,IAAI,SAAS,CAAC,MAAM,CAAC,MAAM,CAAC,CAAC;QACzC,IAAI,MAAM,KAAK,CAAC;YAAE,OAAO,GAAG,CAAC,CAAC,qCAAqC;QACnE,MAAM,KAAK,GAAG,MAAM,GAAG,GAAG,CAAC;QAC3B,MAAM,GAAG,GAAG,CAAC,GAAG,KAAK,CAAC;QACtB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACvC,MAAM,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,MAAM,CAAC,CAAC,CAAC,GAAG,GAAG,CAAC,CAAC;YACtC,GAAG,CAAC,CAAC,CAAC,GAAG,CAAC,GAAG,GAAG,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC;QAC/C,CAAC;QACD,OAAO,GAAG,CAAC;IACb,CAAC;IAED,6DAA6D;IAC7D,UAAU,CAAC,CAAY,EAAE,KAAa;QACpC,IAAI,CAAC,CAAC,IAAI,CAAC,CAAC,MAAM,KAAK,CAAC;YAAE,OAAO,EAAE,CAAC;QACpC,MAAM,GAAG,GAAG,IAAI,KAAK,CAAS,CAAC,CAAC,MAAM,CAAC,CAAC;QACxC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,MAAM,EAAE,CAAC,EAAE;YAAE,GAAG,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,KAAK,CAAC;QACzD,OAAO,GAAG,CAAC;IACb,CAAC;IAED;;;;;;;;;;;OAWG;IACH,eAAe,CAAC,CAAY,EAAE,MAAc,EAAE,CAAY,EAAE,MAAc;QACxE,IAAI,CAAC,CAAC,IAAI,CAAC,CAAC,IAAI,CAAC,CAAC,MAAM,KAAK,CAAC,IAAI,CAAC,CAAC,MAAM,KAAK,CAAC;YAAE,OAAO,CAAC,CAAC;QAC3D,MAAM,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC,MAAM,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC;QACvC,IAAI,GAAG,GAAG,CAAC,CAAC;QACZ,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE;YAAE,GAAG,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC;QAC/C,MAAM,GAAG,GAAG,GAAG,GAAG,MAAM,GAAG,MAAM,CAAC;QAClC,IAAI,GAAG,GAAG,CAAC;YAAE,OAAO,CAAC,CAAC;QACtB,IAAI,GAAG,GAAG,CAAC,CAAC;YAAE,OAAO,CAAC,CAAC,CAAC;QACxB,OAAO,GAAG,CAAC;IACb,CAAC;IAED,2EAA2E;IAC3E,KAAK,CAAC,aAAa,CAAC,OAAe,EAAE,OAAe;QAClD,MAAM,KAAK,GAAG,IAAI,CAAC,GAAG,EAAE,CAAC;QACzB,MAAM,IAAI,GAAG,IAAI,CAAC,aAAa,CAAC,OAAO,CAAC,CAAC;QACzC,IAAI,IAAI,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YACtB,gDAAgD;YAChD,MAAM,MAAM,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,CAAC;YAC1C,KAAK,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC;YACtB,MAAM,CAAC,aAAa,CAAC,OAAO,EAAE,CAAC,CAAC,CAAC;YACjC,MAAM,CAAC,aAAa,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;YAC3B,MAAM,CAAC,aAAa,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC;YAC5B,IAAI,CAAC,WAAW,CAAC,OAAO,EAAE,MAAM,CAAC,CAAC;YAClC,OAAO;gBACL,aAAa,EAAE,IAAI,CAAC,YAAY,CAAC,OAAO,CAAC;gBACzC,eAAe,EAAE,MAAM,CAAC,MAAM;gBAC9B,gBAAgB,EAAE,CAAC;gBACnB,WAAW,EAAE,CAAC;gBACd,GAAG,EAAE,CAAC;gBACN,aAAa,EAAE,EAAE,OAAO,EAAE,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE;gBACrC,UAAU,EAAE,IAAI,CAAC,GAAG,EAAE,GAAG,KAAK;aAC/B,CAAC;QACJ,CAAC;QAED,MAAM,GAAG,GAAG,IAAI,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC,MAAM,CAAC;QAClC,MAAM,SAAS,GAA2E,EAAE,CAAC;QAC7F,KAAK,MAAM,CAAC,IAAI,IAAI,EAAE,CAAC;YACrB,IAAI,CAAC,CAAC,CAAC,MAAM,IAAI,CAAC,CAAC,MAAM,CAAC,MAAM,KAAK,GAAG;gBAAE,SAAS,CAAC,cAAc;YAClE,MAAM,CAAC,GAAG,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC;YAClC,MAAM,KAAK,GAAG,IAAI,CAAC,OAAO,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC;YACrC,SAAS,CAAC,IAAI,CAAC,EAAE,EAAE,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAC,MAAM,EAAE,CAAC,CAAC;QAC7D,CAAC;QAED,2DAA2D;QAC3D,IAAI,SAAS,GAAG,CAAC,CAAC;QAClB,MAAM,MAAM,GAAa,EAAE,CAAC;QAC5B,KAAK,MAAM,CAAC,IAAI,SAAS,EAAE,CAAC;YAC1B,MAAM,KAAK,GAAG,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,EAAE,OAAO,CAAC,CAAC;YACzC,MAAM,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;YACnB,SAAS,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC,GAAG,KAAK,CAAC,MAAM,CAAC;QAC1C,CAAC;QACD,MAAM,UAAU,GAAG,YAAY,GAAG,SAAS,CAAC;QAC5C,MAAM,GAAG,GAAG,MAAM,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;QAErC,KAAK,CAAC,IAAI,CAAC,GAAG,EAAE,CAAC,CAAC,CAAC;QACnB,GAAG,CAAC,aAAa,CAAC,OAAO,EAAE,CAAC,CAAC,CAAC;QAC9B,GAAG,CAAC,aAAa,CAAC,GAAG,EAAE,CAAC,CAAC,CAAC;QAC1B,GAAG,CAAC,aAAa,CAAC,SAAS,CAAC,MAAM,EAAE,EAAE,CAAC,CAAC;QAExC,IAAI,MAAM,GAAG,YAAY,CAAC;QAC1B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,SAAS,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YAC1C,MAAM,CAAC,GAAG,SAAS,CAAC,CAAC,CAAC,CAAC;YACvB,GAAG,CAAC,aAAa,CAAC,CAAC,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC;YAAC,MAAM,IAAI,CAAC,CAAC;YAChD,yDAAyD;YACzD,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,CAAC,UAAU,EAAE,CAAC,CAAC,CAAC,CAAC,UAAU,CAAC,CAAC,IAAI,CAAC,GAAG,EAAE,MAAM,CAAC,CAAC;YAC1E,MAAM,IAAI,GAAG,CAAC;YACd,MAAM,KAAK,GAAG,MAAM,CAAC,CAAC,CAAC,CAAC;YACxB,GAAG,CAAC,aAAa,CAAC,KAAK,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;YAAC,MAAM,IAAI,CAAC,CAAC;YACrD,KAAK,CAAC,IAAI,CAAC,GAAG,EAAE,MAAM,CAAC,CAAC;YAAC,MAAM,IAAI,KAAK,CAAC,MAAM,CAAC;QAClD,CAAC;QAED,IAAI,CAAC,WAAW,CAAC,OAAO,EAAE,GAAG,CAAC,CAAC;QAE/B,kEAAkE;QAClE,MAAM,OAAO,GAAG,IAAI,CAAC,GAAG,CAAC,GAAG,EAAE,SAAS,CAAC,MAAM,CAAC,CAAC;QAChD,IAAI,MAAM,GAAG,CAAC,EAAE,MAAM,GAAG,CAAC,EAAE,KAAK,GAAG,CAAC,CAAC;QACtC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,EAAE,CAAC,EAAE,EAAE,CAAC;YACjC,MAAM,CAAC,GAAG,SAAS,CAAC,CAAC,CAAC,CAAC;YACvB,MAAM,KAAK,GAAG,IAAI,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,KAAK,CAAC,CAAC;YAC5C,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,GAAG,EAAE,CAAC,EAAE,EAAE,CAAC;gBAC7B,MAAM,GAAG,GAAG,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,QAAQ,CAAC,CAAC,CAAC,CAAC,CAAC;gBAC/C,MAAM,IAAI,GAAG,CAAC;gBACd,IAAI,GAAG,GAAG,MAAM;oBAAE,MAAM,GAAG,GAAG,CAAC;gBAC/B,KAAK,EAAE,CAAC;YACV,CAAC;QACH,CAAC;QAED,MAAM,aAAa,GAAG,IAAI,CAAC,YAAY,CAAC,OAAO,CAAC,CAAC;QACjD,OAAO;YACL,aAAa;YACb,eAAe,EAAE,UAAU;YAC3B,gBAAgB,EAAE,aAAa,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,aAAa,GAAG,UAAU,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YAClF,WAAW,EAAE,SAAS,CAAC,MAAM;YAC7B,GAAG;YACH,aAAa,EAAE;gBACb,OAAO,EAAE,KAAK,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,MAAM,GAAG,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;gBACrD,GAAG,EAAE,CAAC,MAAM,CAAC,OAAO,CAAC,CAAC,CAAC;aACxB;YACD,UAAU,EAAE,IAAI,CAAC,GAAG,EAAE,GAAG,KAAK;SAC/B,CAAC;IACJ,CAAC;IAED,8EAA8E;IAC9E,KAAK,CAAC,mBAAmB,CAAC,QAAgB;QACxC,MAAM,GAAG,GAAG,MAAM,EAAE,CAAC,QAAQ,CAAC,QAAQ,CAAC,QAAQ,CAAC,CAAC;QACjD,IAAI,GAAG,CAAC,MAAM,GAAG,YAAY;YAAE,MAAM,IAAI,KAAK,CAAC,uBAAuB,GAAG,CAAC,MAAM,QAAQ,CAAC,CAAC;QAC1F,IAAI,GAAG,CAAC,QAAQ,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,CAAC,KAAK,CAAC,KAAK,CAAC;YAAE,MAAM,IAAI,KAAK,CAAC,qBAAqB,CAAC,CAAC;QACpF,MAAM,OAAO,GAAG,GAAG,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;QACpC,IAAI,OAAO,KAAK,OAAO;YAAE,MAAM,IAAI,KAAK,CAAC,iCAAiC,OAAO,EAAE,CAAC,CAAC;QACrF,MAAM,GAAG,GAAG,GAAG,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;QAChC,MAAM,KAAK,GAAG,GAAG,CAAC,YAAY,CAAC,EAAE,CAAC,CAAC;QAEnC,MAAM,OAAO,GAAG,IAAI,GAAG,EAA0B,CAAC;QAClD,IAAI,MAAM,GAAG,YAAY,CAAC;QAC1B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,KAAK,EAAE,CAAC,EAAE,EAAE,CAAC;YAC/B,IAAI,MAAM,GAAG,CAAC,GAAG,GAAG,GAAG,CAAC,GAAG,GAAG,CAAC,MAAM;gBAAE,MAAM,IAAI,KAAK,CAAC,+BAA+B,CAAC,EAAE,CAAC,CAAC;YAC3F,MAAM,KAAK,GAAG,GAAG,CAAC,YAAY,CAAC,MAAM,CAAC,CAAC;YAAC,MAAM,IAAI,CAAC,CAAC;YACpD,0EAA0E;YAC1E,uDAAuD;YACvD,MAAM,CAAC,GAAG,IAAI,SAAS,CAAC,GAAG,CAAC,CAAC;YAC7B,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC,MAAM,EAAE,CAAC,CAAC,UAAU,EAAE,CAAC,CAAC,UAAU,CAAC,EAAE,CAAC,EAAE,MAAM,EAAE,MAAM,GAAG,GAAG,CAAC,CAAC;YACrF,MAAM,IAAI,GAAG,CAAC;YACd,MAAM,KAAK,GAAG,GAAG,CAAC,YAAY,CAAC,MAAM,CAAC,CAAC;YAAC,MAAM,IAAI,CAAC,CAAC;YACpD,IAAI,MAAM,GAAG,KAAK,GAAG,GAAG,CAAC,MAAM;gBAAE,MAAM,IAAI,KAAK,CAAC,kCAAkC,CAAC,EAAE,CAAC,CAAC;YACxF,MAAM,EAAE,GAAG,GAAG,CAAC,QAAQ,CAAC,MAAM,EAAE,MAAM,GAAG,KAAK,CAAC,CAAC,QAAQ,CAAC,OAAO,CAAC,CAAC;YAClE,MAAM,IAAI,KAAK,CAAC;YAChB,OAAO,CAAC,GAAG,CAAC,EAAE,EAAE,EAAE,CAAC,EAAE,KAAK,EAAE,CAAC,CAAC;QAChC,CAAC;QAED,MAAM,QAAQ,GAAG,CAAC,KAAe,EAAE,IAAY,EAAwC,EAAE;YACvF,IAAI,OAAO,CAAC,IAAI,KAAK,CAAC,IAAI,CAAC,KAAK,IAAI,KAAK,CAAC,MAAM,KAAK,CAAC,IAAI,IAAI,IAAI,CAAC;gBAAE,OAAO,EAAE,CAAC;YAC/E,MAAM,IAAI,GAAG,KAAK,CAAC,MAAM,KAAK,GAAG,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,aAAa,CAAC,KAAK,EAAE,GAAG,CAAC,CAAC;YACtE,MAAM,IAAI,GAAG,IAAI,CAAC,QAAQ,CAAC,IAAI,CAAC,CAAC;YACjC,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC,CAAC;YAClC,MAAM,IAAI,GAAyC,EAAE,CAAC;YACtD,KAAK,MAAM,CAAC,EAAE,EAAE,KAAK,CAAC,IAAI,OAAO,EAAE,CAAC;gBAClC,MAAM,KAAK,GAAG,IAAI,CAAC,eAAe,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,CAAC,CAAC,EAAE,KAAK,CAAC,KAAK,CAAC,CAAC;gBACvE,IAAI,CAAC,IAAI,CAAC,EAAE,EAAE,EAAE,KAAK,EAAE,CAAC,CAAC;YAC3B,CAAC;YACD,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,KAAK,GAAG,CAAC,CAAC,KAAK,CAAC,CAAC;YACvC,OAAO,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,CAAC;QAC7B,CAAC,CAAC;QAEF,OAAO,EAAE,OAAO,EAAE,GAAG,EAAE,QAAQ,EAAE,CAAC;IACpC,CAAC;IAED,+EAA+E;IAC/E,sBAAsB,CAAC,OAAmB;QACxC,IAAI,CAAC,OAAO,IAAI,OAAO,CAAC,MAAM,KAAK,CAAC;YAAE,OAAO,EAAE,OAAO,EAAE,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,CAAC;QAChF,IAAI,GAAG,GAAG,CAAC,EAAE,MAAM,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC;QAC/B,KAAK,MAAM,CAAC,IAAI,OAAO,EAAE,CAAC;YACxB,IAAI,CAAC,CAAC,IAAI,CAAC,CAAC,MAAM,KAAK,CAAC;gBAAE,SAAS;YACnC,MAAM,CAAC,GAAG,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC,CAAC;YAC3B,MAAM,KAAK,GAAG,IAAI,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC;YAC9B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;gBAClC,MAAM,KAAK,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,KAAK,CAAC;gBAC3B,MAAM,GAAG,GAAG,IAAI,CAAC,GAAG,CAAC,KAAK,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;gBACnC,GAAG,IAAI,GAAG,CAAC;gBACX,IAAI,GAAG,GAAG,MAAM;oBAAE,MAAM,GAAG,GAAG,CAAC;gBAC/B,CAAC,EAAE,CAAC;YACN,CAAC;QACH,CAAC;QACD,OAAO;YACL,OAAO,EAAE,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,GAAG,GAAG,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YAC1C,GAAG,EAAE,CAAC,MAAM,CAAC,OAAO,CAAC,CAAC,CAAC;YACvB,OAAO,EAAE,CAAC;SACX,CAAC;IACJ,CAAC;IAED,0EAA0E;IAElE,OAAO,CAAC,MAAgB;QAC9B,IAAI,MAAM,GAAG,CAAC,CAAC;QACf,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACvC,MAAM,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;YAC9B,IAAI,CAAC,GAAG,MAAM;gBAAE,MAAM,GAAG,CAAC,CAAC;QAC7B,CAAC;QACD,OAAO,MAAM,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,MAAM,GAAG,GAAG,CAAC;IACzC,CAAC;IAEO,aAAa,CAAC,OAAe;QACnC,gFAAgF;QAChF,MAAM,IAAI,GAAmB,EAAE,CAAC;QAChC,IAAI,CAAC;YACH,MAAM,IAAI,GAAG,EAAE,CAAC,QAAQ,CAAC,OAAO,CAAC,CAAC;YAClC,IAAI,IAAI,CAAC,WAAW,EAAE,EAAE,CAAC;gBACvB,KAAK,MAAM,CAAC,IAAI,EAAE,CAAC,WAAW,CAAC,OAAO,CAAC,EAAE,CAAC;oBACxC,IAAI,CAAC,CAAC,CAAC,QAAQ,CAAC,OAAO,CAAC;wBAAE,SAAS;oBACnC,MAAM,GAAG,GAAG,EAAE,CAAC,YAAY,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,EAAE,CAAC,CAAC,EAAE,OAAO,CAAC,CAAC;oBAC5D,IAAI,CAAC,WAAW,CAAC,GAAG,EAAE,IAAI,CAAC,CAAC;gBAC9B,CAAC;YACH,CAAC;iBAAM,CAAC;gBACN,MAAM,GAAG,GAAG,EAAE,CAAC,YAAY,CAAC,OAAO,EAAE,OAAO,CAAC,CAAC;gBAC9C,IAAI,CAAC,WAAW,CAAC,GAAG,EAAE,IAAI,CAAC,CAAC;YAC9B,CAAC;QACH,CAAC;QAAC,MAAM,CAAC,CAAC,0CAA0C,CAAC,CAAC;QACtD,OAAO,IAAI,CAAC;IACd,CAAC;IAEO,WAAW,CAAC,GAAW,EAAE,GAAmB;QAClD,IAAI,CAAC;YACH,MAAM,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAuE,CAAC;YACrG,IAAI,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,CAAC,EAAE,CAAC;gBAC/B,KAAK,MAAM,CAAC,IAAI,MAAM,CAAC,IAAI,EAAE,CAAC;oBAC5B,IAAI,CAAC,IAAI,OAAO,CAAC,CAAC,EAAE,KAAK,QAAQ,IAAI,KAAK,CAAC,OAAO,CAAC,CAAC,CAAC,MAAM,CAAC;wBAAE,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,EAAE,CAAC,CAAC,EAAE,EAAE,MAAM,EAAE,CAAC,CAAC,MAAM,EAAE,CAAC,CAAC;gBACzG,CAAC;YACH,CAAC;iBAAM,IAAI,MAAM,CAAC,OAAO,IAAI,OAAO,MAAM,CAAC,OAAO,KAAK,QAAQ,EAAE,CAAC;gBAChE,uDAAuD;gBACvD,KAAK,MAAM,CAAC,EAAE,EAAE,GAAG,CAAC,IAAI,MAAM,CAAC,OAAO,CAAC,MAAM,CAAC,OAAO,CAAC,EAAE,CAAC;oBACvD,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC;wBAAE,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,EAAE,MAAM,EAAE,GAAG,EAAE,CAAC,CAAC;gBACxD,CAAC;YACH,CAAC;QACH,CAAC;QAAC,MAAM,CAAC,CAAC,wBAAwB,CAAC,CAAC;IACtC,CAAC;IAEO,YAAY,CAAC,OAAe;QAClC,IAAI,CAAC;YACH,MAAM,IAAI,GAAG,EAAE,CAAC,QAAQ,CAAC,OAAO,CAAC,CAAC;YAClC,IAAI,IAAI,CAAC,MAAM,EAAE;gBAAE,OAAO,IAAI,CAAC,IAAI,CAAC;YACpC,IAAI,IAAI,CAAC,WAAW,EAAE,EAAE,CAAC;gBACvB,IAAI,KAAK,GAAG,CAAC,CAAC;gBACd,KAAK,MAAM,CAAC,IAAI,EAAE,CAAC,WAAW,CAAC,OAAO,CAAC,EAAE,CAAC;oBACxC,IAAI,CAAC;wBAAC,KAAK,IAAI,EAAE,CAAC,QAAQ,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,EAAE,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC;oBAAC,CAAC;oBAAC,MAAM,CAAC,CAAC,UAAU,CAAC,CAAC;gBAChF,CAAC;gBACD,OAAO,KAAK,CAAC;YACf,CAAC;QACH,CAAC;QAAC,MAAM,CAAC,CAAC,aAAa,CAAC,CAAC;QACzB,OAAO,CAAC,CAAC;IACX,CAAC;IAEO,WAAW,CAAC,QAAgB,EAAE,IAAY;QAChD,EAAE,CAAC,SAAS,CAAC,IAAI,CAAC,OAAO,CAAC,QAAQ,CAAC,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAC;QAC1D,MAAM,GAAG,GAAG,QAAQ,GAAG,MAAM,CAAC;QAC9B,EAAE,CAAC,aAAa,CAAC,GAAG,EAAE,IAAI,CAAC,CAAC;QAC5B,EAAE,CAAC,UAAU,CAAC,GAAG,EAAE,QAAQ,CAAC,CAAC;IAC/B,CAAC;CACF;AAED,SAAS,aAAa,CAAC,CAAW,EAAE,GAAW;IAC7C,IAAI,CAAC,CAAC,MAAM,KAAK,GAAG;QAAE,OAAO,CAAC,CAAC;IAC/B,IAAI,CAAC,CAAC,MAAM,GAAG,GAAG;QAAE,OAAO,CAAC,CAAC,KAAK,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;IAC3C,MAAM,GAAG,GAAG,CAAC,CAAC,KAAK,EAAE,CAAC;IACtB,OAAO,GAAG,CAAC,MAAM,GAAG,GAAG;QAAE,GAAG,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;IACrC,OAAO,GAAG,CAAC;AACb,CAAC;AAED,IAAI,SAAS,GAAyB,IAAI,CAAC;AAC3C,MAAM,UAAU,gBAAgB;IAC9B,IAAI,CAAC,SAAS;QAAE,SAAS,GAAG,IAAI,aAAa,EAAE,CAAC;IAChD,OAAO,SAAS,CAAC;AACnB,CAAC;AACD,MAAM,UAAU,0BAA0B,KAAW,SAAS,GAAG,IAAI,CAAC,CAAC,CAAC"}

package/dist/brain/brain-query-expander.d.ts

@@ -0,0 +1,31 @@
+export type ExpansionTechnique = 'synonyms' | 'abbrev' | 'related' | 'hyde';
+export interface ExpansionOptions {
+    techniques?: ExpansionTechnique[];
+    maxTerms?: number;
+    hydeModel?: string;
+}
+export interface ExpansionResult {
+    original: string;
+    expanded: string;
+    techniques_used: string[];
+    terms_added: string[];
+    hyde_doc?: string;
+    durationMs: number;
+}
+export declare class BrainQueryExpander {
+    /** Expand a query using selected techniques; returns the densified search text. */
+    expand(query: string, opts?: ExpansionOptions): Promise<ExpansionResult>;
+    /**
+     * Run all techniques and return the densest expansion. Best for direct use
+     * by retrievers that don't want to think about technique selection.
+     */
+    combineAll(query: string, opts?: ExpansionOptions): Promise<ExpansionResult>;
+    private lookupRelatedEntities;
+    private generateHyde;
+    private probeHydeModel;
+    private templateHyde;
+    private persist;
+}
+export declare function getBrainQueryExpander(): BrainQueryExpander;
+export declare function resetBrainQueryExpanderForTests(): void;
+//# sourceMappingURL=brain-query-expander.d.ts.map

package/dist/brain/brain-query-expander.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"brain-query-expander.d.ts","sourceRoot":"","sources":["../../src/brain/brain-query-expander.ts"],"names":[],"mappings":"AAkCA,MAAM,MAAM,kBAAkB,GAAG,UAAU,GAAG,QAAQ,GAAG,SAAS,GAAG,MAAM,CAAC;AAE5E,MAAM,WAAW,gBAAgB;IAC/B,UAAU,CAAC,EAAE,kBAAkB,EAAE,CAAC;IAClC,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,SAAS,CAAC,EAAE,MAAM,CAAC;CACpB;AAED,MAAM,WAAW,eAAe;IAC9B,QAAQ,EAAE,MAAM,CAAC;IACjB,QAAQ,EAAE,MAAM,CAAC;IACjB,eAAe,EAAE,MAAM,EAAE,CAAC;IAC1B,WAAW,EAAE,MAAM,EAAE,CAAC;IACtB,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,UAAU,EAAE,MAAM,CAAC;CACpB;AA0LD,qBAAa,kBAAkB;IAC7B,mFAAmF;IAC7E,MAAM,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,CAAC,EAAE,gBAAgB,GAAG,OAAO,CAAC,eAAe,CAAC;IAkF9E;;;OAGG;IACG,UAAU,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,CAAC,EAAE,gBAAgB,GAAG,OAAO,CAAC,eAAe,CAAC;YASpE,qBAAqB;YAiCrB,YAAY;YAyBZ,cAAc;IAsB5B,OAAO,CAAC,YAAY;YAaN,OAAO;CAOtB;AAYD,wBAAgB,qBAAqB,IAAI,kBAAkB,CAG1D;AACD,wBAAgB,+BAA+B,IAAI,IAAI,CAAsB"}