From d44d88ac32a3b2bf5e032f8a7c35f231f07e3b5f Mon Sep 17 00:00:00 2001
From: Kikuo Emoto <kemoto@codemonger.io>
Date: Wed, 11 Oct 2023 21:54:16 +0900
Subject: [PATCH] feat(viewer): add createEmbedding

- Introduces a new submodule `utils/embeddings` that provides utilities
  to access OpenAI API to deal with embeddings. It provides
  `createEmbedding` function that calculates an embedding vector for a
  given text. It obtains the OpenAI API key from the Parameter Store on
  AWS Systems Manager.

issue codemonger-io/mumble#26
---
 cdk/viewer/src/utils/embeddings.ts | 66 ++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 cdk/viewer/src/utils/embeddings.ts

diff --git a/cdk/viewer/src/utils/embeddings.ts b/cdk/viewer/src/utils/embeddings.ts
new file mode 100644
index 0000000..5b20e6f
--- /dev/null
+++ b/cdk/viewer/src/utils/embeddings.ts
@@ -0,0 +1,66 @@
+import { getOpenAiApiKeyParameter } from '~/utils/parameters';
+
+const EMBEDDING_ENDPOINT = 'https://api.openai.com/v1/embeddings';
+const EMBEDDING_MODEL = 'text-embedding-ada-002';
+
+// OpenAI API key.
+//
+// You have to configure the environment variable
+// `OPENAI_API_KEY_PARAMETER_PATH` that specifies the path to the parameter that
+// stores the OpenAI API key for embedding calculation in the Parameter Store on
+// AWS Systems Manager.
+//
+// This variable caches the API key.
+let openaiApiKey: string | null = null;
+
+/**
+ * Creates an embedding vector of a given text.
+ *
+ * @remarks
+ *
+ * You have to configure the environment variable
+ * `OPENAI_API_KEY_PARAMETER_PATH`.
+ */
+export async function createEmbedding(text: string): Promise<[number]> {
+  if (openaiApiKey == null) {
+    openaiApiKey = await getOpenAiApiKeyParameter();
+  }
+  const res = await fetch(EMBEDDING_ENDPOINT, {
+    method: 'POST',
+    headers: {
+      Authorization: `Bearer ${openaiApiKey}`,
+      'Content-Type': 'application/json',
+    },
+    body: JSON.stringify({
+      model: EMBEDDING_MODEL,
+      input: text,
+      user: 'mumble',
+    }),
+  });
+  if (res.status !== 200) {
+    throw new Error(
+      `embedding request failed with ${res.status}: ${await res.text()}`,
+    );
+  }
+  const result = await res.json();
+  if (!Array.isArray(result.data)) {
+    throw new Error(
+      'unexpected response from OpenAI API: data is not an array',
+    );
+  }
+  if (result.data.length === 0) {
+    throw new Error('unexpected response from OpenAI API: data is empty');
+  }
+  if (result.data.length > 1) {
+    console.warn(
+      'unexpected response from OpenAI API: more than one embeddings',
+    );
+  }
+  const embedding = result.data[0];
+  if (!Array.isArray(embedding.embedding)) {
+    throw new Error(
+      'unexpected response from OpenAI API: embedding is not array',
+    );
+  }
+  return embedding.embedding;
+}