@zvec/zvec-grep 0.1.0

package/dist/engine/storage/zvec.js

@@ -0,0 +1,706 @@
+import { ZVecCollectionSchema, ZVecCreateAndOpen, ZVecDataType, ZVecIndexType, ZVecInitialize, ZVecLogLevel, ZVecMetricType, ZVecOpen, } from "@zvec/zvec";
+import { closeSync, existsSync, mkdirSync, openSync } from "node:fs";
+import { join } from "node:path";
+import { EngineError } from "../errors/index.js";
+import { writeJsonFileSync } from "../utils/json.js";
+import { normalizePath } from "../utils/path.js";
+const ENTITY_VECTOR_FIELD = "embedding";
+const ENTITY_TEXT_FIELD = "text";
+const FILES_MANIFEST = "files.json";
+const ZVEC_PATH = "index.zvec";
+const ZVEC_LOCK_FILE = "LOCK";
+const NO_MATCH_FILTER = "file_id = '__zvec_grep_no_match__'";
+const ZVEC_UPSERT_BATCH_SIZE = 1024;
+const ZVEC_OPEN_RETRY_ATTEMPTS = 8;
+const ZVEC_OPEN_RETRY_BASE_DELAY_MS = 100;
+const ZVEC_OPEN_RETRY_MAX_DELAY_MS = 1000;
+let zvecInitialized = false;
+export class ZvecCollectionStorage {
+    collectionPath;
+    embedding;
+    readOnly;
+    collection;
+    filesPath;
+    filesById = new Map();
+    filesByAbsolutePath = new Map();
+    needsOptimize = false;
+    constructor(collectionPath, embedding, files = [], readOnly = false) {
+        this.collectionPath = collectionPath;
+        this.embedding = embedding;
+        this.readOnly = readOnly;
+        initializeZvec();
+        if (!readOnly) {
+            mkdirSync(collectionPath, { recursive: true });
+        }
+        this.filesPath = join(collectionPath, FILES_MANIFEST);
+        for (const file of files) {
+            this.rememberFile(file);
+        }
+        const zvecPath = join(collectionPath, ZVEC_PATH);
+        if (existsSync(zvecPath)) {
+            this.collection = openZvecCollection(zvecPath, readOnly, "open", () => ZVecOpen(zvecPath, { readOnly }));
+        }
+        else if (readOnly) {
+            throw new EngineError("zvec collection storage does not exist", {
+                code: "ZVEC_GREP.ENGINE.STORAGE.ZVEC_COLLECTION_MISSING",
+                context: `path=${zvecPath}`,
+            });
+        }
+        else {
+            this.collection = openZvecCollection(zvecPath, readOnly, "create", () => ZVecCreateAndOpen(zvecPath, createSchema(embedding)));
+        }
+    }
+    static filesManifestPath(collectionPath) {
+        return join(collectionPath, FILES_MANIFEST);
+    }
+    static emptyManifest() {
+        return { files: [] };
+    }
+    getFileById(fileId) {
+        const file = this.filesById.get(fileId);
+        return file ? fileRecordToInfo(file) : null;
+    }
+    getFileByPath(absolutePath) {
+        const file = this.filesByAbsolutePath.get(normalizePath(absolutePath));
+        return file ? fileRecordToInfo(file) : null;
+    }
+    listFiles() {
+        return [...this.filesById.values()].map(fileRecordToInfo);
+    }
+    listEntitiesByFile(fileId, options = {}) {
+        const file = this.filesById.get(fileId);
+        if (!file) {
+            return [];
+        }
+        const offset = options.offset ?? 0;
+        const limit = options.limit ?? file.entityIds.length;
+        const ids = file.entityIds.slice(offset, offset + limit);
+        return this.fetchStoredEntities(ids);
+    }
+    getEntity(entityId, options = {}) {
+        const fragment = this.getFragment(entityId, options);
+        if (!fragment) {
+            return null;
+        }
+        if (fragment.fragment.group && fragment.fragment.group !== fragment.fragment.id) {
+            return this.getEntity(fragment.fragment.group, options);
+        }
+        const entity = fragmentToEntity(fragment.fragment);
+        return fragment.vector
+            ? { file: fragment.file, entity, vector: fragment.vector }
+            : { file: fragment.file, entity };
+    }
+    getFragment(fragmentId, options = {}) {
+        const docs = this.collection.fetchSync({
+            ids: fragmentId,
+            includeVector: options.includeVector ?? false,
+        });
+        const doc = docs[fragmentId];
+        if (!doc) {
+            return null;
+        }
+        const stored = this.docToStoredFragment(doc);
+        if (!stored) {
+            return null;
+        }
+        const vector = options.includeVector
+            ? readVector(doc.vectors[ENTITY_VECTOR_FIELD])
+            : undefined;
+        return vector ? { ...stored, vector } : stored;
+    }
+    upsertFile(file, fragments, vectors) {
+        if (fragments.length !== vectors.length) {
+            throw new EngineError("Storage received mismatched entity/vector counts", {
+                code: "ZVEC_GREP.ENGINE.STORAGE.ENTITY_VECTOR_COUNT_MISMATCH",
+                context: `fileId=${file.id} fragmentCount=${fragments.length} vectorCount=${vectors.length}`,
+            });
+        }
+        validateFragmentGroups(file.id, fragments);
+        this.markFileDirty(file);
+        this.deleteFileDocuments(file.id);
+        const now = Date.now();
+        const entityIds = publicEntityIds(fragments);
+        const indexedFile = {
+            ...file,
+            absolutePath: normalizePath(file.absolutePath),
+            indexStatus: {
+                indexedTime: now,
+                entityCount: entityIds.length,
+            },
+            entityIds,
+        };
+        if (fragments.length > 0) {
+            const docs = fragments.map((fragment, index) => ({
+                id: fragment.id,
+                vectors: {
+                    [ENTITY_VECTOR_FIELD]: vectors[index],
+                },
+                fields: fragmentToFields(indexedFile, fragment, index),
+            }));
+            this.upsertDocs(file.id, docs);
+            this.needsOptimize = true;
+        }
+        this.rememberFile(indexedFile);
+        this.persistFiles();
+    }
+    upsertDocs(fileId, docs) {
+        for (let start = 0; start < docs.length; start += ZVEC_UPSERT_BATCH_SIZE) {
+            const batch = docs.slice(start, start + ZVEC_UPSERT_BATCH_SIZE);
+            const statuses = this.collection.upsertSync(batch);
+            const failed = statuses.find((status) => !status.ok);
+            if (failed) {
+                throw new EngineError("zvec failed to upsert entity documents", {
+                    code: "ZVEC_GREP.ENGINE.STORAGE.ZVEC_UPSERT_FAILED",
+                    context: `fileId=${fileId} batchStart=${start} batchSize=${batch.length} code=${failed.code} message=${failed.message}`,
+                });
+            }
+        }
+    }
+    markFileFailed(file, error) {
+        this.deleteFileDocuments(file.id);
+        this.rememberFile({
+            ...file,
+            absolutePath: normalizePath(file.absolutePath),
+            indexStatus: {
+                indexedTime: null,
+                entityCount: 0,
+                error,
+            },
+            entityIds: [],
+        });
+        this.persistFiles();
+    }
+    deleteFile(fileId) {
+        const existing = this.filesById.get(fileId);
+        this.deleteFileDocuments(fileId);
+        if (existing) {
+            this.filesById.delete(fileId);
+            this.filesByAbsolutePath.delete(normalizePath(existing.absolutePath));
+        }
+        this.persistFiles();
+    }
+    searchFts(query, limit, filter) {
+        const zvecFilter = buildFilter(filter);
+        const docs = this.collection.querySync({
+            fieldName: ENTITY_TEXT_FIELD,
+            fts: { matchString: query },
+            ...(zvecFilter ? { filter: zvecFilter } : {}),
+            topk: limit,
+            includeVector: false,
+        });
+        return docsToHits(docs, "fts", this);
+    }
+    searchVector(vector, limit, filter) {
+        const zvecFilter = buildFilter(filter);
+        const docs = this.collection.querySync({
+            fieldName: ENTITY_VECTOR_FIELD,
+            vector: [...vector],
+            ...(zvecFilter ? { filter: zvecFilter } : {}),
+            topk: limit,
+            includeVector: false,
+        });
+        return docsToHits(docs, "vector", this);
+    }
+    optimize() {
+        if (this.needsOptimize) {
+            this.collection.optimizeSync();
+            this.needsOptimize = false;
+        }
+    }
+    close() {
+        this.optimize();
+        this.collection.closeSync();
+    }
+    fetchStoredEntities(ids) {
+        if (ids.length === 0) {
+            return [];
+        }
+        const docs = this.collection.fetchSync({
+            ids: [...ids],
+            includeVector: false,
+        });
+        const entities = [];
+        for (const id of ids) {
+            const doc = docs[id];
+            const entity = doc ? this.docToStoredEntity(doc) : null;
+            if (entity) {
+                entities.push(entity);
+            }
+        }
+        return entities;
+    }
+    docToStoredEntity(doc) {
+        const stored = this.docToStoredFragment(doc);
+        if (!stored) {
+            return null;
+        }
+        if (stored.fragment.group && stored.fragment.group !== stored.fragment.id) {
+            return this.getEntity(stored.fragment.group);
+        }
+        return {
+            file: stored.file,
+            entity: fragmentToEntity(stored.fragment),
+        };
+    }
+    docToStoredFragment(doc) {
+        const fileId = readStringField(doc, "file_id");
+        const file = this.filesById.get(fileId);
+        if (!file) {
+            return null;
+        }
+        return {
+            file: fileRecordToInfo(file),
+            fragment: {
+                id: doc.id,
+                group: readNullableStringFieldFromFields(doc.fields, "group") ?? undefined,
+                fileId,
+                range: parseRange(readStringField(doc, "range_json")),
+                content: parseContent(doc.fields),
+                metadata: parseMetadata(doc.fields),
+            },
+        };
+    }
+    rememberFile(file) {
+        const normalized = {
+            ...file,
+            absolutePath: normalizePath(file.absolutePath),
+        };
+        this.filesById.set(normalized.id, normalized);
+        this.filesByAbsolutePath.set(normalized.absolutePath, normalized);
+    }
+    markFileDirty(file) {
+        this.rememberFile({
+            ...file,
+            absolutePath: normalizePath(file.absolutePath),
+            indexStatus: {
+                indexedTime: null,
+                entityCount: 0,
+            },
+            entityIds: [],
+        });
+        this.persistFiles();
+    }
+    deleteFileDocuments(fileId) {
+        this.collection.deleteByFilterSync(`file_id = ${quoteFilterString(fileId)}`);
+        this.needsOptimize = true;
+    }
+    persistFiles() {
+        const manifest = {
+            files: [...this.filesById.values()].sort((left, right) => left.relativePath.localeCompare(right.relativePath)),
+        };
+        writeJsonFileSync(this.filesPath, manifest);
+    }
+}
+function fileRecordToInfo(file) {
+    const { entityIds: _entityIds, ...info } = file;
+    return info;
+}
+function fragmentToEntity(fragment) {
+    return {
+        id: publicEntityId(fragment),
+        fileId: fragment.fileId,
+        range: fragment.range,
+        content: fragment.content,
+        metadata: fragment.metadata,
+    };
+}
+function publicEntityId(fragment) {
+    return fragment.group ?? fragment.id;
+}
+function publicEntityIds(fragments) {
+    const ids = new Set();
+    for (const fragment of fragments) {
+        if (!fragment.group || fragment.group === fragment.id) {
+            ids.add(publicEntityId(fragment));
+        }
+    }
+    return [...ids];
+}
+function validateFragmentGroups(fileId, fragments) {
+    const ids = new Set();
+    const groups = new Map();
+    for (const fragment of fragments) {
+        if (fragment.fileId !== fileId) {
+            throw new EngineError("Entity fragment belongs to the wrong file", {
+                code: "ZVEC_GREP.ENGINE.STORAGE.FRAGMENT_FILE_MISMATCH",
+                context: `fileId=${fileId} fragmentId=${fragment.id} fragmentFileId=${fragment.fileId}`,
+            });
+        }
+        if (ids.has(fragment.id)) {
+            throw new EngineError("Duplicate entity fragment id", {
+                code: "ZVEC_GREP.ENGINE.STORAGE.DUPLICATE_FRAGMENT_ID",
+                context: `fileId=${fileId} fragmentId=${fragment.id}`,
+            });
+        }
+        ids.add(fragment.id);
+        if (fragment.group) {
+            const group = groups.get(fragment.group) ?? [];
+            group.push(fragment);
+            groups.set(fragment.group, group);
+        }
+    }
+    for (const [groupId, group] of groups) {
+        const majorCount = group.filter((fragment) => fragment.id === groupId).length;
+        if (majorCount !== 1) {
+            throw new EngineError("Fragment group must have exactly one major fragment", {
+                code: "ZVEC_GREP.ENGINE.STORAGE.INVALID_FRAGMENT_GROUP",
+                context: `fileId=${fileId} group=${groupId} majorCount=${majorCount}`,
+            });
+        }
+    }
+}
+function initializeZvec() {
+    if (zvecInitialized) {
+        return;
+    }
+    ZVecInitialize({ logLevel: ZVecLogLevel.WARN });
+    zvecInitialized = true;
+}
+function openZvecCollection(zvecPath, readOnly, action, open) {
+    const lockWritable = canTouchZvecLock(zvecPath);
+    let lastError;
+    for (let attempt = 0; attempt < ZVEC_OPEN_RETRY_ATTEMPTS; attempt++) {
+        try {
+            return open();
+        }
+        catch (error) {
+            lastError = error;
+            if (attempt + 1 >= ZVEC_OPEN_RETRY_ATTEMPTS
+                || !isRetryableZvecOpenError(error, lockWritable)) {
+                break;
+            }
+            sleepSync(zvecOpenRetryDelayMs(attempt));
+        }
+    }
+    throw new EngineError("Failed to open zvec collection storage", {
+        code: "ZVEC_GREP.ENGINE.STORAGE.ZVEC_OPEN_FAILED",
+        context: `path=${zvecPath} action=${action} readOnly=${readOnly} attempts=${ZVEC_OPEN_RETRY_ATTEMPTS}`,
+        cause: lastError,
+    });
+}
+function canTouchZvecLock(zvecPath) {
+    const lockPath = join(zvecPath, ZVEC_LOCK_FILE);
+    if (!existsSync(lockPath)) {
+        return true;
+    }
+    try {
+        const fd = openSync(lockPath, "a");
+        closeSync(fd);
+        return true;
+    }
+    catch {
+        return false;
+    }
+}
+function isRetryableZvecOpenError(error, lockWritable) {
+    if (!lockWritable) {
+        return false;
+    }
+    const message = error instanceof Error ? error.message : String(error);
+    if (!/lock/i.test(message)) {
+        return false;
+    }
+    return !/permission|access denied|eacces|eperm|read-only/i.test(message);
+}
+function zvecOpenRetryDelayMs(attempt) {
+    const exponential = ZVEC_OPEN_RETRY_BASE_DELAY_MS * 2 ** attempt;
+    return Math.min(exponential, ZVEC_OPEN_RETRY_MAX_DELAY_MS);
+}
+function sleepSync(ms) {
+    const buffer = new SharedArrayBuffer(4);
+    const values = new Int32Array(buffer);
+    Atomics.wait(values, 0, 0, ms);
+}
+function createSchema(embedding) {
+    return new ZVecCollectionSchema({
+        name: "zvec_grep_entities",
+        fields: [
+            indexedStringField("group", true),
+            indexedStringField("file_id"),
+            stringField("content_kind"),
+            stringField("content_hash", true),
+            stringField("metadata_kind", true),
+            indexedStringField("symbol_type", true),
+            indexedStringField("symbol_name", true),
+            stringField("symbol_scope", true),
+            stringField("symbol_signature", true),
+            stringField("symbol_doc", true),
+            stringField("symbol_modifiers", true),
+            stringField("node_type", true),
+            stringField("heading", true),
+            {
+                name: "heading_level",
+                dataType: ZVecDataType.INT32,
+                nullable: true,
+            },
+            {
+                name: ENTITY_TEXT_FIELD,
+                dataType: ZVecDataType.STRING,
+                nullable: false,
+                indexParams: {
+                    indexType: ZVecIndexType.FTS,
+                    tokenizerName: "jieba",
+                    filters: ["lowercase"],
+                },
+            },
+            {
+                name: "fragment_index",
+                dataType: ZVecDataType.INT32,
+                nullable: false,
+            },
+            { name: "range_json", dataType: ZVecDataType.STRING, nullable: false },
+            { name: "content_base64", dataType: ZVecDataType.STRING, nullable: true },
+            { name: "image_format", dataType: ZVecDataType.STRING, nullable: true },
+        ],
+        vectors: {
+            name: ENTITY_VECTOR_FIELD,
+            dataType: ZVecDataType.VECTOR_FP32,
+            dimension: embedding.dimension,
+            indexParams: {
+                indexType: ZVecIndexType.HNSW,
+                metricType: metricToZvec(embedding.metric),
+            },
+        },
+    });
+}
+function indexedStringField(name, nullable = false) {
+    return {
+        name,
+        dataType: ZVecDataType.STRING,
+        nullable,
+        indexParams: { indexType: ZVecIndexType.INVERT },
+    };
+}
+function stringField(name, nullable = false) {
+    return {
+        name,
+        dataType: ZVecDataType.STRING,
+        nullable,
+    };
+}
+function metricToZvec(metric) {
+    if (metric === "cosine") {
+        return ZVecMetricType.COSINE;
+    }
+    if (metric === "dot") {
+        return ZVecMetricType.IP;
+    }
+    if (metric === "euclidean") {
+        return ZVecMetricType.L2;
+    }
+    throw new EngineError("Unsupported vector metric for zvec storage", {
+        code: "ZVEC_GREP.ENGINE.STORAGE.UNSUPPORTED_VECTOR_METRIC",
+        context: `metric=${metric}`,
+    });
+}
+function fragmentToFields(file, fragment, fragmentIndex) {
+    const contentFields = contentToFields(fragment.content);
+    return {
+        ...optionalFields({
+            group: fragment.group,
+            content_hash: file.contentHash,
+        }),
+        file_id: file.id,
+        content_kind: fragment.content.kind,
+        ...metadataToFields(fragment.metadata),
+        fragment_index: fragmentIndex,
+        range_json: JSON.stringify(fragment.range),
+        ...contentFields,
+    };
+}
+function contentToFields(content) {
+    if (content.kind === "text") {
+        return {
+            text: content.text,
+        };
+    }
+    return {
+        text: `[image:${content.format}]`,
+        content_base64: Buffer.from(content.data).toString("base64"),
+        image_format: content.format,
+    };
+}
+function metadataToFields(metadata) {
+    if (!metadata) {
+        return {};
+    }
+    if (metadata.kind === "code") {
+        return {
+            metadata_kind: "code",
+            symbol_type: metadata.symbolType,
+            ...optionalFields({
+                symbol_name: metadata.symbolName,
+                symbol_scope: metadata.scope,
+                symbol_signature: metadata.signature,
+                symbol_doc: metadata.doc,
+                symbol_modifiers: metadata.modifiers.length > 0 ? metadata.modifiers.join(" ") : null,
+                node_type: metadata.nodeType,
+            }),
+        };
+    }
+    return {
+        metadata_kind: "markdown",
+        ...optionalFields({
+            symbol_scope: metadata.scope,
+            heading: metadata.heading,
+            heading_level: metadata.level,
+        }),
+    };
+}
+function optionalFields(fields) {
+    const result = {};
+    for (const [key, value] of Object.entries(fields)) {
+        if (value !== null && value !== undefined) {
+            result[key] = value;
+        }
+    }
+    return result;
+}
+function parseContent(fields) {
+    const kind = fields.content_kind;
+    if (kind === "text") {
+        return {
+            kind: "text",
+            text: readStringFieldFromFields(fields, "text"),
+        };
+    }
+    if (kind === "image") {
+        return {
+            kind: "image",
+            data: Buffer.from(readStringFieldFromFields(fields, "content_base64"), "base64"),
+            format: readStringFieldFromFields(fields, "image_format"),
+        };
+    }
+    throw new EngineError("Stored entity has unsupported content kind", {
+        code: "ZVEC_GREP.ENGINE.STORAGE.UNSUPPORTED_STORED_CONTENT_KIND",
+        context: `contentKind=${String(kind)}`,
+    });
+}
+function parseMetadata(fields) {
+    const kind = fields.metadata_kind;
+    if (kind === "code") {
+        return {
+            kind: "code",
+            symbolType: readStringFieldFromFields(fields, "symbol_type"),
+            symbolName: readNullableStringFieldFromFields(fields, "symbol_name"),
+            scope: readNullableStringFieldFromFields(fields, "symbol_scope"),
+            nodeType: readNullableStringFieldFromFields(fields, "node_type"),
+            signature: readNullableStringFieldFromFields(fields, "symbol_signature"),
+            doc: readNullableStringFieldFromFields(fields, "symbol_doc"),
+            modifiers: readCodeModifiers(readNullableStringFieldFromFields(fields, "symbol_modifiers")),
+        };
+    }
+    if (kind === "markdown") {
+        return {
+            kind: "markdown",
+            heading: readNullableStringFieldFromFields(fields, "heading"),
+            level: readNullableNumberFieldFromFields(fields, "heading_level"),
+            scope: readNullableStringFieldFromFields(fields, "symbol_scope"),
+        };
+    }
+    return undefined;
+}
+function readCodeModifiers(value) {
+    if (!value) {
+        return [];
+    }
+    return value.split(/\s+/).filter((item) => item === "exported"
+        || item === "async"
+        || item === "static"
+        || item === "public"
+        || item === "private"
+        || item === "protected"
+        || item === "internal");
+}
+function parseRange(value) {
+    return JSON.parse(value);
+}
+function docsToHits(docs, path, storage) {
+    const hits = [];
+    for (const doc of docs) {
+        const stored = storage.docToStoredFragment(doc);
+        if (stored) {
+            hits.push({
+                ...stored,
+                path,
+                score: doc.score,
+            });
+        }
+    }
+    return hits;
+}
+function readStringField(doc, field) {
+    return readStringFieldFromFields(doc.fields, field);
+}
+function readStringFieldFromFields(fields, field) {
+    const value = fields[field];
+    if (typeof value === "string") {
+        return value;
+    }
+    if (value === null || value === undefined) {
+        return "";
+    }
+    return String(value);
+}
+function readNullableStringFieldFromFields(fields, field) {
+    const value = fields[field];
+    if (value === null || value === undefined || value === "") {
+        return null;
+    }
+    return String(value);
+}
+function readNullableNumberFieldFromFields(fields, field) {
+    const value = fields[field];
+    if (typeof value === "number" && Number.isFinite(value)) {
+        return value === 0 ? null : value;
+    }
+    if (typeof value === "string" && value.length > 0) {
+        const parsed = Number.parseInt(value, 10);
+        return Number.isFinite(parsed) && parsed !== 0 ? parsed : null;
+    }
+    return null;
+}
+function readVector(vector) {
+    if (!vector) {
+        return undefined;
+    }
+    if (Array.isArray(vector)) {
+        return vector;
+    }
+    if (ArrayBuffer.isView(vector)) {
+        return [...vector];
+    }
+    return undefined;
+}
+function buildFilter(filter) {
+    if (!filter) {
+        return undefined;
+    }
+    const clauses = [];
+    if (filter.fileIds) {
+        clauses.push(buildNonEmptyInFilter("file_id", filter.fileIds));
+    }
+    if (filter.groupIds) {
+        clauses.push(buildNonEmptyInFilter("group", filter.groupIds));
+    }
+    if (filter.symbolNames) {
+        clauses.push(buildNonEmptyInFilter("symbol_name", filter.symbolNames));
+    }
+    if (filter.symbolTypes) {
+        clauses.push(buildNonEmptyInFilter("symbol_type", filter.symbolTypes));
+    }
+    return clauses.length > 0 ? clauses.join(" AND ") : undefined;
+}
+function buildNonEmptyInFilter(field, values) {
+    return values.length > 0 ? buildInFilter(field, values) : NO_MATCH_FILTER;
+}
+function buildInFilter(field, values) {
+    if (values.length === 1) {
+        return `${field} = ${quoteFilterString(values[0])}`;
+    }
+    return `(${values.map((value) => `${field} = ${quoteFilterString(value)}`).join(" OR ")})`;
+}
+function quoteFilterString(value) {
+    return `'${value.replaceAll("\\", "\\\\").replaceAll("'", "\\'")}'`;
+}
+//# sourceMappingURL=zvec.js.map