From e6365c8ac66aa1a23f05b6d40101524ce410ff0a Mon Sep 17 00:00:00 2001
From: ckt1031 <65409152+ckt1031@users.noreply.github.com>
Date: Thu, 2 May 2024 18:55:05 +0800
Subject: [PATCH] Add OPENAI_CUSTOM_CHAT_MODEL configuration option

---
 sample.config.toml   |  1 +
 src/config.ts        |  4 ++++
 src/lib/providers.ts | 16 ++++++++++++++++
 src/routes/config.ts |  3 +++
 4 files changed, 24 insertions(+)

diff --git a/sample.config.toml b/sample.config.toml
index 836b5b5..496a1ab 100644
--- a/sample.config.toml
+++ b/sample.config.toml
@@ -5,6 +5,7 @@ SIMILARITY_MEASURE = "cosine" # "cosine" or "dot"
 [API_KEYS]
 OPENAI = "" # OpenAI API key - sk-1234567890abcdef1234567890abcdef
 OPENAI_BASE_URL = "https://api.openai.com/v1" # OpenAI API base URL
+OPENAI_CUSTOM_CHAT_MODEL = "" # OpenAI custom chat model ID - gpt-3.5-turbo-1106
 GROQ = "" # Groq API key - gsk_1234567890abcdef1234567890abcdef
 
 [API_ENDPOINTS]
diff --git a/src/config.ts b/src/config.ts
index cf3eebf..73da0c2 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -12,6 +12,7 @@ interface Config {
   API_KEYS: {
     OPENAI: string;
     OPENAI_BASE_URL: string;
+    OPENAI_CUSTOM_CHAT_MODEL: string;
     GROQ: string;
   };
   API_ENDPOINTS: {
@@ -39,6 +40,9 @@ export const getOpenaiApiKey = () => loadConfig().API_KEYS.OPENAI;
 export const getOpenaiBaseUrl = () =>
   loadConfig().API_KEYS.OPENAI_BASE_URL || 'https://api.openai.com/v1';
 
+export const getOpenaiCustomChatModel = () =>
+  loadConfig().API_KEYS.OPENAI_CUSTOM_CHAT_MODEL;
+
 export const getGroqApiKey = () => loadConfig().API_KEYS.GROQ;
 
 export const getSearxngApiEndpoint = () => loadConfig().API_ENDPOINTS.SEARXNG;
diff --git a/src/lib/providers.ts b/src/lib/providers.ts
index 6528183..8bd5b24 100644
--- a/src/lib/providers.ts
+++ b/src/lib/providers.ts
@@ -6,12 +6,14 @@ import {
   getOllamaApiEndpoint,
   getOpenaiApiKey,
   getOpenaiBaseUrl,
+  getOpenaiCustomChatModel,
 } from '../config';
 import logger from '../utils/logger';
 
 export const getAvailableProviders = async () => {
   const openAIApiKey = getOpenaiApiKey();
   const openAIBaseUrl = getOpenaiBaseUrl();
+  const openAICustomChatModel = getOpenaiCustomChatModel();
   const groqApiKey = getGroqApiKey();
   const ollamaEndpoint = getOllamaApiEndpoint();
 
@@ -50,6 +52,20 @@ export const getAvailableProviders = async () => {
             baseURL: openAIBaseUrl,
           },
         ),
+        ...(openAICustomChatModel.length > 0
+          ? {
+              openAICustomChatModel: new ChatOpenAI(
+                {
+                  openAIApiKey,
+                  modelName: openAICustomChatModel,
+                  temperature: 0.7,
+                },
+                {
+                  baseURL: openAIBaseUrl,
+                },
+              ),
+            }
+          : {}),
         embeddings: new OpenAIEmbeddings(
           {
             openAIApiKey,
diff --git a/src/routes/config.ts b/src/routes/config.ts
index f9e3545..f14582c 100644
--- a/src/routes/config.ts
+++ b/src/routes/config.ts
@@ -5,6 +5,7 @@ import {
   getOllamaApiEndpoint,
   getOpenaiApiKey,
   getOpenaiBaseUrl,
+  getOpenaiCustomChatModel,
   updateConfig,
 } from '../config';
 
@@ -27,6 +28,7 @@ router.get('/', async (_, res) => {
 
   config['openaiApiKey'] = getOpenaiApiKey();
   config['openaiBaseUrl'] = getOpenaiBaseUrl();
+  config['openaiCustomChatModel'] = getOpenaiCustomChatModel();
   config['ollamaApiUrl'] = getOllamaApiEndpoint();
   config['groqApiKey'] = getGroqApiKey();
 
@@ -40,6 +42,7 @@ router.post('/', async (req, res) => {
     API_KEYS: {
       OPENAI: config.openaiApiKey,
       OPENAI_BASE_URL: config.openaiBaseUrl,
+      OPENAI_CUSTOM_CHAT_MODEL: config.openaiCustomChatModel,
       GROQ: config.groqApiKey,
     },
     API_ENDPOINTS: {