diff --git a/src/indexes.ts b/src/indexes.ts index 3fb051111..d1fe24c0f 100644 --- a/src/indexes.ts +++ b/src/indexes.ts @@ -53,6 +53,7 @@ import { ProximityPrecision, Embedders, SearchCutoffMs, + SearchSimilarDocumentsParams, } from './types'; import { removeUndefinedFromObject } from './utils'; import { HttpRequests } from './http-requests'; @@ -177,6 +178,25 @@ class Index = Record> { ); } + /** + * Search for similar documents + * + * @param params - Parameters used to search for similar documents + * @returns Promise containing the search response + */ + async searchSimilarDocuments< + D extends Record = T, + S extends SearchParams = SearchParams, + >(params: SearchSimilarDocumentsParams): Promise> { + const url = `indexes/${this.uid}/similar`; + + return await this.httpRequest.post( + url, + removeUndefinedFromObject(params), + undefined, + ); + } + /// /// INDEX /// diff --git a/src/types/types.ts b/src/types/types.ts index c031a1cca..4d1789062 100644 --- a/src/types/types.ts +++ b/src/types/types.ts @@ -265,6 +265,18 @@ export type FieldDistribution = { [field: string]: number; }; +export type SearchSimilarDocumentsParams = { + id: string | number; + offset?: number; + limit?: number; + filter?: Filter; + embedder?: string; + attributesToRetrieve?: string[]; + showRankingScore?: boolean; + showRankingScoreDetails?: boolean; + rankingScoreThreshold?: number; +}; + /* ** Documents */ diff --git a/tests/embedders.test.ts b/tests/embedders.test.ts index d6691581f..60748756f 100644 --- a/tests/embedders.test.ts +++ b/tests/embedders.test.ts @@ -14,6 +14,39 @@ const index = { uid: 'movies_test', }; +const datasetSimilarSearch = [ + { + title: 'Shazam!', + release_year: 2019, + id: '287947', + _vectors: { manual: [0.8, 0.4, -0.5] }, + }, + { + title: 'Captain Marvel', + release_year: 2019, + id: '299537', + _vectors: { manual: [0.6, 0.8, -0.2] }, + }, + { + title: 'Escape Room', + release_year: 2019, + id: '522681', + _vectors: { manual: [0.1, 0.6, 0.8] }, + }, + { + title: 'How to Train Your Dragon: The Hidden World', + release_year: 2019, + id: '166428', + _vectors: { manual: [0.7, 0.7, -0.4] }, + }, + { + title: 'All Quiet on the Western Front', + release_year: 1930, + id: '143', + _vectors: { manual: [-0.5, 0.3, 0.85] }, + }, +]; + jest.setTimeout(100 * 1000); afterAll(() => { @@ -223,6 +256,38 @@ describe.each([{ permission: 'Master' }, { permission: 'Admin' }])( expect(response).toEqual(null); }); + + test(`${permission} key: search for similar documents`, async () => { + const client = await getClient(permission); + + const newEmbedder: Embedders = { + manual: { + source: 'userProvided', + dimensions: 3, + }, + }; + const { taskUid: updateEmbeddersTask }: EnqueuedTask = await client + .index(index.uid) + .updateEmbedders(newEmbedder); + + await client.waitForTask(updateEmbeddersTask); + + const { taskUid: documentAdditionTask } = await client + .index(index.uid) + .addDocuments(datasetSimilarSearch); + + await client.waitForTask(documentAdditionTask); + + const response = await client.index(index.uid).searchSimilarDocuments({ + id: '143', + }); + + expect(response).toHaveProperty('hits'); + expect(response.hits.length).toEqual(4); + expect(response).toHaveProperty('offset', 0); + expect(response).toHaveProperty('limit', 20); + expect(response).toHaveProperty('estimatedTotalHits', 4); + }); }, );