Add support for configuring muvera encoding with multivectors

Author: tsmith023

src/collections/config/integration.test.ts

@@ -389,11 +389,7 @@ describe('Testing of the collection.config namespace', () => {
     ]);
   });
 
-  requireAtLeast(
-    1,
-    31,
-    0
-  )('Mutable named vectors', () => {
+  requireAtLeast(1, 31, 0)(describe)('Mutable named vectors', () => {
     it('should be able to add named vectors to a collection', async () => {
       const collectionName = 'TestCollectionConfigAddVector' as const;
       const collection = await client.collections.create({
@@ -715,4 +711,51 @@ describe('Testing of the collection.config namespace', () => {
       },
     });
   });
+
+  requireAtLeast(1, 31, 0)(it)(
+    'should be able to create and get a multi-vector collection with encoding',
+    async () => {
+      const collectionName = 'TestCollectionConfigCreateWithMuveraEncoding';
+      const collection = await client.collections.create({
+        name: collectionName,
+        vectorizers: weaviate.configure.vectorizer.none({
+          vectorIndexConfig: weaviate.configure.vectorIndex.hnsw({
+            multiVector: weaviate.configure.vectorIndex.multiVector.multiVector({
+              aggregation: 'maxSim',
+              encoding: weaviate.configure.vectorIndex.multiVector.encoding.muvera(),
+            }),
+          }),
+        }),
+      });
+      const config = await collection.config.get();
+      expect(config.name).toEqual(collectionName);
+
+      const indexConfig = config.vectorizers.default.indexConfig as VectorIndexConfigHNSW;
+      expect(indexConfig.multiVector).toBeDefined();
+      expect(indexConfig.multiVector?.aggregation).toEqual('maxSim');
+      expect(indexConfig.multiVector?.encoding).toBeDefined();
+    }
+  );
+
+  requireAtLeast(1, 31, 0)(it)(
+    'should be able to create and get a multi-vector collection without encoding',
+    async () => {
+      const collectionName = 'TestCollectionConfigCreateWithoutMuveraEncoding';
+      const collection = await client.collections.create({
+        name: collectionName,
+        vectorizers: weaviate.configure.vectorizer.none({
+          vectorIndexConfig: weaviate.configure.vectorIndex.hnsw({
+            multiVector: weaviate.configure.vectorIndex.multiVector.multiVector(),
+          }),
+        }),
+      });
+      const config = await collection.config.get();
+      expect(config.name).toEqual(collectionName);
+
+      const indexConfig = config.vectorizers.default.indexConfig as VectorIndexConfigHNSW;
+      expect(indexConfig.multiVector).toBeDefined();
+      expect(indexConfig.multiVector?.aggregation).toEqual('maxSim');
+      expect(indexConfig.multiVector?.encoding).toBeUndefined();
+    }
+  );
 });

src/collections/config/types/vectorIndex.ts

@@ -64,8 +64,18 @@ export type PQConfig = {
 
 export type MultiVectorConfig = {
   aggregation: 'maxSim' | string;
+  encoding?: MultiVectorEncodingConfig;
 };
 
+export type MuveraEncodingConfig = {
+  ksim?: number;
+  dprojections?: number;
+  repetitions?: number;
+  type: 'muvera';
+};
+
+export type MultiVectorEncodingConfig = MuveraEncodingConfig | Record<string, any>;
+
 export type PQEncoderConfig = {
   type: PQEncoderType;
   distribution: PQEncoderDistribution;

src/collections/config/utils.ts

@@ -18,7 +18,7 @@ import {
   WeaviateVectorsConfig,
 } from '../../openapi/types.js';
 import { DbVersionSupport } from '../../utils/dbVersion.js';
-import { QuantizerGuards, VectorIndexGuards } from '../configure/parsing.js';
+import { MultiVectorEncodingGuards, QuantizerGuards, VectorIndexGuards } from '../configure/parsing.js';
 import {
   PropertyConfigCreate,
   ReferenceConfigCreate,
@@ -39,6 +39,7 @@ import {
   ModuleConfig,
   MultiTenancyConfig,
   MultiVectorConfig,
+  MultiVectorEncodingConfig,
   PQConfig,
   PQEncoderConfig,
   PQEncoderDistribution,
@@ -147,18 +148,34 @@ export const parseVectorIndex = (module: ModuleConfig<VectorIndexType, VectorInd
     };
   }
 
-  let multiVector;
+  let multivector: any;
   if (VectorIndexGuards.isHNSW(module.config) && module.config.multiVector !== undefined) {
-    multiVector = {
-      ...module.config.multiVector,
+    multivector = {
+      aggregation: module.config.multiVector.aggregation,
       enabled: true,
     };
+    if (
+      module.config.multiVector.encoding !== undefined &&
+      MultiVectorEncodingGuards.isMuvera(module.config.multiVector.encoding)
+    ) {
+      multivector.muvera = {
+        enabled: true,
+        ksim: module.config.multiVector.encoding.ksim,
+        dprojections: module.config.multiVector.encoding.dprojections,
+        repetitions: module.config.multiVector.encoding.repetitions,
+      };
+    }
   }
 
-  const { quantizer, ...conf } = module.config as
+  const { quantizer, ...rest } = module.config as
     | VectorIndexConfigFlatCreate
     | VectorIndexConfigHNSWCreate
     | Record<string, any>;
+
+  const conf = {
+    ...rest,
+    multivector,
+  };
   if (quantizer === undefined) return conf;
   if (QuantizerGuards.isBQCreate(quantizer)) {
     const { type, ...quant } = quantizer;
@@ -476,6 +493,7 @@ class ConfigMapping {
     } else {
       quantizer = undefined;
     }
+
     return {
       cleanupIntervalSeconds: v.cleanupIntervalSeconds,
       distance: v.distance,
@@ -502,8 +520,25 @@ class ConfigMapping {
     if (v.enabled === false) return undefined;
     if (!exists<string>(v.aggregation))
       throw new WeaviateDeserializationError('Multi vector aggregation was not returned by Weaviate');
+    let encoding: MultiVectorEncodingConfig | undefined;
+    if (
+      exists<{
+        ksim: number;
+        dprojections: number;
+        repetitions: number;
+        enabled: boolean;
+      }>(v.muvera)
+    ) {
+      encoding = v.muvera.enabled
+        ? {
+            type: 'muvera',
+            ...v.muvera,
+          }
+        : undefined;
+    }
     return {
       aggregation: v.aggregation,
+      encoding,
     };
   }
   static bq(v?: Record<string, unknown>): BQConfig | undefined {

src/collections/configure/parsing.ts

@@ -1,3 +1,4 @@
+import { MuveraEncodingConfigCreate } from '../index.js';
 import {
   BQConfigCreate,
   BQConfigUpdate,
@@ -57,6 +58,12 @@ export class VectorIndexGuards {
   }
 }
 
+export class MultiVectorEncodingGuards {
+  static isMuvera(config?: Record<string, any>): config is MuveraEncodingConfigCreate {
+    return (config as { type: string })?.type === 'muvera';
+  }
+}
+
 export function parseWithDefault<D>(value: D | undefined, defaultValue: D): D {
   return value !== undefined ? value : defaultValue;
 }

src/collections/configure/types/vectorIndex.ts

@@ -2,6 +2,7 @@ import {
   BQConfig,
   ModuleConfig,
   MultiVectorConfig,
+  MuveraEncodingConfig,
   PQConfig,
   PQEncoderDistribution,
   PQEncoderType,
@@ -49,6 +50,8 @@ export type SQConfigUpdate = {
 
 export type MultiVectorConfigCreate = RecursivePartial<MultiVectorConfig>;
 
+export type MuveraEncodingConfigCreate = RecursivePartial<MuveraEncodingConfig>;
+
 export type VectorIndexConfigHNSWCreate = RecursivePartial<VectorIndexConfigHNSW>;
 
 export type VectorIndexConfigDynamicCreate = RecursivePartial<VectorIndexConfigDynamic>;

src/collections/configure/vectorIndex.ts

@@ -8,6 +8,7 @@ import {
   BQConfigCreate,
   BQConfigUpdate,
   MultiVectorConfigCreate,
+  MuveraEncodingConfigCreate,
   PQConfigCreate,
   PQConfigUpdate,
   SQConfigCreate,
@@ -101,15 +102,45 @@ const configure = {
    * Define the configuration for a multi-vector index.
    */
   multiVector: {
+    /**
+     * Specify the encoding configuration for a multi-vector index.
+     */
+    encoding: {
+      /**
+       * Create an object of type `MuveraEncodingConfigCreate` to be used when defining the encoding configuration of a multi-vector index using MUVERA.
+       *
+       * @param {number} [options.ksim] The number of nearest neighbors to consider for similarity. Default is undefined.
+       * @param {number} [options.dprojections] The number of projections to use. Default is undefined.
+       * @param {number} [options.repetitions] The number of repetitions to use. Default is undefined.
+       * @returns {MuveraEncodingConfigCreate} The object of type `MuveraEncodingConfigCreate`.
+       */
+      muvera: (options?: {
+        ksim?: number;
+        dprojections?: number;
+        repetitions?: number;
+      }): MuveraEncodingConfigCreate => {
+        return {
+          ksim: options?.ksim,
+          dprojections: options?.dprojections,
+          repetitions: options?.repetitions,
+          type: 'muvera',
+        };
+      },
+    },
     /**
      * Create an object of type `MultiVectorConfigCreate` to be used when defining the configuration of a multi-vector index.
      *
-     * @param {object} [options.aggregation] The aggregation method to use. Default is 'maxSim'.
+     * @param {string} [options.aggregation] The aggregation method to use. Default is 'maxSim'.
+     * @param {MultiVectorConfig['encoding']} [options.encoding] The encoding configuration for the multi-vector index. Default is undefined.
      * @returns {MultiVectorConfigCreate} The object of type `MultiVectorConfigCreate`.
      */
-    multiVector: (options?: { aggregation?: 'maxSim' | string }): MultiVectorConfigCreate => {
+    multiVector: (options?: {
+      aggregation?: 'maxSim' | string;
+      encoding?: MultiVectorConfigCreate['encoding'];
+    }): MultiVectorConfigCreate => {
       return {
         aggregation: options?.aggregation,
+        encoding: options?.encoding,
       };
     },
   },

src/collections/query/integration.test.ts

@@ -134,11 +134,7 @@ describe('Testing of the collection.query methods with a simple collection', ()
     expect(ret.objects[0].uuid).toEqual(id);
   });
 
-  requireAtLeast(
-    1,
-    31,
-    0
-  )('bm25 search operator (minimum_should_match)', () => {
+  requireAtLeast(1, 31, 0)(describe)('bm25 search operator (minimum_should_match)', () => {
     it('should query with bm25 + operator', async () => {
       const ret = await collection.query.bm25('carrot', {
         limit: 1,

src/roles/integration.test.ts

@@ -279,11 +279,7 @@ const testCases: TestCase[] = [
   },
 ];
 
-requireAtLeast(
-  1,
-  29,
-  0
-)('Integration testing of the roles namespace', () => {
+requireAtLeast(1, 29, 0)(describe)('Integration testing of the roles namespace', () => {
   let client: WeaviateClient;
 
   beforeAll(async () => {
@@ -317,11 +313,7 @@ requireAtLeast(
     expect(exists).toBeFalsy();
   });
 
-  requireAtLeast(
-    1,
-    30,
-    0
-  )('namespaced users', () => {
+  requireAtLeast(1, 30, 0)(describe)('namespaced users', () => {
     it('retrieves assigned users with/without namespace', async () => {
       await client.roles.create('landlord', {
         collection: 'Buildings',

src/users/integration.test.ts

@@ -3,11 +3,7 @@ import { requireAtLeast } from '../../test/version.js';
 import { WeaviateUserTypeDB } from '../openapi/types.js';
 import { GetUserOptions, UserDB } from './types.js';
 
-requireAtLeast(
-  1,
-  29,
-  0
-)('Integration testing of the users namespace', () => {
+requireAtLeast(1, 29, 0)(describe)('Integration testing of the users namespace', () => {
   const makeClient = (key: string) =>
     weaviate.connectToLocal({
       port: 8091,
@@ -62,11 +58,7 @@ requireAtLeast(
     expect(roles.test).toBeUndefined();
   });
 
-  requireAtLeast(
-    1,
-    30,
-    0
-  )('dynamic user management', () => {
+  requireAtLeast(1, 30, 0)(describe)('dynamic user management', () => {
     /** List dynamic DB users. */
     const listDBUsers = (c: WeaviateClient, opts?: GetUserOptions) =>
       c.users.db.listAll(opts).then((all) => all.filter((u) => u.userType == 'db_user'));
@@ -172,11 +164,7 @@ requireAtLeast(
       expect(roles.Permissioner.nodesPermissions).toHaveLength(1);
     });
 
-    requireAtLeast(
-      1,
-      30,
-      1
-    )('additional DUM features', () => {
+    requireAtLeast(1, 30, 1)(describe)('additional DUM features', () => {
       it('should be able to fetch additional user info', async () => {
         const admin = await makeClient('admin-key');
         const timKey = await admin.users.db.create('timely-tim');

test/version.ts

@@ -3,5 +3,7 @@ import { DbVersion } from '../src/utils/dbVersion';
 const version = DbVersion.fromString(`v${process.env.WEAVIATE_VERSION!}`);
 
 /** Run the suite / test only for Weaviate version above this. */
-export const requireAtLeast = (...semver: [...Parameters<DbVersion['isAtLeast']>]) =>
-  version.isAtLeast(...semver) ? describe : describe.skip;
+export const requireAtLeast =
+  (...semver: [...Parameters<DbVersion['isAtLeast']>]) =>
+  (type: jest.Describe | jest.It) =>
+    version.isAtLeast(...semver) ? type : type.skip;