feat: Adjust embedding types and access

SAP · Oct 14, 2024 · 09c16b4 · 09c16b4
1 parent 461bf37
commit 09c16b4
Show file tree

Hide file tree

Showing 7 changed files with 92 additions and 47 deletions.
diff --git a/.changeset/grumpy-apes-beam.md b/.changeset/grumpy-apes-beam.md
@@ -0,0 +1,5 @@
+---
+'@sap-ai-sdk/foundation-models': minor
+---
+
+[New Functionality] Add convenience method to access all embeddings in an Azure OpenAI response (`AzureOpenAiEmbeddingResponse`).
diff --git a/.changeset/many-zebras-repair.md b/.changeset/many-zebras-repair.md
@@ -0,0 +1,5 @@
+---
+'@sap-ai-sdk/foundation-models': minor
+---
+
+[Compatibility Note] Adjust `AzureOpenAiEmbeddingOutput` type to include multiple embedding responses as opposed to one.
diff --git a/packages/foundation-models/src/azure-openai/azure-openai-embedding-response.test.ts b/packages/foundation-models/src/azure-openai/azure-openai-embedding-response.test.ts
@@ -2,35 +2,57 @@ import { createLogger } from '@sap-cloud-sdk/util';
 import { jest } from '@jest/globals';
 import { parseMockResponse } from '../../../../test-util/mock-http.js';
 import { AzureOpenAiEmbeddingResponse } from './azure-openai-embedding-response.js';
+import type { HttpResponse } from '@sap-cloud-sdk/http-client';
+import type { AzureOpenAiEmbeddingOutput } from './azure-openai-embedding-types.js';
 
 describe('Azure OpenAI embedding response', () => {
-  const mockResponse = parseMockResponse<AzureOpenAiEmbeddingResponse>(
-    'foundation-models',
-    'azure-openai-embeddings-success-response.json'
-  );
-  const rawResponse = {
-    data: mockResponse,
-    status: 200,
-    headers: {},
-    request: {}
-  };
-  const embeddingResponse = new AzureOpenAiEmbeddingResponse(rawResponse);
+  let embeddingResponse: AzureOpenAiEmbeddingResponse;
+  let rawResponse: HttpResponse;
+  let mockedData: AzureOpenAiEmbeddingOutput;
+  beforeAll(async () => {
+    mockedData = await parseMockResponse<AzureOpenAiEmbeddingOutput>(
+      'foundation-models',
+      'azure-openai-embeddings-success-response.json'
+    );
+
+    rawResponse = {
+      data: mockedData,
+      status: 200,
+      headers: {},
+      request: {}
+    };
+    embeddingResponse = new AzureOpenAiEmbeddingResponse(rawResponse);
+  });
 
   it('should return the embedding response', () => {
-    expect(embeddingResponse.data).toStrictEqual(mockResponse);
+    expect(embeddingResponse.data).toStrictEqual(mockedData);
   });
 
   it('should return raw response', () => {
     expect(embeddingResponse.rawResponse).toBe(rawResponse);
   });
+
+  it('should return the first embedding', () => {
+    expect(embeddingResponse.getEmbedding()).toEqual(
+      mockedData.data[0].embedding
+    );
+  });
+
   it('should return undefined when convenience function is called with incorrect index', () => {
     const logger = createLogger({
       package: 'foundation-models',
       messageContext: 'azure-openai-embedding-response'
     });
     const errorSpy = jest.spyOn(logger, 'error');
-    expect(embeddingResponse.getEmbedding(1)).toBeUndefined();
-    expect(errorSpy).toHaveBeenCalledWith('Data index 1 is out of bounds.');
+    expect(embeddingResponse.getEmbedding(2)).toBeUndefined();
+    expect(errorSpy).toHaveBeenCalledWith('Data index 2 is out of bounds.');
     expect(errorSpy).toHaveBeenCalledTimes(1);
   });
+
+  it('should return all embeddings', () => {
+    expect(embeddingResponse.getEmbeddings()).toEqual([
+      mockedData.data[0].embedding,
+      mockedData.data[1]?.embedding
+    ]);
+  });
 });
