Merge ee9d1d2471 into 1fcd64ad42

src/lib/providers/index.ts

@@ -42,5 +42,7 @@ export const getAvailableEmbeddingModelProviders = async () => {
     }
   }
 
+  models['custom_openai'] = {};
+
   return models;
 };

src/websocket/connectionManager.ts

@@ -8,7 +8,7 @@ import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { Embeddings } from '@langchain/core/embeddings';
 import type { IncomingMessage } from 'http';
 import logger from '../utils/logger';
-import { ChatOpenAI } from '@langchain/openai';
+import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai';
 
 export const handleConnection = async (
   ws: WebSocket,
@@ -61,11 +61,20 @@ export const handleConnection = async (
 
     if (
       embeddingModelProviders[embeddingModelProvider] &&
-      embeddingModelProviders[embeddingModelProvider][embeddingModel]
+      embeddingModelProviders[embeddingModelProvider][embeddingModel] &&
+      embeddingModelProvider != 'custom_openai'
     ) {
       embeddings = embeddingModelProviders[embeddingModelProvider][
         embeddingModel
       ] as Embeddings | undefined;
+    } else if (embeddingModelProvider == 'custom_openai') {
+      embeddings = new OpenAIEmbeddings({
+        modelName: embeddingModel,
+        openAIApiKey: searchParams.get('openAIApiKey'),
+        configuration: {
+          baseURL: searchParams.get('openAIBaseURL'),
+        },
+      }) as unknown as Embeddings
     }
 
     if (!llm || !embeddings) {

ui/components/ChatWindow.tsx

@@ -130,6 +130,7 @@ const useSocket = (
 
           if (
             embeddingModelProvider &&
+            embeddingModelProvider != 'custom_openai' &&
             !embeddingModelProviders[embeddingModelProvider][embeddingModel]
           ) {
             embeddingModel = Object.keys(
@@ -159,6 +160,17 @@ const useSocket = (
         searchParams.append('embeddingModel', embeddingModel!);
         searchParams.append('embeddingModelProvider', embeddingModelProvider);
 
+        if (embeddingModelProvider === 'custom_openai') {
+          searchParams.append(
+            'openAIApiKey',
+            localStorage.getItem('openAIApiKey')!,
+          );
+          searchParams.append(
+            'openAIBaseURL',
+            localStorage.getItem('openAIBaseURL')!,
+          );
+        }
+
         wsURL.search = searchParams.toString();
 
         const ws = new WebSocket(wsURL.toString());

ui/components/SettingsDialog.tsx

@@ -9,7 +9,7 @@ import React, {
 } from 'react';
 import ThemeSwitcher from './theme/Switcher';
 
-interface InputProps extends React.InputHTMLAttributes<HTMLInputElement> {}
+interface InputProps extends React.InputHTMLAttributes<HTMLInputElement> { }
 
 const Input = ({ className, ...restProps }: InputProps) => {
   return (
@@ -258,30 +258,30 @@ const SettingsDialog = ({
                             options={(() => {
                               const chatModelProvider =
                                 config.chatModelProviders[
-                                  selectedChatModelProvider
+                                selectedChatModelProvider
                                 ];
 
                               return chatModelProvider
                                 ? chatModelProvider.length > 0
                                   ? chatModelProvider.map((model) => ({
-                                      value: model,
+                                    value: model,
-                                      label: model,
+                                    label: model,
-                                    }))
+                                  }))
                                   : [
-                                      {
-                                        value: '',
-                                        label: 'No models available',
-                                        disabled: true,
-                                      },
-                                    ]
-                                : [
                                     {
                                       value: '',
-                                      label:
+                                      label: 'No models available',
-                                        'Invalid provider, please check backend logs',
                                       disabled: true,
                                     },
-                                  ];
+                                  ]
+                                : [
+                                  {
+                                    value: '',
+                                    label:
+                                      'Invalid provider, please check backend logs',
+                                    disabled: true,
+                                  },
+                                ];
                             })()}
                           />
                         </div>
@@ -340,9 +340,13 @@ const SettingsDialog = ({
                           value={selectedEmbeddingModelProvider ?? undefined}
                           onChange={(e) => {
                             setSelectedEmbeddingModelProvider(e.target.value);
-                            setSelectedEmbeddingModel(
+                            if (e.target.value === 'custom_openai') {
-                              config.embeddingModelProviders[e.target.value][0],
+                              setSelectedEmbeddingModel('');
-                            );
+                            } else {
+                              setSelectedEmbeddingModel(
+                                config.embeddingModelProviders[e.target.value][0],
+                              );
+                            }
                           }}
                           options={Object.keys(
                             config.embeddingModelProviders,
@@ -355,7 +359,7 @@ const SettingsDialog = ({
                         />
                       </div>
                     )}
-                    {selectedEmbeddingModelProvider && (
+                    {selectedEmbeddingModelProvider && selectedEmbeddingModelProvider != 'custom_openai' && (
                       <div className="flex flex-col space-y-1">
                         <p className="text-black/70 dark:text-white/70 text-sm">
                           Embedding Model
@@ -368,34 +372,49 @@ const SettingsDialog = ({
                           options={(() => {
                             const embeddingModelProvider =
                               config.embeddingModelProviders[
-                                selectedEmbeddingModelProvider
+                              selectedEmbeddingModelProvider
                               ];
 
                             return embeddingModelProvider
                               ? embeddingModelProvider.length > 0
                                 ? embeddingModelProvider.map((model) => ({
-                                    label: model,
+                                  label: model,
-                                    value: model,
+                                  value: model,
-                                  }))
+                                }))
                                 : [
-                                    {
-                                      label: 'No embedding models available',
-                                      value: '',
-                                      disabled: true,
-                                    },
-                                  ]
-                              : [
                                   {
-                                    label:
+                                    label: 'No embedding models available',
-                                      'Invalid provider, please check backend logs',
                                     value: '',
                                     disabled: true,
                                   },
-                                ];
+                                ]
+                              : [
+                                {
+                                  label:
+                                    'Invalid provider, please check backend logs',
+                                  value: '',
+                                  disabled: true,
+                                },
+                              ];
                           })()}
                         />
                       </div>
                     )}
+                    {selectedEmbeddingModelProvider && selectedEmbeddingModelProvider === 'custom_openai' && (
+                      <div className="flex flex-col space-y-1">
+                        <p className="text-black/70 dark:text-white/70 text-sm">
+                          Embedding Model
+                        </p>
+                        <Input
+                          type="text"
+                          placeholder="Embedding Model name"
+                          defaultValue={selectedEmbeddingModel!}
+                          onChange={(e) =>
+                            setSelectedEmbeddingModel(e.target.value)
+                          }
+                        />
+                      </div>
+                    )}
                     <div className="flex flex-col space-y-1">
                       <p className="text-black/70 dark:text-white/70 text-sm">
                         OpenAI API Key