feat: add support for defining custom models in config, switched toml library for proper serialization

.gitignore

@@ -2,10 +2,13 @@
 node_modules/
 npm-debug.log
 yarn-error.log
+.yarnrc.yml
 
 # Build output
 /.next/
 /out/
+dist/
+.yarn/
 
 # IDE/Editor specific
 .vscode/

package.json

@@ -20,7 +20,6 @@
     "typescript": "^5.4.3"
   },
   "dependencies": {
-    "@iarna/toml": "^2.2.5",
     "@langchain/openai": "^0.0.25",
     "@xenova/transformers": "^2.17.1",
     "axios": "^1.6.8",
@@ -30,8 +29,10 @@
     "dotenv": "^16.4.5",
     "express": "^4.19.2",
     "langchain": "^0.1.30",
+    "smol-toml": "^1.0.0",
     "winston": "^3.13.0",
     "ws": "^8.16.0",
     "zod": "^3.22.4"
-  }
+  },
+  "packageManager": "yarn@3.6.1+sha512.de524adec81a6c3d7a26d936d439d2832e351cdfc5728f9d91f3fc85dd20b04391c038e9b4ecab11cae2b0dd9f0d55fd355af766bc5c1a7f8d25d96bb2a0b2ca"
 }

sample.config.toml

@@ -9,3 +9,23 @@ GROQ = "" # Groq API key - gsk_1234567890abcdef1234567890abcdef
 [API_ENDPOINTS]
 SEARXNG = "http://localhost:32768" # SearxNG API URL
 OLLAMA = "" # Ollama API URL - http://host.docker.internal:11434
+
+[[MODELS]]
+name = "text-generation-webui"
+api_key = "blah"
+base_url = "http://localhost:5000/v1"
+provider = "openai"
+
+[[EMBEDDINGS]]
+name = "text-generation-webui-small"
+model = "text-embedding-3-small"
+api_key = "blah"
+base_url = "http://localhost:5000/v1"
+provider = "openai"
+
+[[EMBEDDINGS]]
+name = "text-generation-webui-large"
+model = "text-embedding-3-large"
+api_key = "blah"
+base_url = "http://localhost:5000/v1"
+provider = "openai"

src/config.ts

@@ -1,6 +1,6 @@
 import fs from 'fs';
 import path from 'path';
-import toml from '@iarna/toml';
+import {parse, stringify} from "smol-toml";
 
 const configFileName = 'config.toml';
 
@@ -17,6 +17,23 @@ interface Config {
         SEARXNG: string;
         OLLAMA: string;
     };
