feat(agents): support local LLMs

src/agents/academicSearchAgent.ts

@@ -11,7 +11,7 @@ import {
 } from '@langchain/core/runnables';
 import { StringOutputParser } from '@langchain/core/output_parsers';
 import { Document } from '@langchain/core/documents';
-import { searchSearxng } from '../core/searxng';
+import { searchSearxng } from '../lib/searxng';
 import type { StreamEvent } from '@langchain/core/tracers/log_stream';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { Embeddings } from '@langchain/core/embeddings';

src/agents/imageSearchAgent.ts

@@ -7,7 +7,7 @@ import { PromptTemplate } from '@langchain/core/prompts';
 import formatChatHistoryAsString from '../utils/formatHistory';
 import { BaseMessage } from '@langchain/core/messages';
 import { StringOutputParser } from '@langchain/core/output_parsers';
-import { searchSearxng } from '../core/searxng';
+import { searchSearxng } from '../lib/searxng';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 
 const imageSearchChainPrompt = `

src/agents/redditSearchAgent.ts

@@ -11,7 +11,7 @@ import {
 } from '@langchain/core/runnables';
 import { StringOutputParser } from '@langchain/core/output_parsers';
 import { Document } from '@langchain/core/documents';
-import { searchSearxng } from '../core/searxng';
+import { searchSearxng } from '../lib/searxng';
 import type { StreamEvent } from '@langchain/core/tracers/log_stream';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { Embeddings } from '@langchain/core/embeddings';

src/agents/webSearchAgent.ts

@@ -11,7 +11,7 @@ import {
 } from '@langchain/core/runnables';
 import { StringOutputParser } from '@langchain/core/output_parsers';
 import { Document } from '@langchain/core/documents';
-import { searchSearxng } from '../core/searxng';
+import { searchSearxng } from '../lib/searxng';
 import type { StreamEvent } from '@langchain/core/tracers/log_stream';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { Embeddings } from '@langchain/core/embeddings';

src/agents/wolframAlphaSearchAgent.ts

@@ -11,7 +11,7 @@ import {
 } from '@langchain/core/runnables';
 import { StringOutputParser } from '@langchain/core/output_parsers';
 import { Document } from '@langchain/core/documents';
-import { searchSearxng } from '../core/searxng';
+import { searchSearxng } from '../lib/searxng';
 import type { StreamEvent } from '@langchain/core/tracers/log_stream';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { Embeddings } from '@langchain/core/embeddings';

src/agents/youtubeSearchAgent.ts

@@ -11,7 +11,7 @@ import {
 } from '@langchain/core/runnables';
 import { StringOutputParser } from '@langchain/core/output_parsers';
 import { Document } from '@langchain/core/documents';
-import { searchSearxng } from '../core/searxng';
+import { searchSearxng } from '../lib/searxng';
 import type { StreamEvent } from '@langchain/core/tracers/log_stream';
 import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import type { Embeddings } from '@langchain/core/embeddings';

src/config.ts

@@ -1,6 +1,6 @@
 import fs from 'fs';
 import path from 'path';
-import toml, { JsonMap } from '@iarna/toml';
+import toml from '@iarna/toml';
 
 const configFileName = 'config.toml';
 
@@ -8,18 +8,21 @@ interface Config {
   GENERAL: {
     PORT: number;
     SIMILARITY_MEASURE: string;
+    CHAT_MODEL_PROVIDER: string;
+    CHAT_MODEL: string;
   };
   API_KEYS: {
     OPENAI: string;
   };
   API_ENDPOINTS: {
     SEARXNG: string;
+    OLLAMA: string;
   };
 }
 
 const loadConfig = () =>
   toml.parse(
-    fs.readFileSync(path.join(process.cwd(), `${configFileName}`), 'utf-8'),
+    fs.readFileSync(path.join(__dirname, `../${configFileName}`), 'utf-8'),
   ) as any as Config;
 
 export const getPort = () => loadConfig().GENERAL.PORT;
@@ -27,6 +30,13 @@ export const getPort = () => loadConfig().GENERAL.PORT;
 export const getSimilarityMeasure = () =>
   loadConfig().GENERAL.SIMILARITY_MEASURE;
 
+export const getChatModelProvider = () =>
+  loadConfig().GENERAL.CHAT_MODEL_PROVIDER;
+
+export const getChatModel = () => loadConfig().GENERAL.CHAT_MODEL;
+
 export const getOpenaiApiKey = () => loadConfig().API_KEYS.OPENAI;
 
 export const getSearxngApiEndpoint = () => loadConfig().API_ENDPOINTS.SEARXNG;
+
+export const getOllamaApiEndpoint = () => loadConfig().API_ENDPOINTS.OLLAMA;

src/core/agentPicker.ts

@@ -1,69 +0,0 @@
-import { z } from 'zod';
-import { OpenAI } from '@langchain/openai';
-import { RunnableSequence } from '@langchain/core/runnables';
-import { StructuredOutputParser } from 'langchain/output_parsers';
-import { PromptTemplate } from '@langchain/core/prompts';
-
-const availableAgents = [
-  {
-    name: 'webSearch',
-    description:
-      'It is expert is searching the web for information and answer user queries',
-  },
-  /* {
-    name: 'academicSearch',
-    description:
-      'It is expert is searching the academic databases for information and answer user queries. It is particularly good at finding research papers and articles on topics like science, engineering, and technology. Use this instead of wolframAlphaSearch if the user query is not mathematical or scientific in nature',
-  },
-  {
-    name: 'youtubeSearch',
-    description:
-      'This model is expert at finding videos on youtube based on user queries',
-  },
-  {
-    name: 'wolframAlphaSearch',
-    description:
-      'This model is expert at finding answers to mathematical and scientific questions based on user queries.',
-  },
-  {
-    name: 'redditSearch',
-    description:
-      'This model is expert at finding posts and discussions on reddit based on user queries',
-  },
-  {
-    name: 'writingAssistant',
-    description:
-      'If there is no need for searching, this model is expert at generating text based on user queries',
-  }, */
-];
-
-const parser = StructuredOutputParser.fromZodSchema(
-  z.object({
-    agent: z.string().describe('The name of the selected agent'),
-  }),
-);
-
-const prompt = `
-    You are an AI model who is expert at finding suitable agents for user queries. The available agents are:
-    ${availableAgents.map((agent) => `- ${agent.name}: ${agent.description}`).join('\n')}
-
-    Your task is to find the most suitable agent for the following query: {query}
-
-    {format_instructions}
-`;
-
-const chain = RunnableSequence.from([
-  PromptTemplate.fromTemplate(prompt),
-  new OpenAI({ temperature: 0 }),
-  parser,
-]);
-
-const pickSuitableAgent = async (query: string) => {
-  const res = await chain.invoke({
-    query,
-    format_instructions: parser.getFormatInstructions(),
-  });
-  return res.agent;
-};
-
-export default pickSuitableAgent;

src/lib/providers.ts

@@ -0,0 +1,58 @@
+import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai';
+import { ChatOllama } from '@langchain/community/chat_models/ollama';
+import { OllamaEmbeddings } from '@langchain/community/embeddings/ollama';
+import { getOllamaApiEndpoint, getOpenaiApiKey } from '../config';
+
+export const getAvailableProviders = async () => {
+  const openAIApiKey = getOpenaiApiKey();
+  const ollamaEndpoint = getOllamaApiEndpoint();
+
+  const models = {};
+
+  if (openAIApiKey) {
+    models['openai'] = {
+      'gpt-3.5-turbo': new ChatOpenAI({
+        openAIApiKey,
+        modelName: 'gpt-3.5-turbo',
+        temperature: 0.7,
+      }),
+      'gpt-4': new ChatOpenAI({
+        openAIApiKey,
+        modelName: 'gpt-4',
+        temperature: 0.7,
+      }),
+      embeddings: new OpenAIEmbeddings({
+        openAIApiKey,
+        modelName: 'text-embedding-3-large',
+      }),
+    };
+  }
+
+  if (ollamaEndpoint) {
+    try {
+      const response = await fetch(`${ollamaEndpoint}/api/tags`);
+
+      const { models: ollamaModels } = (await response.json()) as any;
+
+      models['ollama'] = ollamaModels.reduce((acc, model) => {
+        acc[model.model] = new ChatOllama({
+          baseUrl: ollamaEndpoint,
+          model: model.model,
+          temperature: 0.7,
+        });
+        return acc;
+      }, {});
+
+      if (Object.keys(models['ollama']).length > 0) {
+        models['ollama']['embeddings'] = new OllamaEmbeddings({
+          baseUrl: ollamaEndpoint,
+          model: models['ollama'][Object.keys(models['ollama'])[0]].model,
+        });
+      }
+    } catch (err) {
+      console.log(err);
+    }
+  }
+
+  return models;
+};

src/websocket/connectionManager.ts

@@ -1,18 +1,32 @@
 import { WebSocket } from 'ws';
 import { handleMessage } from './messageHandler';
-import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai';
-import { getOpenaiApiKey } from '../config';
+import { getChatModel, getChatModelProvider } from '../config';
+import { getAvailableProviders } from '../lib/providers';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import type { Embeddings } from '@langchain/core/embeddings';
 
-export const handleConnection = (ws: WebSocket) => {
-  const llm = new ChatOpenAI({
-    temperature: 0.7,
-    openAIApiKey: getOpenaiApiKey(),
-  });
+export const handleConnection = async (ws: WebSocket) => {
+  const models = await getAvailableProviders();
+  const provider = getChatModelProvider();
+  const chatModel = getChatModel();
 
-  const embeddings = new OpenAIEmbeddings({
-    openAIApiKey: getOpenaiApiKey(),
-    modelName: 'text-embedding-3-large',
-  });
+  let llm: BaseChatModel | undefined;
+  let embeddings: Embeddings | undefined;
+
+  if (models[provider] && models[provider][chatModel]) {
+    llm = models[provider][chatModel] as BaseChatModel | undefined;
+    embeddings = models[provider].embeddings as Embeddings | undefined;
+  }
+
+  if (!llm || !embeddings) {
+    ws.send(
+      JSON.stringify({
+        type: 'error',
+        data: 'Invalid LLM or embeddings model selected',
+      }),
+    );
+    ws.close();
+  }
 
   ws.on(
     'message',