diff --git a/src/collections/query/index.ts b/src/collections/query/index.ts index f1b23531..8bc87ba6 100644 --- a/src/collections/query/index.ts +++ b/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'; diff --git a/src/collections/query/types.ts b/src/collections/query/types.ts index 62e5b95c..0ba27346 100644 --- a/src/collections/query/types.ts +++ b/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 = { distance?: number; /** Specify which vector to search on if using named vectors. */ targetVector?: TargetVectorInputType; + /** Specify diversity selection algorithm. */ + diversity?: DiversityConfig; }; /** Base options for the near search queries. */ diff --git a/src/collections/query/utils.ts b/src/collections/query/utils.ts index a46334da..885d4abd 100644 --- a/src/collections/query/utils.ts +++ b/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): DiversityConfig { + return { + ...args, + type: 'mmr', + }; + } +} diff --git a/src/collections/serialize/index.ts b/src/collections/serialize/index.ts index 04413230..7d806841 100644 --- a/src/collections/serialize/index.ts +++ b/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 = (groupBy?: GroupByOptions): 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; 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; }): Promise => 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, }); diff --git a/src/collections/types/query.ts b/src/collections/types/query.ts index a257cc99..930a813b 100644 --- a/src/collections/types/query.ts +++ b/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'; + } +} diff --git a/test/collections/query/integration.test.ts b/test/collections/query/integration.test.ts index e2cb3520..e442f6bc 100644 --- a/test/collections/query/integration.test.ts +++ b/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,