rag-lite-ts 2.0.5 → 2.1.1

package/dist/core/binary-index-format.js

@@ -17,7 +17,7 @@
 import { readFileSync, writeFileSync } from 'fs';
 export class BinaryIndexFormat {
     /**
-     * Save index data to binary format
+     * Save index data to binary format (original format for backward compatibility)
      *
      * File structure:
      * - Header (24 bytes): dimensions, maxElements, M, efConstruction, seed, currentSize
@@ -66,7 +66,115 @@ export class BinaryIndexFormat {
         writeFileSync(indexPath, Buffer.from(buffer));
     }
     /**
-     * Load index data from binary format
+     * Save index data to grouped binary format
+     *
+     * File structure:
+     * - Extended Header (40 bytes):
+     *   - Original 6 fields (24 bytes)
+     *   - hasGroups flag (4 bytes)
+     *   - textOffset (4 bytes)
+     *   - textCount (4 bytes)
+     *   - imageOffset (4 bytes)
+     *   - imageCount (4 bytes)
+     * - Data section: [text vectors...][image vectors...]
+     *
+     * @param indexPath Path to save the binary index file
+     * @param data Index data to serialize
+     */
+    static async saveGrouped(indexPath, data) {
+        if (!data.hasContentTypeGroups || !data.textVectors || !data.imageVectors) {
+            // Fallback to original format
+            return this.save(indexPath, data);
+        }
+        const headerSize = 44; // Extended header: 24 + 20 bytes (hasGroups + textOffset + textCount + imageOffset + imageCount)
+        const vectorSize = 4 + (data.dimensions * 4); // id + vector
+        // Calculate offsets and total size
+        const textOffset = headerSize;
+        const imageOffset = textOffset + (data.textVectors.length * vectorSize);
+        const totalSize = imageOffset + (data.imageVectors.length * vectorSize);
+        const buffer = new ArrayBuffer(totalSize);
+        const view = new DataView(buffer);
+        let offset = 0;
+        // Write extended header (40 bytes, all little-endian)
+        if (offset + 40 > buffer.byteLength) {
+            throw new Error(`Header write would exceed buffer bounds: offset=${offset}, headerSize=40, bufferSize=${buffer.byteLength}`);
+        }
+        view.setUint32(offset, data.dimensions, true);
+        offset += 4;
+        view.setUint32(offset, data.maxElements, true);
+        offset += 4;
+        view.setUint32(offset, data.M, true);
+        offset += 4;
+        view.setUint32(offset, data.efConstruction, true);
+        offset += 4;
+        view.setUint32(offset, data.seed, true);
+        offset += 4;
+        view.setUint32(offset, data.currentSize, true);
+        offset += 4;
+        // Extended fields
+        view.setUint32(offset, 1, true);
+        offset += 4; // hasGroups = 1
+        view.setUint32(offset, textOffset, true);
+        offset += 4;
+        view.setUint32(offset, data.textVectors.length, true);
+        offset += 4;
+        view.setUint32(offset, imageOffset, true);
+        offset += 4;
+        view.setUint32(offset, data.imageVectors.length, true);
+        offset += 4;
+        // Write text vectors
+        for (const item of data.textVectors) {
+            // Ensure 4-byte alignment
+            if (offset % 4 !== 0) {
+                throw new Error(`Offset ${offset} is not 4-byte aligned`);
+            }
+            // Check bounds before writing
+            if (offset + 4 > buffer.byteLength) {
+                throw new Error(`ID write would exceed buffer bounds: offset=${offset}, bufferSize=${buffer.byteLength}`);
+            }
+            // Write vector ID
+            view.setUint32(offset, item.id, true);
+            offset += 4;
+            // Check bounds for vector data
+            const vectorDataSize = item.vector.length * 4;
+            if (offset + vectorDataSize > buffer.byteLength) {
+                throw new Error(`Vector data write would exceed buffer bounds: offset=${offset}, dataSize=${vectorDataSize}, bufferSize=${buffer.byteLength}`);
+            }
+            // Write vector data
+            for (let i = 0; i < item.vector.length; i++) {
+                view.setFloat32(offset, item.vector[i], true);
+                offset += 4;
+            }
+        }
+        // Write image vectors
+        for (const item of data.imageVectors) {
+            // Ensure 4-byte alignment
+            if (offset % 4 !== 0) {
+                throw new Error(`Offset ${offset} is not 4-byte aligned`);
+            }
+            // Check bounds before writing
+            if (offset + 4 > buffer.byteLength) {
+                throw new Error(`ID write would exceed buffer bounds: offset=${offset}, bufferSize=${buffer.byteLength}`);
+            }
+            // Write vector ID
+            view.setUint32(offset, item.id, true);
+            offset += 4;
+            // Check bounds for vector data
+            const vectorDataSize = item.vector.length * 4;
+            if (offset + vectorDataSize > buffer.byteLength) {
+                throw new Error(`Vector data write would exceed buffer bounds: offset=${offset}, dataSize=${vectorDataSize}, bufferSize=${buffer.byteLength}`);
+            }
+            // Write vector data
+            for (let i = 0; i < item.vector.length; i++) {
+                view.setFloat32(offset, item.vector[i], true);
+                offset += 4;
+            }
+        }
+        // Write to file
+        writeFileSync(indexPath, Buffer.from(buffer));
+    }
+    /**
+     * Load index data from binary format (supports both original and grouped formats)
      *
      * Uses zero-copy Float32Array views for efficient loading.
      * Copies the views to ensure data persistence after buffer lifecycle.
@@ -78,7 +186,7 @@ export class BinaryIndexFormat {
         const buffer = readFileSync(indexPath);
         const view = new DataView(buffer.buffer, buffer.byteOffset, buffer.byteLength);
         let offset = 0;
-        // Read header (24 bytes, all little-endian)
+        // Read basic header (24 bytes, all little-endian)
         const dimensions = view.getUint32(offset, true);
         offset += 4;
         const maxElements = view.getUint32(offset, true);
@@ -91,32 +199,93 @@ export class BinaryIndexFormat {
         offset += 4;
         const currentSize = view.getUint32(offset, true);
         offset += 4;
-        // Read vectors
-        const vectors = [];
-        for (let i = 0; i < currentSize; i++) {
-            // Ensure 4-byte alignment (should always be true with our format)
-            if (offset % 4 !== 0) {
-                throw new Error(`Offset ${offset} is not 4-byte aligned`);
+        // Check if this is the extended grouped format (40+ bytes header)
+        const hasGroups = buffer.byteLength >= 40 ? view.getUint32(offset, true) : 0;
+        if (hasGroups === 1 && buffer.byteLength >= 40) {
+            // Load grouped format
+            const textOffset = view.getUint32(offset + 4, true);
+            const textCount = view.getUint32(offset + 8, true);
+            const imageOffset = view.getUint32(offset + 12, true);
+            const imageCount = view.getUint32(offset + 16, true);
+            // Load text vectors
+            const textVectors = [];
+            offset = textOffset;
+            for (let i = 0; i < textCount; i++) {
+                // Ensure 4-byte alignment
+                if (offset % 4 !== 0) {
+                    throw new Error(`Offset ${offset} is not 4-byte aligned`);
+                }
+                // Read vector ID
+                const id = view.getUint32(offset, true);
+                offset += 4;
+                // Zero-copy Float32Array view
+                const vectorView = new Float32Array(buffer.buffer, buffer.byteOffset + offset, dimensions);
+                // Copy to avoid buffer lifecycle issues
+                const vector = new Float32Array(vectorView);
+                offset += dimensions * 4;
+                textVectors.push({ id, vector });
             }
-            // Read vector ID
-            const id = view.getUint32(offset, true);
-            offset += 4;
-            // Zero-copy Float32Array view (fast!)
-            const vectorView = new Float32Array(buffer.buffer, buffer.byteOffset + offset, dimensions);
-            // Copy to avoid buffer lifecycle issues
-            const vector = new Float32Array(vectorView);
-            offset += dimensions * 4;
-            vectors.push({ id, vector });
+            // Load image vectors
+            const imageVectors = [];
+            offset = imageOffset;
+            for (let i = 0; i < imageCount; i++) {
+                // Ensure 4-byte alignment
+                if (offset % 4 !== 0) {
+                    throw new Error(`Offset ${offset} is not 4-byte aligned`);
+                }
+                // Read vector ID
+                const id = view.getUint32(offset, true);
+                offset += 4;
+                // Zero-copy Float32Array view
+                const vectorView = new Float32Array(buffer.buffer, buffer.byteOffset + offset, dimensions);
+                // Copy to avoid buffer lifecycle issues
+                const vector = new Float32Array(vectorView);
+                offset += dimensions * 4;
+                imageVectors.push({ id, vector });
+            }
+            // Combine all vectors for backward compatibility
+            const allVectors = [...textVectors, ...imageVectors];
+            return {
+                dimensions,
+                maxElements,
+                M,
+                efConstruction,
+                seed,
+                currentSize,
+                vectors: allVectors,
+                hasContentTypeGroups: true,
+                textVectors,
+                imageVectors
+            };
+        }
+        else {
+            // Load original format
+            const vectors = [];
+            for (let i = 0; i < currentSize; i++) {
+                // Ensure 4-byte alignment (should always be true with our format)
+                if (offset % 4 !== 0) {
+                    throw new Error(`Offset ${offset} is not 4-byte aligned`);
+                }
+                // Read vector ID
+                const id = view.getUint32(offset, true);
+                offset += 4;
+                // Zero-copy Float32Array view (fast!)
+                const vectorView = new Float32Array(buffer.buffer, buffer.byteOffset + offset, dimensions);
+                // Copy to avoid buffer lifecycle issues
+                const vector = new Float32Array(vectorView);
+                offset += dimensions * 4;
+                vectors.push({ id, vector });
+            }
+            return {
+                dimensions,
+                maxElements,
+                M,
+                efConstruction,
+                seed,
+                currentSize,
+                vectors
+            };
         }
-        return {
-            dimensions,
-            maxElements,
-            M,
-            efConstruction,
-            seed,
-            currentSize,
-            vectors
-        };
     }
 }
 //# sourceMappingURL=binary-index-format.js.map