From fb344e74426508f2d2789afe18a85fb468cffd4a Mon Sep 17 00:00:00 2001
From: sagitchu <601096721@qq.com>
Date: Sat, 24 Aug 2024 11:29:52 +0800
Subject: [PATCH] add baseurl

---
 src/config.ts               |  5 +++++
 src/lib/providers/openai.ts | 25 ++++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/config.ts b/src/config.ts
index 9ebc182..3c42cd3 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -18,6 +18,9 @@ interface Config {
     SEARXNG: string;
     OLLAMA: string;
   };
+  BASE_URL: {
+    OPENAI: string;
+  };
 }
 
 type RecursivePartial<T> = {
@@ -36,6 +39,8 @@ export const getSimilarityMeasure = () =>
 
 export const getOpenaiApiKey = () => loadConfig().API_KEYS.OPENAI;
 
+export const getOpenaiBaseUrl = () => loadConfig().BASE_URL.OPENAI;
+
 export const getGroqApiKey = () => loadConfig().API_KEYS.GROQ;
 
 export const getAnthropicApiKey = () => loadConfig().API_KEYS.ANTHROPIC;
diff --git a/src/lib/providers/openai.ts b/src/lib/providers/openai.ts
index 8673954..56669ac 100644
--- a/src/lib/providers/openai.ts
+++ b/src/lib/providers/openai.ts
@@ -1,9 +1,10 @@
 import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai';
-import { getOpenaiApiKey } from '../../config';
+import { getOpenaiApiKey, getOpenaiBaseUrl } from '../../config';
 import logger from '../../utils/logger';
 
 export const loadOpenAIChatModels = async () => {
   const openAIApiKey = getOpenaiApiKey();
+  const openaibaseURL = getOpenaiBaseUrl();
 
   if (!openAIApiKey) return {};
 
@@ -11,26 +12,41 @@ export const loadOpenAIChatModels = async () => {
     const chatModels = {
       'GPT-3.5 turbo': new ChatOpenAI({
         openAIApiKey,
+        configuration: {
+          baseURL: openaibaseURL,
+        },
         modelName: 'gpt-3.5-turbo',
         temperature: 0.7,
       }),
       'GPT-4': new ChatOpenAI({
         openAIApiKey,
+        configuration: {
+          baseURL: openaibaseURL,
+        },
         modelName: 'gpt-4',
         temperature: 0.7,
       }),
       'GPT-4 turbo': new ChatOpenAI({
         openAIApiKey,
+        configuration: {
+          baseURL: openaibaseURL,
+        },
         modelName: 'gpt-4-turbo',
         temperature: 0.7,
       }),
       'GPT-4 omni': new ChatOpenAI({
         openAIApiKey,
+        configuration: {
+          baseURL: openaibaseURL,
+        },
         modelName: 'gpt-4o',
         temperature: 0.7,
       }),
       'GPT-4 omni mini': new ChatOpenAI({
         openAIApiKey,
+        configuration: {
+          baseURL: openaibaseURL,
+        },
         modelName: 'gpt-4o-mini',
         temperature: 0.7,
       }),
@@ -45,6 +61,7 @@ export const loadOpenAIChatModels = async () => {
 
 export const loadOpenAIEmbeddingsModels = async () => {
   const openAIApiKey = getOpenaiApiKey();
+  const openaibaseURL = getOpenaiBaseUrl();
 
   if (!openAIApiKey) return {};
 
@@ -52,10 +69,16 @@ export const loadOpenAIEmbeddingsModels = async () => {
     const embeddingModels = {
       'Text embedding 3 small': new OpenAIEmbeddings({
         openAIApiKey,
+        configuration: {
+          baseURL: openaibaseURL,
+        },
         modelName: 'text-embedding-3-small',
       }),
       'Text embedding 3 large': new OpenAIEmbeddings({
         openAIApiKey,
+        configuration: {
+          baseURL: openaibaseURL,
+        },
         modelName: 'text-embedding-3-large',
       }),
     };