From f96d57153df623d6d68530012709cc1b352b0fe9 Mon Sep 17 00:00:00 2001
From: samlhuillier <sam.jansa123@gmail.com>
Date: Tue, 5 Dec 2023 16:37:51 -0600
Subject: [PATCH] fully custom local model path speicifcation working

---
 electron/main/llm/Types.ts              |  8 +----
 electron/main/llm/llmSessionHandlers.ts |  2 +-
 electron/main/llm/models/LlamaCpp.ts    | 43 +++++++++++--------------
 electron/main/llm/models/OpenAI.ts      |  2 +-
 4 files changed, 22 insertions(+), 33 deletions(-)

diff --git a/electron/main/llm/Types.ts b/electron/main/llm/Types.ts
index 7901c646..fb21c2ac 100644
--- a/electron/main/llm/Types.ts
+++ b/electron/main/llm/Types.ts
@@ -1,15 +1,9 @@
-export interface IModel {
-  loadModel(): Promise<void>;
-  unloadModel(): Promise<void>;
-  isModelLoaded(): boolean;
-}
-
 export interface ISessionService {
   /**
    * Initializes the session.
    * @returns A promise that resolves when the initialization is complete.
    */
-  init(): Promise<void>;
+  // init(): Promise<void>;
 
   /**
    * Handles the streaming of prompts.
diff --git a/electron/main/llm/llmSessionHandlers.ts b/electron/main/llm/llmSessionHandlers.ts
index 501b423d..e107c02f 100644
--- a/electron/main/llm/llmSessionHandlers.ts
+++ b/electron/main/llm/llmSessionHandlers.ts
@@ -81,7 +81,7 @@ async function createSession(
     );
     sessions[sessionId] = sessionService;
   } else {
-    const sessionService = new LlamaCPPSessionService();
+    const sessionService = new LlamaCPPSessionService(currentConfig.localPath);
     sessions[sessionId] = sessionService;
   }
   return sessionId;
diff --git a/electron/main/llm/models/LlamaCpp.ts b/electron/main/llm/models/LlamaCpp.ts
index b4497690..defe17c2 100644
--- a/electron/main/llm/models/LlamaCpp.ts
+++ b/electron/main/llm/models/LlamaCpp.ts
@@ -1,26 +1,35 @@
 import path from "path";
 import os from "os";
-import { IModel, ISendFunctionImplementer, ISessionService } from "../Types";
+import { ISendFunctionImplementer, ISessionService } from "../Types";
 
 export class LlamaCPPSessionService implements ISessionService {
   private session: any;
   public context: any;
   private model: any; // Model instance
 
-  constructor() {
-    this.init();
+  constructor(localModelPath: string) {
+    this.init(localModelPath);
   }
 
-  private async loadModel(): Promise<void> {
+  async init(localModelPath: string): Promise<void> {
+    await this.loadModel(localModelPath);
+    if (!this.isModelLoaded()) {
+      throw new Error("Model not loaded");
+    }
+
+    import("node-llama-cpp").then(async (nodeLLamaCpp: any) => {
+      this.context = new nodeLLamaCpp.LlamaContext({ model: this.model });
+      this.session = new nodeLLamaCpp.LlamaChatSession({
+        context: this.context,
+      });
+    });
+  }
+
+  private async loadModel(localModelPath: string): Promise<void> {
     // Load model logic - similar to what was in LlamaCPPModelLoader
     const nodeLLamaCpp = await import("node-llama-cpp");
     this.model = new nodeLLamaCpp.LlamaModel({
-      modelPath: path.join(
-        os.homedir(),
-        "Downloads",
-        "tinyllama-2-1b-miniguanaco.Q2_K.gguf"
-        // "mistral-7b-v0.1.Q4_K_M.gguf"
-      ),
+      modelPath: localModelPath,
       gpuLayers: 0,
     });
 
@@ -46,20 +55,6 @@ export class LlamaCPPSessionService implements ISessionService {
     return !!this.model;
   }
 
-  async init() {
-    await this.loadModel();
-    if (!this.isModelLoaded()) {
-      throw new Error("Model not loaded");
-    }
-
-    import("node-llama-cpp").then(async (nodeLLamaCpp: any) => {
-      this.context = new nodeLLamaCpp.LlamaContext({ model: this.model });
-      this.session = new nodeLLamaCpp.LlamaChatSession({
-        context: this.context,
-      });
-    });
-  }
-
   public async streamingPrompt(
     prompt: string,
     sendFunctionImplementer: ISendFunctionImplementer
diff --git a/electron/main/llm/models/OpenAI.ts b/electron/main/llm/models/OpenAI.ts
index 7367f95a..9a198778 100644
--- a/electron/main/llm/models/OpenAI.ts
+++ b/electron/main/llm/models/OpenAI.ts
@@ -1,5 +1,5 @@
 import OpenAI from "openai";
-import { IModel, ISendFunctionImplementer, ISessionService } from "../Types";
+import { ISendFunctionImplementer, ISessionService } from "../Types";
 
 export class OpenAIModelSessionService implements ISessionService {
   private openai: OpenAI;