implement vlinks_withscores

Author: nkaradzhov

packages/client/lib/commands/VLINKS_WITHSCORES.spec.ts

@@ -0,0 +1,72 @@
+import { strict as assert } from 'node:assert';
+import testUtils, { GLOBAL } from '../test-utils';
+import { parseArgs } from './generic-transformers';
+import VLINKS_WITHSCORES from './VLINKS_WITHSCORES';
+
+describe('VLINKS WITHSCORES', () => {
+  it('transformArguments', () => {
+    assert.deepEqual(parseArgs(VLINKS_WITHSCORES, 'key', 'element'), [
+      'VLINKS',
+      'key',
+      'element',
+      'WITHSCORES'
+    ]);
+  });
+
+  testUtils.testAll(
+    'vLinksWithScores',
+    async client => {
+      // Create a vector set with multiple elements to build HNSW graph layers
+      await client.vAdd('key', [1.0, 2.0, 3.0], 'element1');
+      await client.vAdd('key', [1.1, 2.1, 3.1], 'element2');
+      await client.vAdd('key', [1.2, 2.2, 3.2], 'element3');
+      await client.vAdd('key', [2.0, 3.0, 4.0], 'element4');
+
+      const result = await client.vLinksWithScores('key', 'element1');
+
+      assert.ok(Array.isArray(result));
+
+      for (const layer of result) {
+        assert.equal(
+          typeof layer,
+          'object'
+        );
+      }
+
+      assert.ok(result.length >= 1, 'Should have at least layer 0');
+    },
+    {
+      client: GLOBAL.SERVERS.OPEN,
+      cluster: GLOBAL.CLUSTERS.OPEN
+    }
+  );
+
+  testUtils.testWithClient(
+    'vLinksWithScores with RESP3',
+    async client => {
+      await client.vAdd('resp3-key', [1.0, 2.0, 3.0], 'element1');
+      await client.vAdd('resp3-key', [1.1, 2.1, 3.1], 'element2');
+      await client.vAdd('resp3-key', [1.2, 2.2, 3.2], 'element3');
+      await client.vAdd('resp3-key', [2.0, 3.0, 4.0], 'element4');
+
+      const result = await client.vLinksWithScores('resp3-key', 'element1');
+
+      assert.ok(Array.isArray(result));
+
+      for (const layer of result) {
+        assert.equal(
+          typeof layer,
+          'object'
+        );
+      }
+
+      assert.ok(result.length >= 1, 'Should have at least layer 0');
+    },
+    {
+      ...GLOBAL.SERVERS.OPEN,
+      clientOptions: {
+        RESP: 3
+      }
+    }
+  );
+});

packages/client/lib/commands/VLINKS_WITHSCORES.ts

@@ -0,0 +1,42 @@
+import { BlobStringReply, Command, DoubleReply, MapReply } from '../RESP/types';
+import { transformDoubleReply } from './generic-transformers';
+import VLINKS from './VLINKS';
+
+
+function transformVLinksWithScoresReply(reply: any): Array<Record<string, DoubleReply>> {
+  const layers: Array<Record<string, DoubleReply>> = [];
+
+  for (const layer of reply) {
+    const obj: Record<string, DoubleReply> = Object.create(null);
+
+    // Each layer contains alternating element names and scores
+    for (let i = 0; i < layer.length; i += 2) {
+      const element = layer[i];
+      const score = transformDoubleReply[2](layer[i + 1]);
+      obj[element.toString()] = score;
+    }
+
+    layers.push(obj);
+  }
+
+  return layers;
+}
+
+export default {
+  IS_READ_ONLY: VLINKS.IS_READ_ONLY,
+  /**
+   * Get the connections for each layer of the HNSW graph with similarity scores
+   * @param args - Same parameters as the VLINKS command
+   * @see https://redis.io/commands/vlinks/
+   */
+  parseCommand(...args: Parameters<typeof VLINKS.parseCommand>) {
+    const parser = args[0];
+
+    VLINKS.parseCommand(...args);
+    parser.push('WITHSCORES');
+  },
+  transformReply: {
+    2: transformVLinksWithScoresReply,
+    3: undefined as unknown as () => Array<MapReply<BlobStringReply, DoubleReply>>
+  }
+} as const satisfies Command;

packages/client/lib/commands/VSIM_WITHSCORES.ts

@@ -4,7 +4,6 @@ import {
   Command,
   DoubleReply,
   MapReply,
-  TuplesReply,
   UnwrapReply
 } from '../RESP/types';
 import { transformDoubleReply } from './generic-transformers';

packages/client/lib/commands/index.ts

@@ -351,6 +351,7 @@ import VEMB from './VEMB';
 import VGETATTR from './VGETATTR';
 import VINFO from './VINFO';
 import VLINKS from './VLINKS';
+import VLINKS_WITHSCORES from './VLINKS_WITHSCORES';
 import VRANDMEMBER from './VRANDMEMBER';
 import VREM from './VREM';
 import VSETATTR from './VSETATTR';
@@ -1064,6 +1065,8 @@ export default {
   vInfo: VINFO,
   VLINKS,
   vLinks: VLINKS,
+  VLINKS_WITHSCORES,
+  vLinksWithScores: VLINKS_WITHSCORES,
   VRANDMEMBER,
   vRandMember: VRANDMEMBER,
   VREM,