diff --git a/packages/foundation-models/src/azure-openai/azure-openai-embedding-response.ts b/packages/foundation-models/src/azure-openai/azure-openai-embedding-response.ts
@@ -30,6 +30,14 @@ export class AzureOpenAiEmbeddingResponse {
     return this.data.data[dataIndex]?.embedding;
   }
 
+  /**
+   * Parses the Azure OpenAI response and returns all embeddings.
+   * @returns The embedding vectors.
+   */
+  getEmbeddings(): number[][] | undefined {
+    return this.data.data.map(({ embedding }) => embedding);
+  }
+
   private logInvalidDataIndex(dataIndex: number): void {
     if (dataIndex < 0 || dataIndex >= this.data.data.length) {
       logger.error(`Data index ${dataIndex} is out of bounds.`);

diff --git a/packages/foundation-models/src/azure-openai/azure-openai-embedding-types.ts b/packages/foundation-models/src/azure-openai/azure-openai-embedding-types.ts
@@ -27,22 +27,20 @@ export interface AzureOpenAiEmbeddingOutput {
   /**
    * Array of result candidates.
    */
-  data: [
-    {
-      /**
-       * Embedding object.
-       */
-      object: 'embedding';
-      /**
-       * Array of size `1536` (Azure OpenAI's embedding size) containing embedding vector.
-       */
-      embedding: number[];
-      /**
-       * Index of choice.
-       */
-      index: number;
-    }
-  ];
+  data: {
+    /**
+     * Embedding object.
+     */
+    object: 'embedding';
+    /**
+     * Array of size `1536` (Azure OpenAI's embedding size) containing embedding vector.
+     */
+    embedding: number[];
+    /**
+     * Index of choice.
+     */
+    index: number;
+  }[];
   /**
    * Token Usage.
    */

diff --git a/test-util/data/foundation-models/azure-openai-embeddings-success-response.json b/test-util/data/foundation-models/azure-openai-embeddings-success-response.json
@@ -8,6 +8,15 @@
       ],
       "index": 0,
       "object": "embedding"
+    },
+    {
+      "embedding": [
+        -0.011352594, -0.006521842, 0.0059352037, -0.021066532, -0.0062957075,
+        0.024606025, -0.02332132, 0.016792923, 0.0029676019, -0.04000937,
+        0.0143283885, 0.025274595, -0.0006419426, -0.0014493892, -0.006269489
+      ],
+      "index": 1,
+      "object": "embedding"
     }
   ],
   "model": "ada",

diff --git a/test-util/mock-http.ts b/test-util/mock-http.ts
@@ -1,18 +1,18 @@
-import fs from 'fs';
-import path from 'path';
-import { fileURLToPath } from 'url';
-import {
-  DestinationAuthToken,
-  HttpDestination,
-  ServiceCredentials
-} from '@sap-cloud-sdk/connectivity';
+import { readFile } from 'node:fs/promises';
+import path from 'node:path';
+import { fileURLToPath } from 'node:url';
 import nock from 'nock';
 import { type EndpointOptions } from '@sap-ai-sdk/core';
 import {
   type FoundationModel,
   type DeploymentResolutionOptions
 } from '@sap-ai-sdk/ai-api/internal.js';
 import { dummyToken } from './mock-jwt.js';
+import type {
+  DestinationAuthToken,
+  HttpDestination,
+  ServiceCredentials
+} from '@sap-cloud-sdk/connectivity';
 
 // Get the directory of this file
 const __filename = fileURLToPath(import.meta.url);
@@ -121,7 +121,7 @@ export function mockDeploymentsList(
 ): nock.Scope {
   const nockOpts = {
     reqheaders: {
-      'ai-resource-group': opts?.resourceGroup  ?? 'default',
+      'ai-resource-group': opts?.resourceGroup ?? 'default'
     }
   };
   const query = {
@@ -143,14 +143,12 @@ export function mockDeploymentsList(
 /**
  * @internal
  */
-export function parseMockResponse<T>(client: string, fileName: string): T {
-  const fileContent = fs.readFileSync(
-    path.join(
-      __dirname,
-      'data',
-      client,
-      fileName
-    ),
+export async function parseMockResponse<T>(
+  client: string,
+  fileName: string
+): Promise<T> {
+  const fileContent = await readFile(
+    path.join(__dirname, 'data', client, fileName),
     'utf-8'
   );
   return JSON.parse(fileContent);