From 8cb2e99d128732ac5225b227345d5d087fdebe4b Mon Sep 17 00:00:00 2001
From: Tomas Pilar <thomas7pilar@gmail.com>
Date: Wed, 19 Jun 2024 08:35:32 +0200
Subject: [PATCH] feat(similarity): add sentence similarity service (#109)

Signed-off-by: Tomas Pilar <tomas.pilar@ibm.com>
---
 examples/shared/constants.ts                  |  1 +
 examples/similarity.ts                        | 16 ++++++++
 package.json                                  |  2 +-
 src/schema.ts                                 |  9 +++++
 .../text/TextSentenceSimilarityService.ts     | 37 +++++++++++++++++++
 src/services/text/TextService.ts              | 11 ++++++
 6 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 examples/similarity.ts
 create mode 100644 src/services/text/TextSentenceSimilarityService.ts

diff --git a/examples/shared/constants.ts b/examples/shared/constants.ts
index b336f0d..ce0714e 100644
--- a/examples/shared/constants.ts
+++ b/examples/shared/constants.ts
@@ -1,2 +1,3 @@
 export const MODEL = 'google/flan-ul2';
 export const CHAT_MODEL = 'meta-llama/llama-3-70b-instruct';
+export const EMBEDDING_MODEL = 'sentence-transformers/all-minilm-l6-v2';
diff --git a/examples/similarity.ts b/examples/similarity.ts
new file mode 100644
index 0000000..174f055
--- /dev/null
+++ b/examples/similarity.ts
@@ -0,0 +1,16 @@
+import { Client } from '../src/index.js';
+
+import { EMBEDDING_MODEL } from './shared/constants.js';
+
+const client = new Client({
+  apiKey: process.env.GENAI_API_KEY,
+});
+
+{
+  const output = await client.text.experimental.sentenceSimilarity.create({
+    model_id: EMBEDDING_MODEL,
+    source_sentence: 'Good morning',
+    sentences: ['How are you?', 'Get lost!'],
+  });
+  console.log(output);
+}
diff --git a/package.json b/package.json
index e27bd7c..80a2a10 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@ibm-generative-ai/node-sdk",
-  "version": "3.0.0",
+  "version": "3.1.0",
   "description": "IBM Generative AI Node.js SDK (Tech Preview)",
   "keywords": [
     "ai",
diff --git a/src/schema.ts b/src/schema.ts
index 0a383d3..7fd5f01 100644
--- a/src/schema.ts
+++ b/src/schema.ts
@@ -55,6 +55,15 @@ export type TextEmbeddingCreateOutput = OutputWrapper<
   ApiClientResponse<'POST', '/v2/text/embeddings'>['data']
 >;
 
+// TextSentenceSimilarityService
+
+export type TextSentenceSimilarityCreateInput = InputBodyWrapper<
+  ApiClientOptions<'POST', '/v2/beta/text/sentence-similarity'>['body']
+>;
+export type TextSentenceSimilarityCreateOutput = OutputWrapper<
+  ApiClientResponse<'POST', '/v2/beta/text/sentence-similarity'>['data']
+>;
+
 // ModelService
 
 export type ModelServiceListInput = InputQueryWrapper<
diff --git a/src/services/text/TextSentenceSimilarityService.ts b/src/services/text/TextSentenceSimilarityService.ts
new file mode 100644
index 0000000..549e9fb
--- /dev/null
+++ b/src/services/text/TextSentenceSimilarityService.ts
@@ -0,0 +1,37 @@
+import { Options } from '../../client.js';
+import { clientErrorWrapper } from '../../utils/errors.js';
+import {
+  TextSentenceSimilarityCreateInput,
+  TextSentenceSimilarityCreateOutput,
+} from '../../schema.js';
+import { ApiClient } from '../../api/client.js';
+import { SteamingApiClient } from '../../api/streaming-client.js';
+import { ConcurrencyLimiter } from '../../utils/concurrency.js';
+import { BaseService } from '../BaseService.js';
+
+export class TextSentenceSimilarityService extends BaseService {
+  constructor(
+    protected readonly _client: ApiClient,
+    protected readonly _streamingClient: SteamingApiClient,
+    protected readonly _limiter: ConcurrencyLimiter,
+  ) {
+    super(_client, _streamingClient);
+  }
+
+  create(
+    input: TextSentenceSimilarityCreateInput,
+    opts?: Options,
+  ): Promise<TextSentenceSimilarityCreateOutput> {
+    return this._limiter.execute(
+      () =>
+        clientErrorWrapper(
+          this._client.POST('/v2/beta/text/sentence-similarity', {
+            ...opts,
+            params: { query: { version: '2023-11-22' } },
+            body: input,
+          }),
+        ),
+      { signal: opts?.signal },
+    );
+  }
+}
diff --git a/src/services/text/TextService.ts b/src/services/text/TextService.ts
index e9bad6e..a8919f5 100644
--- a/src/services/text/TextService.ts
+++ b/src/services/text/TextService.ts
@@ -8,12 +8,16 @@ import { TextGenerationService } from './TextGenerationService.js';
 import { TextTokenizationService } from './TextTokenizationService.js';
 import { TextEmbeddingService } from './TextEmbeddingService.js';
 import { TextChatService } from './TextChatService.js';
+import { TextSentenceSimilarityService } from './TextSentenceSimilarityService.js';
 
 export class TextService extends BaseService {
   public readonly generation: TextGenerationService;
   public readonly tokenization: TextTokenizationService;
   public readonly embedding: TextEmbeddingService;
   public readonly chat: TextChatService;
+  public readonly experimental: {
+    sentenceSimilarity: TextSentenceSimilarityService;
+  };
 
   constructor(client: ApiClient, streamingClient: SteamingApiClient) {
     super(client, streamingClient);
@@ -51,5 +55,12 @@ export class TextService extends BaseService {
       embeddingLimiter,
     );
     this.chat = new TextChatService(client, streamingClient, generationLimiter);
+    this.experimental = {
+      sentenceSimilarity: new TextSentenceSimilarityService(
+        client,
+        streamingClient,
+        embeddingLimiter,
+      ),
+    };
   }
 }