Add Ollama Model for local usage of yorkie intelligence

backend/.env.development

@@ -37,9 +37,14 @@ YORKIE_PROJECT_NAME=default
 YORKIE_PROJECT_SECRET_KEY=""
 
 # YORKIE_INTELLIGENCE: Whether to enable Yorkie Intelligence for collaborative editing.
-# Set to true if Yorkie Intelligence is required.
+# Set to ollama modelname if Yorkie Intelligence is required
+# you can find llm model in https://ollama.com/library
 # If set to false, OPENAI_API_KEY is not required.
-YORKIE_INTELLIGENCE=false
+YORKIE_INTELLIGENCE="ollama:gemma2:2b"
+
+# OLLAMA_HOST_URL: yorkie-intelligence ollama url
+OLLAMA_HOST_URL=http://localhost:11434
+
 # OPENAI_API_KEY: API key for using the gpt-3.5-turbo model by Yorkie Intelligence.
 # To obtain an API key, visit OpenAI: https://help.openai.com/en/articles/4936850-where-do-i-find-my-api-key
 OPENAI_API_KEY=your_openai_api_key_here

backend/docker/docker-compose-full.yml

@@ -40,6 +40,12 @@ services:
             - "8080:8080"
             - "8081:8081"
 
+    yorkie-intelligence:
+        image: "ollama/ollama:latest"
+        restart: always
+        expose:
+            - "11434"
+
     mongo:
         build:
             context: ./mongodb_replica

backend/docker/docker-compose.yml

@@ -9,6 +9,12 @@ services:
             - "8080:8080"
             - "8081:8081"
 
+    yorkie-intelligence:
+        image: "ollama/ollama:latest"
+        restart: always
+        expose:
+            - "11434"
+
     mongo:
         build:
             context: ./mongodb_replica

backend/package.json

@@ -32,6 +32,7 @@
 		"@aws-sdk/s3-request-presigner": "^3.509.0",
 		"@langchain/community": "^0.0.21",
 		"@langchain/core": "^0.1.18",
+		"@langchain/ollama": "^0.0.4",
 		"@nestjs/common": "^10.0.0",
 		"@nestjs/config": "^3.1.1",
 		"@nestjs/core": "^10.0.0",

backend/src/langchain/langchain.module.ts

@@ -1,10 +1,54 @@
 import { Module } from "@nestjs/common";
 import { ChatOpenAI } from "@langchain/openai";
+import { ChatOllama } from "@langchain/ollama";
 import { BaseChatModel } from "@langchain/core/language_models/chat_models";
 
+type ModelList = {
+	[key: string]: string[];
+};
+
+const modelList: ModelList = {
+	ollama: [
+		"lamma3.1",
+		"gemma2",
+		"gemma2:2b",
+		"phi3",
+		"mistral",
+		"neural-chat",
+		"starling-lm",
+		"solar",
+	],
+	openai: ["gpt-3.5-turbo", "gpt-4o-mini"],
+};
+
 const chatModelFactory = {
 	provide: "ChatModel",
-	useFactory: () => new ChatOpenAI({ modelName: "gpt-4o-mini" }) as BaseChatModel,
+	useFactory: () => {
+		const modelType = process.env.YORKIE_INTELLIGENCE;
+		try {
+			const [provider, model] = modelType.split(":", 2);
+			let chatModel: BaseChatModel | ChatOllama;
+
+			if (modelList[provider]?.includes(model)) {
+				if (provider === "ollama") {
+					chatModel = new ChatOllama({
+						model: model,
+						baseUrl: process.env.OLLAMA_HOST_URL,
+						checkOrPullModel: true,
+						streaming: true,
+					});
+				} else if (provider === "openai") {
+					chatModel = new ChatOpenAI({ modelName: model });
+				}
+			}
+
+			if (!chatModel) throw new Error();
+
+			return chatModel;
+		} catch {
+			throw new Error(`${modelType} is not found. please check your model name`);
+		}
+	},
 };
 
 @Module({

backend/src/settings/settings.service.ts

@@ -10,7 +10,7 @@ export class SettingsService {
 	async getSettings(): Promise<GetSettingsResponse> {
 		return {
 			yorkieIntelligence: {
-				enable: this.configService.get("YORKIE_INTELLIGENCE") === "true",
+				enable: this.configService.get("YORKIE_INTELLIGENCE") !== "false",
 				config: {
 					features: generateFeatureList(this.configService),
 				},