feat(query): Support diversity selection for near- queries

src/collections/query/index.ts

@@ -356,4 +356,4 @@ export {
   SingleVectorType,
 } from './types.js';
 
-export { Bm25Operator } from './utils.js';
+export { Bm25Operator, Diversity } from './utils.js';

src/collections/query/types.ts

@@ -1,5 +1,5 @@
 import { FilterValue } from '../filters/index.js';
-import { CallOptions, MultiTargetVectorJoin, ReturnVectors } from '../index.js';
+import { CallOptions, DiversityConfig, MultiTargetVectorJoin, ReturnVectors } from '../index.js';
 import { Sorting } from '../sort/classes.js';
 import {
   GroupByOptions,
@@ -158,6 +158,8 @@ export type NearSearchOptions<V> = {
   distance?: number;
   /** Specify which vector to search on if using named vectors. */
   targetVector?: TargetVectorInputType<V>;
+  /** Specify diversity selection algorithm. */
+  diversity?: DiversityConfig;
 };
 
 /** Base options for the near search queries. */

src/collections/query/utils.ts

@@ -1,4 +1,4 @@
-import { MultiTargetVectorJoin, Vectors } from '../index.js';
+import { DiversityConfig, MMR, MultiTargetVectorJoin, Vectors } from '../index.js';
 import {
   Bm25OperatorOptions,
   Bm25OperatorOr,
@@ -76,3 +76,13 @@ export class Bm25Operator {
     return { ...opts, operator: 'Or' };
   }
 }
+
+export class Diversity {
+  /** Use Maximal Marginal Relevance diversity selection. */
+  static mmr(args?: Omit<MMR, 'type'>): DiversityConfig {
+    return {
+      ...args,
+      type: 'mmr',
+    };
+  }
+}

src/collections/serialize/index.ts

@@ -17,6 +17,8 @@ import {
   NearVideoSearch,
   SearchOperatorOptions,
   SearchOperatorOptions_Operator,
+  Selection,
+  Selection_MMR,
   Targets,
   VectorForTarget,
   WeightsForTarget,
@@ -111,6 +113,8 @@ import {
   AggregateHybridOptions,
   AggregateNearOptions,
   BatchReference,
+  DiversityConfig,
+  DiversityGuards,
   GenerativeConfigRuntime,
   GroupByAggregate,
   GroupedTask,
@@ -630,6 +634,18 @@ class Search {
     });
   };
 
+  public static selection = (diversity?: DiversityConfig): Selection | undefined => {
+    if (DiversityGuards.isMMR(diversity)) {
+      return {
+        mmr: Selection_MMR.fromPartial({
+          balance: diversity.balance,
+          limit: diversity.limit,
+        }),
+      };
+    }
+    return undefined;
+  };
+
   public static groupBy = <T>(groupBy?: GroupByOptions<T>): GroupBy => {
     return GroupBy.fromPartial({
       path: groupBy?.property ? [groupBy.property as string] : undefined,
@@ -1162,6 +1178,7 @@ export class Serialize {
       audio: args.audio,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targetVectors,
       targets,
     });
@@ -1175,6 +1192,7 @@ export class Serialize {
       depth: args.depth,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targetVectors,
       targets,
     });
@@ -1188,6 +1206,7 @@ export class Serialize {
       image: args.image,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targetVectors,
       targets,
     });
@@ -1199,6 +1218,7 @@ export class Serialize {
       imu: args.imu,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targetVectors,
       targets,
     });
@@ -1210,6 +1230,7 @@ export class Serialize {
       id: args.id,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targetVectors,
       targets,
     });
@@ -1220,6 +1241,7 @@ export class Serialize {
     targetVector?: TargetVectorInputType<V>;
     certainty?: number;
     distance?: number;
+    diversity?: DiversityConfig;
     moveAway?: { concepts?: string[]; force?: number; objects?: string[] };
     moveTo?: { concepts?: string[]; force?: number; objects?: string[] };
   }) => {
@@ -1228,6 +1250,7 @@ export class Serialize {
       query: typeof args.query === 'string' ? [args.query] : args.query,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targets,
       targetVectors,
       moveAway: args.moveAway
@@ -1255,15 +1278,12 @@ export class Serialize {
       thermal: args.thermal,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targetVectors,
       targets,
     });
   };
 
-  private static vectorToBuffer = (vector: number[]): ArrayBufferLike => {
-    return new Float32Array(vector).buffer;
-  };
-
   private static vectorToBytes = (vector: number[]): Uint8Array => {
     const uint32len = 4;
     const dv = new DataView(new ArrayBuffer(vector.length * uint32len));
@@ -1307,11 +1327,13 @@ export class Serialize {
     supportsVectors: boolean;
     certainty?: number;
     distance?: number;
+    diversity?: DiversityConfig;
     targetVector?: TargetVectorInputType<V>;
   }): Promise<NearVector> =>
     NearVector.fromPartial({
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       ...(await Serialize.vectors({
         ...args,
         argumentName: 'nearVector',
@@ -1539,6 +1561,7 @@ export class Serialize {
       video: args.video,
       certainty: args.certainty,
       distance: args.distance,
+      selection: Search.selection(args.diversity),
       targetVectors,
       targets,
     });

src/collections/types/query.ts

@@ -189,3 +189,17 @@ export type SortBy = {
   property: string;
   ascending?: boolean;
 };
+
+export type DiversityConfig = MMR;
+
+export type MMR = {
+  type: 'mmr';
+  limit?: number;
+  balance?: number;
+};
+
+export class DiversityGuards {
+  static isMMR(diversity?: DiversityConfig): diversity is MMR {
+    return (diversity as MMR)?.type === 'mmr';
+  }
+}

test/collections/query/integration.test.ts

@@ -7,6 +7,7 @@ import weaviate, {
   Bm25Operator,
   Collection,
   CrossReference,
+  Diversity,
   GroupByOptions,
   Reference,
   WeaviateClient,
@@ -182,6 +183,21 @@ describe('Testing of the collection.query methods with a simple collection', ()
     });
   });
 
+  requireAtLeast(1, 37, 8).describe('near text query with diversity selection', () => {
+    it('using Diversity factory', async () => {
+      const ret = await collection.query.nearText(['carrot'], {
+        diversity: Diversity.mmr({ limit: 1 }),
+      });
+      expect(ret.objects).toHaveLength(1);
+    });
+    it('using manual object construction', async () => {
+      const ret = await collection.query.nearText(['carrot'], {
+        diversity: { type: 'mmr', limit: 1 },
+      });
+      expect(ret.objects).toHaveLength(1);
+    });
+  });
+
   it('should query with hybrid and vector', async () => {
     const ret = await collection.query.hybrid('carrot', {
       limit: 1,