From 2d8a62d0596127e9459da30adf7128ab7712f8bb Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sun, 17 Dec 2023 02:56:47 +0200
Subject: [PATCH] Refactor openai vectors

---
 src/openai-vectors.js | 42 +++++++++++++++++++++++++++++++-----------
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/src/openai-vectors.js b/src/openai-vectors.js
index 40c54ae2f..3b19a4c96 100644
--- a/src/openai-vectors.js
+++ b/src/openai-vectors.js
@@ -1,21 +1,41 @@
 const fetch = require('node-fetch').default;
 const { SECRET_KEYS, readSecret } = require('./endpoints/secrets');
 
+const SOURCES = {
+    'mistral': {
+        secretKey: SECRET_KEYS.MISTRAL,
+        url: 'api.mistral.ai',
+        model: 'mistral-embed',
+    },
+    'openai': {
+        secretKey: SECRET_KEYS.OPENAI,
+        url: 'api.openai.com',
+        model: 'text-embedding-ada-002',
+    },
+};
+
 /**
  * Gets the vector for the given text from an OpenAI compatible endpoint.
  * @param {string} text - The text to get the vector for
+ * @param {string} source - The source of the vector
  * @returns {Promise<number[]>} - The vector for the text
  */
 async function getOpenAIVector(text, source) {
-    const isMistral = source === 'mistral';
-    const key = readSecret(isMistral ? SECRET_KEYS.MISTRALAI : SECRET_KEYS.OPENAI);
+    const config = SOURCES[source];
 
-    if (!key) {
-        console.log('No OpenAI key found');
-        throw new Error('No OpenAI key found');
+    if (!config) {
+        console.log('Unknown source', source);
+        throw new Error('Unknown source');
     }
 
-    const url = isMistral ? 'api.mistral.ai' : 'api.openai.com';
+    const key = readSecret(config.secretKey);
+
+    if (!key) {
+        console.log('No API key found');
+        throw new Error('No API key found');
+    }
+
+    const url = config.url;
     const response = await fetch(`https://${url}/v1/embeddings`, {
         method: 'POST',
         headers: {
@@ -24,22 +44,22 @@ async function getOpenAIVector(text, source) {
         },
         body: JSON.stringify({
             input: text,
-            model: isMistral ? 'mistral-embed' : 'text-embedding-ada-002',
+            model: config.model,
         }),
     });
 
     if (!response.ok) {
         const text = await response.text();
-        console.log('OpenAI request failed', response.statusText, text);
-        throw new Error('OpenAI request failed');
+        console.log('API request failed', response.statusText, text);
+        throw new Error('API request failed');
     }
 
     const data = await response.json();
     const vector = data?.data[0]?.embedding;
 
     if (!Array.isArray(vector)) {
-        console.log('OpenAI response was not an array');
-        throw new Error('OpenAI response was not an array');
+        console.log('API response was not an array');
+        throw new Error('API response was not an array');
     }
 
     return vector;