From 9dfa334a8344cd45a6cdc7704b60c97e498cee0f Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Fri, 31 May 2024 11:11:28 -0700 Subject: [PATCH] feat: CURRENT_DATE, USER_NAME prompt variable support --- src/lib/apis/ollama/index.ts | 4 +-- src/lib/apis/openai/index.ts | 4 +-- src/lib/components/chat/Chat.svelte | 14 ++++++--- src/lib/utils/index.test.ts | 34 +++++++++++----------- src/lib/utils/index.ts | 44 +++++++++++++++++++++++++++-- 5 files changed, 73 insertions(+), 27 deletions(-) diff --git a/src/lib/apis/ollama/index.ts b/src/lib/apis/ollama/index.ts index 7bd0ed20a..5ab2363cb 100644 --- a/src/lib/apis/ollama/index.ts +++ b/src/lib/apis/ollama/index.ts @@ -1,5 +1,5 @@ import { OLLAMA_API_BASE_URL } from '$lib/constants'; -import { promptTemplate } from '$lib/utils'; +import { titleGenerationTemplate } from '$lib/utils'; export const getOllamaConfig = async (token: string = '') => { let error = null; @@ -212,7 +212,7 @@ export const generateTitle = async ( ) => { let error = null; - template = promptTemplate(template, prompt); + template = titleGenerationTemplate(template, prompt); console.log(template); diff --git a/src/lib/apis/openai/index.ts b/src/lib/apis/openai/index.ts index ddeab495c..fc9365d12 100644 --- a/src/lib/apis/openai/index.ts +++ b/src/lib/apis/openai/index.ts @@ -1,5 +1,5 @@ import { OPENAI_API_BASE_URL } from '$lib/constants'; -import { promptTemplate } from '$lib/utils'; +import { titleGenerationTemplate } from '$lib/utils'; import { type Model, models, settings } from '$lib/stores'; export const getOpenAIConfig = async (token: string = '') => { @@ -340,7 +340,7 @@ export const generateTitle = async ( ) => { let error = null; - template = promptTemplate(template, prompt); + template = titleGenerationTemplate(template, prompt); console.log(template); diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte index 348d79d75..ab1b482de 100644 --- a/src/lib/components/chat/Chat.svelte +++ b/src/lib/components/chat/Chat.svelte @@ -16,9 +16,15 @@ showSidebar, tags as _tags, WEBUI_NAME, - banners + banners, + user } from '$lib/stores'; - import { convertMessagesToHistory, copyToClipboard, splitStream } from '$lib/utils'; + import { + convertMessagesToHistory, + copyToClipboard, + promptTemplate, + splitStream + } from '$lib/utils'; import { cancelOllamaRequest, generateChatCompletion } from '$lib/apis/ollama'; import { @@ -515,7 +521,7 @@ $settings.system || (responseMessage?.userContext ?? null) ? { role: 'system', - content: `${$settings?.system ?? ''}${ + content: `${promptTemplate($settings?.system ?? '', $user.name)}${ responseMessage?.userContext ?? null ? `\n\nUser Context:\n${(responseMessage?.userContext ?? []).join('\n')}` : '' @@ -816,7 +822,7 @@ $settings.system || (responseMessage?.userContext ?? null) ? { role: 'system', - content: `${$settings?.system ?? ''}${ + content: `${promptTemplate($settings?.system ?? '', $user.name)}${ responseMessage?.userContext ?? null ? `\n\nUser Context:\n${(responseMessage?.userContext ?? []).join('\n')}` : '' diff --git a/src/lib/utils/index.test.ts b/src/lib/utils/index.test.ts index f67449dcc..f69ef2260 100644 --- a/src/lib/utils/index.test.ts +++ b/src/lib/utils/index.test.ts @@ -1,66 +1,66 @@ -import { promptTemplate } from '$lib/utils/index'; +import { titleGenerationTemplate } from '$lib/utils/index'; import { expect, test } from 'vitest'; -test('promptTemplate correctly replaces {{prompt}} placeholder', () => { +test('titleGenerationTemplate correctly replaces {{prompt}} placeholder', () => { const template = 'Hello {{prompt}}!'; const prompt = 'world'; const expected = 'Hello world!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); -test('promptTemplate correctly replaces {{prompt:start:}} placeholder', () => { +test('titleGenerationTemplate correctly replaces {{prompt:start:}} placeholder', () => { const template = 'Hello {{prompt:start:3}}!'; const prompt = 'world'; const expected = 'Hello wor!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); -test('promptTemplate correctly replaces {{prompt:end:}} placeholder', () => { +test('titleGenerationTemplate correctly replaces {{prompt:end:}} placeholder', () => { const template = 'Hello {{prompt:end:3}}!'; const prompt = 'world'; const expected = 'Hello rld!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); -test('promptTemplate correctly replaces {{prompt:middletruncate:}} placeholder when prompt length is greater than length', () => { +test('titleGenerationTemplate correctly replaces {{prompt:middletruncate:}} placeholder when prompt length is greater than length', () => { const template = 'Hello {{prompt:middletruncate:4}}!'; const prompt = 'world'; const expected = 'Hello wo...ld!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); -test('promptTemplate correctly replaces {{prompt:middletruncate:}} placeholder when prompt length is less than or equal to length', () => { +test('titleGenerationTemplate correctly replaces {{prompt:middletruncate:}} placeholder when prompt length is less than or equal to length', () => { const template = 'Hello {{prompt:middletruncate:5}}!'; const prompt = 'world'; const expected = 'Hello world!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); -test('promptTemplate returns original template when no placeholders are present', () => { +test('titleGenerationTemplate returns original template when no placeholders are present', () => { const template = 'Hello world!'; const prompt = 'world'; const expected = 'Hello world!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); -test('promptTemplate does not replace placeholders inside of replaced placeholders', () => { +test('titleGenerationTemplate does not replace placeholders inside of replaced placeholders', () => { const template = 'Hello {{prompt}}!'; const prompt = 'World, {{prompt}} injection'; const expected = 'Hello World, {{prompt}} injection!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); -test('promptTemplate correctly replaces multiple placeholders', () => { +test('titleGenerationTemplate correctly replaces multiple placeholders', () => { const template = 'Hello {{prompt}}! This is {{prompt:start:3}}!'; const prompt = 'world'; const expected = 'Hello world! This is wor!'; - const actual = promptTemplate(template, prompt); + const actual = titleGenerationTemplate(template, prompt); expect(actual).toBe(expected); }); diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts index 8594ff8c9..15ac73f1b 100644 --- a/src/lib/utils/index.ts +++ b/src/lib/utils/index.ts @@ -449,6 +449,42 @@ export const blobToFile = (blob, fileName) => { return file; }; +/** + * @param {string} template - The template string containing placeholders. + * @returns {string} The template string with the placeholders replaced by the prompt. + */ +export const promptTemplate = ( + template: string, + user_name?: string, + current_location?: string +): string => { + // Get the current date + const currentDate = new Date(); + + // Format the date to YYYY-MM-DD + const formattedDate = + currentDate.getFullYear() + + '-' + + String(currentDate.getMonth() + 1).padStart(2, '0') + + '-' + + String(currentDate.getDate()).padStart(2, '0'); + + // Replace {{CURRENT_DATE}} in the template with the formatted date + template = template.replace('{{CURRENT_DATE}}', formattedDate); + + if (user_name) { + // Replace {{USER_NAME}} in the template with the user's name + template = template.replace('{{USER_NAME}}', user_name); + } + + if (current_location) { + // Replace {{CURRENT_LOCATION}} in the template with the current location + template = template.replace('{{CURRENT_LOCATION}}', current_location); + } + + return template; +}; + /** * This function is used to replace placeholders in a template string with the provided prompt. * The placeholders can be in the following formats: @@ -461,8 +497,8 @@ export const blobToFile = (blob, fileName) => { * @param {string} prompt - The string to replace the placeholders with. * @returns {string} The template string with the placeholders replaced by the prompt. */ -export const promptTemplate = (template: string, prompt: string): string => { - return template.replace( +export const titleGenerationTemplate = (template: string, prompt: string): string => { + template = template.replace( /{{prompt}}|{{prompt:start:(\d+)}}|{{prompt:end:(\d+)}}|{{prompt:middletruncate:(\d+)}}/g, (match, startLength, endLength, middleLength) => { if (match === '{{prompt}}') { @@ -482,6 +518,10 @@ export const promptTemplate = (template: string, prompt: string): string => { return ''; } ); + + template = promptTemplate(template); + + return template; }; export const approximateToHumanReadable = (nanoseconds: number) => {