feat: use prompt template

tiddly-gittly · Apr 13, 2024 · 08f7cb4 · 08f7cb4
1 parent 47e0535
commit 08f7cb4
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git a/src/services/languageModel/index.ts b/src/services/languageModel/index.ts
@@ -135,7 +135,7 @@ export class LanguageModel implements ILanguageModelService {
 
   public runLanguageModel$(runner: LanguageModelRunner.llamaCpp, options: IRunLLAmaOptions): Observable<ILanguageModelAPIResponse>;
   public runLanguageModel$(runner: LanguageModelRunner, options: IRunLLAmaOptions): Observable<ILanguageModelAPIResponse> {
-    const { id: conversationID, completionOptions, loadConfig: config } = options;
+    const { id: conversationID, loadConfig: config } = options;
     this.updateModelLoaded({ [runner]: null });
     return new Observable<ILanguageModelAPIResponse>((subscriber) => {
       const runLanguageModelObserverIIFE = async () => {
@@ -158,9 +158,10 @@ export class LanguageModel implements ILanguageModelService {
         } else {
           // load and run model
           const texts = { timeout: i18n.t('LanguageModel.GenerationTimeout'), disposed: i18n.t('LanguageModel.ModelDisposed') };
+          logger.info(options.sessionOptions?.systemPrompt ?? '', { tag: 'options.sessionOptions?.systemPrompt' });
           switch (runner) {
             case LanguageModelRunner.llamaCpp: {
-              observable = worker.runLLama({ completionOptions, loadConfig: { ...config, modelPath }, conversationID }, texts);
+              observable = worker.runLLama({ ...options, loadConfig: { ...config, modelPath }, conversationID }, texts);
               break;
             }
           }

diff --git a/src/services/languageModel/interface.ts b/src/services/languageModel/interface.ts
@@ -1,6 +1,6 @@
 import { LanguageModelChannel } from '@/constants/channels';
 import { ProxyPropertyType } from 'electron-ipc-cat/common';
-import type { LLamaChatPromptOptions, LlamaModelOptions } from 'node-llama-cpp';
+import type { LLamaChatPromptOptions, LlamaChatSessionOptions, LlamaModelOptions, JinjaTemplateChatWrapperOptions } from 'node-llama-cpp';
 import type { Observable } from 'rxjs';
 
 export enum LanguageModelRunner {
@@ -60,6 +60,8 @@ export interface IRunLLAmaOptions extends ILLMResultBase {
    * Without generating text.
    */
   loadModelOnly?: boolean;
+  sessionOptions?: Pick<LlamaChatSessionOptions, 'systemPrompt'>;
+  templates?: Partial<Pick<JinjaTemplateChatWrapperOptions, 'template' | 'systemRoleName'>>;
 }
 
 /**

diff --git a/src/services/languageModel/llmWorker/llamaCpp.ts b/src/services/languageModel/llmWorker/llamaCpp.ts
@@ -1,7 +1,7 @@
 import { LLAMA_PREBUILT_BINS_DIRECTORY } from './preload';
 
 import debounce from 'lodash/debounce';
-import { getLlama, Llama, LlamaChatSession, LlamaContext, LlamaContextSequence, LlamaModel, LlamaModelOptions } from 'node-llama-cpp';
+import { getLlama, Llama, LlamaChatSession, LlamaContext, LlamaContextSequence, LlamaModel, LlamaModelOptions, JinjaTemplateChatWrapper } from 'node-llama-cpp';
 import { Observable, Subscriber } from 'rxjs';
 import { ILanguageModelWorkerResponse, IRunLLAmaOptions } from '../interface';
 import { DEFAULT_TIMEOUT_DURATION } from './constants';
@@ -90,14 +90,10 @@ export async function unloadLLama() {
 }
 const runnerAbortControllers = new Map<string, AbortController>();
 export function runLLama(
-  options: {
-    completionOptions: IRunLLAmaOptions['completionOptions'];
-    conversationID: IRunLLAmaOptions['id'];
-    loadConfig: IRunLLAmaOptions['loadConfig'] & Pick<LlamaModelOptions, 'modelPath'>;
-  },
+  options: IRunLLAmaOptions & { conversationID: IRunLLAmaOptions['id']; loadConfig: IRunLLAmaOptions['loadConfig'] & Pick<LlamaModelOptions, 'modelPath'> },
   texts: { disposed: string; timeout: string },
 ): Observable<ILanguageModelWorkerResponse> {
-  const { conversationID, completionOptions, loadConfig } = options;
+  const { conversationID, completionOptions, sessionOptions, loadConfig, templates } = options;
 
   const loggerCommonMeta = { level: 'debug' as const, meta: { function: 'llmWorker.runLLama' }, id: conversationID };
   return new Observable<ILanguageModelWorkerResponse>((subscriber) => {
@@ -141,9 +137,15 @@ export function runLLama(
         if (contextSequenceInstance === undefined) {
           contextSequenceInstance = contextInstance.getSequence();
         }
+        const chatWrapper = new JinjaTemplateChatWrapper({
+          template: "{{ 'System: ' + systemPrompt if systemPrompt else '' }}{{ 'User: ' + userInput if userInput else '' }}",
+          ...templates,
+        });
         const session = new LlamaChatSession({
           contextSequence: contextSequenceInstance,
           autoDisposeSequence: false,
+          systemPrompt: sessionOptions?.systemPrompt,
+          chatWrapper,
         });
         await session.prompt(completionOptions.prompt, {
           ...completionOptions,
@@ -167,7 +169,7 @@ export function runLLama(
         subscriber.next({ message: 'createCompletion completed', ...loggerCommonMeta });
       } catch (error) {
         if ((error as Error).message.includes('aborted')) {
-          console.info('abortLLama', conversationID);
+          console.info(`abortLLama ${(error as Error).message}`, conversationID);
         } else {
           runnerAbortControllers.delete(conversationID);
           subscriber.error(error);