Refactor Azure OpenAI integration: replace direct usage with new embedder and client utilities; remove unused web source constructor. (#2851)

Author: TheSonOfThomp

apps/chatbot-server/src/ingest/ingest.config.ts

@@ -1,12 +1,11 @@
 import {
   makeMongoDbEmbeddedContentStore,
   makeMongoDbPageStore,
-  makeOpenAiEmbedder,
 } from 'mongodb-rag-core';
-import { AzureOpenAI } from 'mongodb-rag-core/openai';
 import { Config, makeIngestMetaStore } from 'mongodb-rag-ingest';
 
 import { loadEnvVars } from '../utils/loadEnv';
+import { makeEmbedder } from '../utils/makeEmbedder';
 
 import { leafygreenGithubSourceConstructor } from './sources/github-leafygreen-ui';
 import { mongoDbChatbotFrameworkDocsDataSourceConstructor } from './sources/github-mdb-chatbot-framework';
@@ -16,24 +15,11 @@ import { webSourceConstructor } from './utils/webSourceConstructor';
 const {
   MONGODB_CONNECTION_URI,
   MONGODB_DATABASE_NAME,
-  AZURE_OPENAI_API_KEY,
-  AZURE_OPENAI_ENDPOINT,
   AZURE_OPENAI_DEPLOYMENT,
 } = loadEnvVars();
 
-const azureClient = new AzureOpenAI({
-  endpoint: AZURE_OPENAI_ENDPOINT,
-  apiKey: AZURE_OPENAI_API_KEY,
-  apiVersion: '2024-04-01-preview',
-  deployment: AZURE_OPENAI_DEPLOYMENT,
-});
-
 export default {
-  embedder: () =>
-    makeOpenAiEmbedder({
-      openAiClient: azureClient,
-      deployment: AZURE_OPENAI_DEPLOYMENT,
-    }),
+  embedder: () => makeEmbedder(),
   embeddedContentStore: () =>
     makeMongoDbEmbeddedContentStore({
       connectionUri: MONGODB_CONNECTION_URI,

apps/chatbot-server/src/init.ts

@@ -11,13 +11,13 @@ import {
   makeMongoDbConversationsService,
   makeMongoDbEmbeddedContentStore,
   makeOpenAiChatLlm,
-  makeOpenAiEmbedder,
   MongoDbEmbeddedContentStore,
 } from 'mongodb-rag-core';
 import { MongoClient } from 'mongodb-rag-core/mongodb';
 import { AzureOpenAI } from 'mongodb-rag-core/openai';
 
 import { loadEnvVars } from './utils/loadEnv';
+import { makeEmbedder } from './utils/makeEmbedder';
 
 export async function initChatBot(): Promise<{
   llm: ChatLlm;
@@ -37,13 +37,6 @@ export async function initChatBot(): Promise<{
     AZURE_OPENAI_CHAT_COMPLETION_MODEL,
   } = loadEnvVars();
 
-  const azureOpenAIEmbeddingClient = new AzureOpenAI({
-    endpoint: AZURE_OPENAI_ENDPOINT,
-    apiKey: AZURE_OPENAI_API_KEY,
-    apiVersion: '2024-04-01-preview',
-    deployment: AZURE_OPENAI_EMBEDDING_MODEL,
-  });
-
   const azureOpenAIChatClient = new AzureOpenAI({
     endpoint: AZURE_OPENAI_ENDPOINT,
     apiKey: AZURE_OPENAI_API_KEY,
@@ -62,11 +55,7 @@ export async function initChatBot(): Promise<{
 
   // Creates vector embeddings for user queries to find matching content
   // in the embeddedContentStore using Atlas Vector Search.
-  const embedder: Embedder = makeOpenAiEmbedder({
-    openAiClient: azureOpenAIEmbeddingClient,
-    deployment: AZURE_OPENAI_EMBEDDING_MODEL,
-    backoffOptions: {},
-  });
+  const embedder: Embedder = makeEmbedder();
 
   // MongoDB data source for the content used in RAG.
   // Generated with the Ingest CLI.
@@ -85,7 +74,7 @@ export async function initChatBot(): Promise<{
     store: embeddedContentStore,
     findNearestNeighborsOptions: {
       k: 5,
-      path: 'embeddings.text-embedding-3-small',
+      path: `embeddings.${AZURE_OPENAI_EMBEDDING_MODEL}`,
       indexName: VECTOR_SEARCH_INDEX_NAME,
       // Note: you may want to adjust the minScore depending
       // on the embedding model you use. We've found 0.9 works well

apps/chatbot-server/src/utils/makeEmbedder.ts

@@ -0,0 +1,21 @@
+import { makeOpenAiEmbedder } from 'mongodb-rag-core';
+
+import { loadEnvVars } from './loadEnv';
+import { makeEmbeddingClient } from './makeEmbeddingClient';
+
+/**
+ * Returns an consistent Azure OpenAI embedder
+ * used in the ingest process
+ * and in the chatbot server.
+ */
+export const makeEmbedder = () => {
+  const { AZURE_OPENAI_DEPLOYMENT } = loadEnvVars();
+
+  const azureClient = makeEmbeddingClient();
+
+  return makeOpenAiEmbedder({
+    openAiClient: azureClient,
+    deployment: AZURE_OPENAI_DEPLOYMENT,
+    backoffOptions: {},
+  });
+};

apps/chatbot-server/src/utils/makeEmbeddingClient.ts

@@ -0,0 +1,28 @@
+import { AzureOpenAI } from 'mongodb-rag-core/openai';
+
+import { loadEnvVars } from './loadEnv';
+
+/**
+ * Returns an consistent Azure OpenAI client
+ * that can be used to generate embeddings.
+ *
+ * Used in the ingest process
+ * and in the chatbot server.
+ */
+export const makeEmbeddingClient = () => {
+  // Load project environment variables
+  const {
+    AZURE_OPENAI_API_KEY,
+    AZURE_OPENAI_ENDPOINT,
+    AZURE_OPENAI_DEPLOYMENT,
+  } = loadEnvVars();
+
+  const azureClient = new AzureOpenAI({
+    endpoint: AZURE_OPENAI_ENDPOINT,
+    apiKey: AZURE_OPENAI_API_KEY,
+    apiVersion: '2024-04-01-preview',
+    deployment: AZURE_OPENAI_DEPLOYMENT,
+  });
+
+  return azureClient;
+};