+    MODELS: [
+        {
+            "name": string;
+            "api_key": string;
+            "base_url": string;
+            "provider": string;
+        }
+    ];
+    EMBEDDINGS: [
+        {
+            "name": string;
+            "model": string;
+            "api_key": string;
+            "base_url": string;
+            "provider": string;
+        }
+    ];
 }
 
 type RecursivePartial<T> = {
@@ -24,7 +41,7 @@ type RecursivePartial<T> = {
 };
 
 const loadConfig = () =>
-  toml.parse(
+    parse(
         fs.readFileSync(path.join(__dirname, `../${configFileName}`), 'utf-8'),
     ) as any as Config;
 
@@ -41,29 +58,22 @@ export const getSearxngApiEndpoint = () => loadConfig().API_ENDPOINTS.SEARXNG;
 
 export const getOllamaApiEndpoint = () => loadConfig().API_ENDPOINTS.OLLAMA;
 
+export const getCustomModels = () => loadConfig().MODELS;
+
+export const getCustomEmbeddingModels = () => loadConfig().EMBEDDINGS;
+
 export const updateConfig = (config: RecursivePartial<Config>) => {
     const currentConfig = loadConfig();
 
-  for (const key in currentConfig) {
+    const updatedConfig = {
-    if (!config[key]) config[key] = {};
+        ...currentConfig,
+        ...config
+    };
 
-    if (typeof currentConfig[key] === 'object' && currentConfig[key] !== null) {
+    const toml = stringify(updatedConfig);
-      for (const nestedKey in currentConfig[key]) {
-        if (
-          !config[key][nestedKey] &&
-          currentConfig[key][nestedKey] &&
-          config[key][nestedKey] !== ''
-        ) {
-          config[key][nestedKey] = currentConfig[key][nestedKey];
-        }
-      }
-    } else if (currentConfig[key] && config[key] !== '') {
-      config[key] = currentConfig[key];
-    }
-  }
 
     fs.writeFileSync(
         path.join(__dirname, `../${configFileName}`),
-    toml.stringify(config),
+        toml,
     );
 };

src/lib/providers.ts

@@ -3,6 +3,8 @@ import { ChatOllama } from '@langchain/community/chat_models/ollama';
 import {OllamaEmbeddings} from '@langchain/community/embeddings/ollama';
 import {HuggingFaceTransformersEmbeddings} from './huggingfaceTransformer';
 import {
+    getCustomEmbeddingModels,
+    getCustomModels,
     getGroqApiKey,
     getOllamaApiEndpoint,
     getOpenaiApiKey,
@@ -13,6 +15,7 @@ export const getAvailableChatModelProviders = async () => {
     const openAIApiKey = getOpenaiApiKey();
     const groqApiKey = getGroqApiKey();
     const ollamaEndpoint = getOllamaApiEndpoint();
+    const customModels = getCustomModels();
 
     const models = {};
 
@@ -119,12 +122,36 @@ export const getAvailableChatModelProviders = async () => {
 
     models['custom_openai'] = {};
 
+    if (customModels && customModels.length > 0) {
+        models['custom'] = {};
+        try {
+            customModels.forEach((model) => {
+                if (model.provider === "openai") {
+                    models['custom'] = {
+                        ...models['custom'],
+                        [model.name]: new ChatOpenAI({
+                            openAIApiKey: model.api_key,
+                            modelName: model.name,
+                            temperature: 0.7,
+                            configuration: {
+                                baseURL: model.base_url,
+                            }
+                        })
+                    }
+                }
+            });
+        } catch (err) {
+            logger.error(`Error loading custom models: ${err}`);
+        }
+    }
+
     return models;
 };
 
 export const getAvailableEmbeddingModelProviders = async () => {
     const openAIApiKey = getOpenaiApiKey();
     const ollamaEndpoint = getOllamaApiEndpoint();
+    const customEmbeddingModels = getCustomEmbeddingModels();
 
     const models = {};
 
@@ -134,7 +161,7 @@ export const getAvailableEmbeddingModelProviders = async () => {
                 'Text embedding 3 small': new OpenAIEmbeddings({
                     openAIApiKey,
                     modelName: 'text-embedding-3-small',
-        }),
+                }, {baseURL: "http://10.0.1.2:5000/v1"}),
                 'Text embedding 3 large': new OpenAIEmbeddings({
                     openAIApiKey,
                     modelName: 'text-embedding-3-large',
@@ -167,6 +194,28 @@ export const getAvailableEmbeddingModelProviders = async () => {
         }
     }
 
+    if (customEmbeddingModels && customEmbeddingModels.length > 0) {
+        models['custom'] = {};
+        try {
+            customEmbeddingModels.forEach((model) => {
+                if (model.provider === "openai") {
+                    models['custom'] = {
+                        ...models['custom'],
+                        [model.name]: new OpenAIEmbeddings({
+                                openAIApiKey: model.api_key,
+                                modelName: model.model,
+                            },
+                            {
+                                baseURL: model.base_url,
+                            }),
+                    }
+                }
+            });
+        } catch (err) {
+            logger.error(`Error loading custom models: ${err}`);
+        }
+    }
+
     try {
         models['local'] = {
             'BGE Small': new HuggingFaceTransformersEmbeddings({