Test insert 10 items with 1 layer

Signed-off-by: Jay Wang <[email protected]>

src/mememo.ts

@@ -200,6 +200,10 @@ export class HNSW<T = string> {
 
     // Randomly determine the max level of this node
     const level = this._getRandomLevel();
+    console.log('random level:', level);
+
+    // Add this node to the node index first
+    this.nodes.set(key, new Node(key, value));
 
     if (this.entryPointKey !== null) {
       // (1): Search closest point from layers above

test/data/insert-10-layer-0.json

@@ -0,0 +1 @@
+{"1": {"2": 0.1246081047999994, "3": 0.18881115149999972, "4": 0.1741076064999999, "5": 0.15061682040000002, "6": 0.16244587589999981, "7": 0.1346536822999993, "8": 0.12481225549999975, "9": 0.17475681300000034, "10": 0.07815748529999955}, "2": {"1": 0.1246081047999994, "3": 0.12941816250000038, "4": 0.09398657370000041, "5": 0.1260238756000006, "6": 0.15637710450000042, "7": 0.12558542940000061, "8": 0.11024280929999997, "9": 0.15551240220000007, "10": 0.11433565549999969}, "3": {"1": 0.18881115149999972, "2": 0.12941816250000038, "4": 0.12779392509999998, "5": 0.2008425863000003, "6": 0.1572827808999997, "7": 0.17000606770000037, "8": 0.15120083390000005, "9": 0.1945165847000001, "10": 0.1787358560000002}, "4": {"1": 0.1741076064999999, "2": 0.09398657370000041, "3": 0.12779392509999998, "5": 0.1653011399999993, "6": 0.14476001560000018, "7": 0.1559536399999999, "8": 0.16000778209999988, "9": 0.1670294960000005, "10": 0.16677922930000022}, "5": {"3": 0.2008425863000003, "4": 0.1653011399999993, "1": 0.15061682040000002, "2": 0.1260238756000006, "6": 0.1521421429000004, "7": 0.12707044099999976, "8": 0.1286279065999999, "9": 0.1997314214999999, "10": 0.13303564049999994}, "6": {"1": 0.16244587589999981, "2": 0.15637710450000042, "3": 0.1572827808999997, "4": 0.14476001560000018, "5": 0.1521421429000004, "7": 0.15823557500000018, "8": 0.13482364030000005, "9": 0.19377408259999984, "10": 0.13754302259999984}, "7": {"3": 0.17000606770000037, "4": 0.1559536399999999, "6": 0.15823557500000018, "2": 0.12558542940000061, "5": 0.12707044099999976, "1": 0.1346536822999993, "8": 0.10591566870000002, "9": 0.18840105240000038, "10": 0.1465513723000007}, "8": {"4": 0.16000778209999988, "3": 0.15120083390000005, "6": 0.13482364030000005, "2": 0.11024280929999997, "5": 0.1286279065999999, "1": 0.12481225549999975, "7": 0.10591566870000002, "9": 0.15914553510000018, "10": 0.13158060270000005}, "9": {"5": 0.1997314214999999, "3": 0.1945165847000001, "6": 0.19377408259999984, "8": 0.15914553510000018, "4": 0.1670294960000005, "1": 0.17475681300000034, "7": 0.18840105240000038, "2": 0.15551240220000007, "10": 0.18176799640000008}, "10": {"9": 0.18176799640000008, "3": 0.1787358560000002, "7": 0.1465513723000007, "4": 0.16677922930000022, "5": 0.13303564049999994, "2": 0.11433565549999969, "6": 0.13754302259999984, "1": 0.07815748529999955, "8": 0.13158060270000005}}

test/mememo.test.ts

@@ -1,13 +1,19 @@
 import { describe, test, expect, beforeEach } from 'vitest';
 import { add, HNSW } from '../src/mememo';
 import embeddingDataJSON from './data/accident-report-embeddings-100.json';
+import graph10Layer0JSON from './data/insert-10-layer-0.json';
 
 interface EmbeddingData {
   embeddings: number[][];
   reportNumbers: number[];
 }
 const embeddingData = embeddingDataJSON as EmbeddingData;
 
+const graph10Layer0 = graph10Layer0JSON as Record<
+  string,
+  Record<string, number>
+>;
+
 test('add()', () => {
   expect(add(10, 1)).toBe(11);
 });
@@ -20,15 +26,50 @@ test('constructor', () => {
   });
 });
 
-test('insert()', () => {
+//==========================================================================||
+//                                 Insert                                   ||
+//==========================================================================||
+
+test('insert() 10 items, 1 layer', () => {
   const hnsw = new HNSW({
-    distanceFunction: 'cosine'
+    distanceFunction: 'cosine',
+    seed: 20240101
   });
 
-  const embedding = embeddingData.embeddings[0];
-  hnsw.insert('name', embedding);
+  // Insert 10 embeddings
+  const size = 10;
+
+  // The random levels with this seed is 0,0,0,0,0,0,0,0,0,0
+  const reportIDs: string[] = [];
+  for (let i = 0; i < size; i++) {
+    const curReportID = String(embeddingData.reportNumbers[i]);
+    reportIDs.push(curReportID);
+    hnsw.insert(curReportID, embeddingData.embeddings[i]);
+  }
+
+  // There should be only one layer, and all nodes are fully connected
+  expect(hnsw.graphLayers.length).toBe(1);
+
+  for (const reportID of reportIDs) {
+    const curNode = hnsw.graphLayers[0].graph.get(reportID);
+    expect(curNode).to.not.toBeUndefined();
+    expect(curNode!.size).toBe(9);
+
+    // Test the distances
+    const expectedNeighbors = graph10Layer0[reportID];
+    for (const [neighborKey, neighborDistance] of curNode!.entries()) {
+      expect(neighborDistance).toBeCloseTo(
+        expectedNeighbors[neighborKey],
+        1e-6
+      );
+    }
+  }
 });
 
+//==========================================================================||
+//                          Helper Functions                                ||
+//==========================================================================||
+
 test('_getRandomLevel()', () => {
   const hnsw = new HNSW({
     distanceFunction: 'cosine',