grepmax 0.7.0

package/dist/lib/workers/orchestrator.js

@@ -0,0 +1,350 @@
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
+    };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.WorkerOrchestrator = void 0;
+const fs = __importStar(require("node:fs"));
+const path = __importStar(require("node:path"));
+const transformers_1 = require("@huggingface/transformers");
+const ort = __importStar(require("onnxruntime-node"));
+const uuid_1 = require("uuid");
+const config_1 = require("../../config");
+const chunker_1 = require("../index/chunker");
+const skeleton_1 = require("../skeleton");
+const file_utils_1 = require("../utils/file-utils");
+const colbert_math_1 = require("./colbert-math");
+const colbert_1 = require("./embeddings/colbert");
+const granite_1 = require("./embeddings/granite");
+const mlx_client_1 = require("./embeddings/mlx-client");
+const CACHE_DIR = config_1.PATHS.models;
+const LOG_MODELS = process.env.OSGREP_DEBUG_MODELS === "1" ||
+    process.env.OSGREP_DEBUG_MODELS === "true";
+const log = (...args) => {
+    if (LOG_MODELS)
+        console.log(...args);
+};
+transformers_1.env.cacheDir = CACHE_DIR;
+transformers_1.env.allowLocalModels = true;
+transformers_1.env.allowRemoteModels = true;
+const PROJECT_ROOT = process.env.OSGREP_PROJECT_ROOT
+    ? path.resolve(process.env.OSGREP_PROJECT_ROOT)
+    : process.cwd();
+const LOCAL_MODELS = path.join(PROJECT_ROOT, "models");
+if (fs.existsSync(LOCAL_MODELS)) {
+    transformers_1.env.localModelPath = LOCAL_MODELS;
+    log(`Worker: Using local models from ${LOCAL_MODELS}`);
+}
+class WorkerOrchestrator {
+    constructor() {
+        this.granite = new granite_1.GraniteModel();
+        this.colbert = new colbert_1.ColbertModel();
+        this.chunker = new chunker_1.TreeSitterChunker();
+        this.skeletonizer = new skeleton_1.Skeletonizer();
+        this.initPromise = null;
+        this.vectorDimensions = config_1.CONFIG.VECTOR_DIM;
+    }
+    ensureReady() {
+        return __awaiter(this, void 0, void 0, function* () {
+            if (this.granite.isReady() && this.colbert.isReady()) {
+                return;
+            }
+            if (this.initPromise)
+                return this.initPromise;
+            this.initPromise = (() => __awaiter(this, void 0, void 0, function* () {
+                yield Promise.all([
+                    this.chunker.init(),
+                    this.skeletonizer.init(),
+                    this.granite.load(),
+                    this.colbert.load(),
+                ]);
+            }))().finally(() => {
+                this.initPromise = null;
+            });
+            return this.initPromise;
+        });
+    }
+    computeHybrid(texts, onProgress) {
+        return __awaiter(this, void 0, void 0, function* () {
+            var _a, _b;
+            if (!texts.length)
+                return [];
+            yield this.ensureReady();
+            const results = [];
+            const envBatch = Number.parseInt((_a = process.env.OSGREP_WORKER_BATCH_SIZE) !== null && _a !== void 0 ? _a : "", 10);
+            const BATCH_SIZE = Number.isFinite(envBatch) && envBatch > 0
+                ? Math.max(4, Math.min(16, envBatch))
+                : 16;
+            for (let i = 0; i < texts.length; i += BATCH_SIZE) {
+                if (i > 0)
+                    onProgress === null || onProgress === void 0 ? void 0 : onProgress();
+                const batchTexts = texts.slice(i, i + BATCH_SIZE);
+                // Try MLX GPU server first, fall back to ONNX CPU
+                const denseBatch = (_b = (yield (0, mlx_client_1.mlxEmbed)(batchTexts))) !== null && _b !== void 0 ? _b : (yield this.granite.runBatch(batchTexts));
+                const colbertBatch = yield this.colbert.runBatch(batchTexts, denseBatch, this.vectorDimensions);
+                results.push(...colbertBatch);
+            }
+            onProgress === null || onProgress === void 0 ? void 0 : onProgress();
+            return results;
+        });
+    }
+    chunkFile(pathname, content) {
+        return __awaiter(this, void 0, void 0, function* () {
+            yield this.ensureReady();
+            const { chunks: parsedChunks, metadata } = yield this.chunker.chunk(pathname, content);
+            const anchorChunk = (0, chunker_1.buildAnchorChunk)(pathname, content, metadata);
+            const baseChunks = anchorChunk
+                ? [anchorChunk, ...parsedChunks]
+                : parsedChunks;
+            return baseChunks.map((chunk, idx) => {
+                const chunkWithContext = chunk;
+                return Object.assign(Object.assign({}, chunkWithContext), { context: Array.isArray(chunkWithContext.context)
+                        ? chunkWithContext.context
+                        : [], chunkIndex: typeof chunkWithContext.chunkIndex === "number"
+                        ? chunkWithContext.chunkIndex
+                        : anchorChunk
+                            ? idx - 1
+                            : idx, isAnchor: chunkWithContext.isAnchor === true ||
+                        (anchorChunk ? idx === 0 : false), imports: metadata.imports });
+            });
+        });
+    }
+    toPreparedChunks(path, hash, chunks, skeleton) {
+        var _a, _b;
+        const texts = chunks.map((chunk) => (0, chunker_1.formatChunkText)(chunk, path));
+        const prepared = [];
+        for (let i = 0; i < texts.length; i++) {
+            const chunk = chunks[i];
+            const { content, displayText } = texts[i];
+            const prev = (_a = texts[i - 1]) === null || _a === void 0 ? void 0 : _a.displayText;
+            const next = (_b = texts[i + 1]) === null || _b === void 0 ? void 0 : _b.displayText;
+            prepared.push({
+                id: (0, uuid_1.v4)(),
+                path,
+                hash,
+                content: content, // Now minimal
+                display_text: displayText, // Now rich
+                context_prev: typeof prev === "string" ? prev : undefined,
+                context_next: typeof next === "string" ? next : undefined,
+                start_line: chunk.startLine,
+                end_line: chunk.endLine,
+                chunk_index: chunk.chunkIndex,
+                is_anchor: chunk.isAnchor === true,
+                chunk_type: typeof chunk.type === "string" ? chunk.type : undefined,
+                complexity: chunk.complexity,
+                is_exported: chunk.isExported,
+                defined_symbols: chunk.definedSymbols,
+                referenced_symbols: chunk.referencedSymbols,
+                role: chunk.role,
+                parent_symbol: chunk.parentSymbol,
+                file_skeleton: chunk.isAnchor ? skeleton : undefined,
+            });
+        }
+        return prepared;
+    }
+    processFile(input, onProgress) {
+        return __awaiter(this, void 0, void 0, function* () {
+            const absolutePath = path.isAbsolute(input.path)
+                ? input.path
+                : input.absolutePath
+                    ? input.absolutePath
+                    : path.join(PROJECT_ROOT, input.path);
+            const { buffer, mtimeMs, size } = yield (0, file_utils_1.readFileSnapshot)(absolutePath);
+            const hash = (0, file_utils_1.computeBufferHash)(buffer);
+            if (!(0, file_utils_1.isIndexableFile)(absolutePath, size)) {
+                return { vectors: [], hash, mtimeMs, size, shouldDelete: true };
+            }
+            if (buffer.length === 0 || (0, file_utils_1.hasNullByte)(buffer)) {
+                return { vectors: [], hash, mtimeMs, size, shouldDelete: true };
+            }
+            onProgress === null || onProgress === void 0 ? void 0 : onProgress();
+            yield this.ensureReady();
+            onProgress === null || onProgress === void 0 ? void 0 : onProgress();
+            const content = buffer.toString("utf-8");
+            const chunksPromise = this.chunkFile(input.path, content);
+            // Generate skeleton in parallel
+            const skeletonPromise = this.skeletonizer.skeletonizeFile(input.path, content, {
+                includeSummary: true,
+            });
+            const [chunks, skeletonResult] = yield Promise.all([
+                chunksPromise,
+                skeletonPromise,
+            ]);
+            onProgress === null || onProgress === void 0 ? void 0 : onProgress();
+            if (!chunks.length)
+                return { vectors: [], hash, mtimeMs, size };
+            const preparedChunks = this.toPreparedChunks(input.path, hash, chunks, skeletonResult.success ? skeletonResult.skeleton : undefined);
+            const hybrids = yield this.computeHybrid(preparedChunks.map((chunk) => chunk.content), onProgress);
+            const vectors = preparedChunks.map((chunk, idx) => {
+                var _a;
+                const hybrid = (_a = hybrids[idx]) !== null && _a !== void 0 ? _a : {
+                    dense: new Float32Array(),
+                    colbert: new Int8Array(),
+                    scale: 1,
+                };
+                return Object.assign(Object.assign({}, chunk), { vector: hybrid.dense, colbert: Buffer.from(hybrid.colbert), colbert_scale: hybrid.scale, pooled_colbert_48d: hybrid.pooled_colbert_48d, doc_token_ids: hybrid.token_ids });
+            });
+            onProgress === null || onProgress === void 0 ? void 0 : onProgress();
+            return { vectors, hash, mtimeMs, size };
+        });
+    }
+    encodeQuery(text) {
+        return __awaiter(this, void 0, void 0, function* () {
+            var _a;
+            yield this.ensureReady();
+            // Try MLX GPU server first, fall back to ONNX CPU
+            const mlxResult = yield (0, mlx_client_1.mlxEmbed)([text]);
+            const denseVector = (_a = mlxResult === null || mlxResult === void 0 ? void 0 : mlxResult[0]) !== null && _a !== void 0 ? _a : (yield this.granite.runBatch([text]))[0];
+            const encoded = yield this.colbert.encodeQuery(text);
+            const feeds = {
+                input_ids: new ort.Tensor("int64", encoded.input_ids, [
+                    1,
+                    encoded.input_ids.length,
+                ]),
+                attention_mask: new ort.Tensor("int64", encoded.attention_mask, [
+                    1,
+                    encoded.attention_mask.length,
+                ]),
+            };
+            const sessionOut = yield this.colbert.runSession(feeds);
+            const outputName = this.colbert.getOutputName();
+            const output = sessionOut[outputName];
+            if (!output) {
+                throw new Error("ColBERT session output missing embeddings tensor");
+            }
+            const data = output.data;
+            const [, seq, dim] = output.dims;
+            const matrix = [];
+            for (let s = 0; s < seq; s++) {
+                let sumSq = 0;
+                const offset = s * dim;
+                for (let d = 0; d < dim; d++) {
+                    const val = data[offset + d];
+                    sumSq += val * val;
+                }
+                const norm = Math.sqrt(sumSq);
+                const row = [];
+                if (norm > 1e-9) {
+                    for (let d = 0; d < dim; d++) {
+                        row.push(data[offset + d] / norm);
+                    }
+                }
+                else {
+                    for (let d = 0; d < dim; d++) {
+                        row.push(data[offset + d]);
+                    }
+                }
+                matrix.push(row);
+            }
+            // Compute pooled embedding (mean of tokens)
+            const pooled = new Float32Array(dim);
+            for (const row of matrix) {
+                for (let d = 0; d < dim; d++) {
+                    pooled[d] += row[d];
+                }
+            }
+            // Normalize pooled
+            let sumSq = 0;
+            for (let d = 0; d < dim; d++) {
+                pooled[d] /= matrix.length || 1;
+                sumSq += pooled[d] * pooled[d];
+            }
+            const norm = Math.sqrt(sumSq);
+            if (norm > 1e-9) {
+                for (let d = 0; d < dim; d++) {
+                    pooled[d] /= norm;
+                }
+            }
+            return {
+                dense: Array.from(denseVector !== null && denseVector !== void 0 ? denseVector : []),
+                colbert: matrix,
+                colbertDim: dim,
+                pooled_colbert_48d: Array.from(pooled),
+            };
+        });
+    }
+    rerank(input) {
+        return __awaiter(this, void 0, void 0, function* () {
+            yield this.ensureReady();
+            const queryMatrix = input.query.map((row) => row instanceof Float32Array ? row : new Float32Array(row));
+            return input.docs.map((doc) => {
+                const col = doc.colbert;
+                let colbert;
+                if (col instanceof Int8Array) {
+                    colbert = col;
+                }
+                else if (Buffer.isBuffer(col)) {
+                    colbert = new Int8Array(col.buffer, col.byteOffset, col.byteLength);
+                }
+                else if (col &&
+                    typeof col === "object" &&
+                    "type" in col &&
+                    col.type === "Buffer" &&
+                    Array.isArray(col.data)) {
+                    // IPC serialization fallback (still copies, but unavoidable without SharedArrayBuffer)
+                    colbert = new Int8Array(col.data);
+                }
+                else if (Array.isArray(col)) {
+                    colbert = new Int8Array(col);
+                }
+                else {
+                    colbert = new Int8Array(0);
+                }
+                const seqLen = Math.floor(colbert.length / input.colbertDim);
+                const docMatrix = [];
+                for (let i = 0; i < seqLen; i++) {
+                    const start = i * input.colbertDim;
+                    const row = new Float32Array(input.colbertDim);
+                    for (let d = 0; d < input.colbertDim; d++) {
+                        row[d] = (colbert[start + d] * doc.scale) / 127.0;
+                    }
+                    docMatrix.push(row);
+                }
+                const tokenIds = Array.isArray(doc.token_ids) && doc.token_ids.length === seqLen
+                    ? doc.token_ids
+                    : undefined;
+                return (0, colbert_math_1.maxSim)(queryMatrix, docMatrix, tokenIds);
+            });
+        });
+    }
+}
+exports.WorkerOrchestrator = WorkerOrchestrator;