feat: get similar documents (#1677)

Author: mdubus

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<T extends Record<string, any> = Record<string, any>> {
     );
   }
 
+  /**
+   * Search for similar documents
+   *
+   * @param params - Parameters used to search for similar documents
+   * @returns Promise containing the search response
+   */
+  async searchSimilarDocuments<
+    D extends Record<string, any> = T,
+    S extends SearchParams = SearchParams,
+  >(params: SearchSimilarDocumentsParams): Promise<SearchResponse<D, S>> {
+    const url = `indexes/${this.uid}/similar`;
+
+    return await this.httpRequest.post(
+      url,
+      removeUndefinedFromObject(params),
+      undefined,
+    );
+  }
+
   ///
   /// INDEX
   ///

src/types/types.ts

@@ -267,6 +267,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
  */

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);
+    });
   },
 );