feat: CURRENT_DATE, USER_NAME prompt variable support

This commit is contained in:
Timothy J. Baek 2024-05-31 11:11:28 -07:00
parent 7674229e3a
commit 9dfa334a83
5 changed files with 73 additions and 27 deletions

View File

@ -1,5 +1,5 @@
import { OLLAMA_API_BASE_URL } from '$lib/constants'; import { OLLAMA_API_BASE_URL } from '$lib/constants';
import { promptTemplate } from '$lib/utils'; import { titleGenerationTemplate } from '$lib/utils';
export const getOllamaConfig = async (token: string = '') => { export const getOllamaConfig = async (token: string = '') => {
let error = null; let error = null;
@ -212,7 +212,7 @@ export const generateTitle = async (
) => { ) => {
let error = null; let error = null;
template = promptTemplate(template, prompt); template = titleGenerationTemplate(template, prompt);
console.log(template); console.log(template);

View File

@ -1,5 +1,5 @@
import { OPENAI_API_BASE_URL } from '$lib/constants'; 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'; import { type Model, models, settings } from '$lib/stores';
export const getOpenAIConfig = async (token: string = '') => { export const getOpenAIConfig = async (token: string = '') => {
@ -340,7 +340,7 @@ export const generateTitle = async (
) => { ) => {
let error = null; let error = null;
template = promptTemplate(template, prompt); template = titleGenerationTemplate(template, prompt);
console.log(template); console.log(template);

View File

@ -16,9 +16,15 @@
showSidebar, showSidebar,
tags as _tags, tags as _tags,
WEBUI_NAME, WEBUI_NAME,
banners banners,
user
} from '$lib/stores'; } 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 { cancelOllamaRequest, generateChatCompletion } from '$lib/apis/ollama';
import { import {
@ -515,7 +521,7 @@
$settings.system || (responseMessage?.userContext ?? null) $settings.system || (responseMessage?.userContext ?? null)
? { ? {
role: 'system', role: 'system',
content: `${$settings?.system ?? ''}${ content: `${promptTemplate($settings?.system ?? '', $user.name)}${
responseMessage?.userContext ?? null responseMessage?.userContext ?? null
? `\n\nUser Context:\n${(responseMessage?.userContext ?? []).join('\n')}` ? `\n\nUser Context:\n${(responseMessage?.userContext ?? []).join('\n')}`
: '' : ''
@ -816,7 +822,7 @@
$settings.system || (responseMessage?.userContext ?? null) $settings.system || (responseMessage?.userContext ?? null)
? { ? {
role: 'system', role: 'system',
content: `${$settings?.system ?? ''}${ content: `${promptTemplate($settings?.system ?? '', $user.name)}${
responseMessage?.userContext ?? null responseMessage?.userContext ?? null
? `\n\nUser Context:\n${(responseMessage?.userContext ?? []).join('\n')}` ? `\n\nUser Context:\n${(responseMessage?.userContext ?? []).join('\n')}`
: '' : ''

View File

@ -1,66 +1,66 @@
import { promptTemplate } from '$lib/utils/index'; import { titleGenerationTemplate } from '$lib/utils/index';
import { expect, test } from 'vitest'; import { expect, test } from 'vitest';
test('promptTemplate correctly replaces {{prompt}} placeholder', () => { test('titleGenerationTemplate correctly replaces {{prompt}} placeholder', () => {
const template = 'Hello {{prompt}}!'; const template = 'Hello {{prompt}}!';
const prompt = 'world'; const prompt = 'world';
const expected = 'Hello world!'; const expected = 'Hello world!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); expect(actual).toBe(expected);
}); });
test('promptTemplate correctly replaces {{prompt:start:<length>}} placeholder', () => { test('titleGenerationTemplate correctly replaces {{prompt:start:<length>}} placeholder', () => {
const template = 'Hello {{prompt:start:3}}!'; const template = 'Hello {{prompt:start:3}}!';
const prompt = 'world'; const prompt = 'world';
const expected = 'Hello wor!'; const expected = 'Hello wor!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); expect(actual).toBe(expected);
}); });
test('promptTemplate correctly replaces {{prompt:end:<length>}} placeholder', () => { test('titleGenerationTemplate correctly replaces {{prompt:end:<length>}} placeholder', () => {
const template = 'Hello {{prompt:end:3}}!'; const template = 'Hello {{prompt:end:3}}!';
const prompt = 'world'; const prompt = 'world';
const expected = 'Hello rld!'; const expected = 'Hello rld!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); expect(actual).toBe(expected);
}); });
test('promptTemplate correctly replaces {{prompt:middletruncate:<length>}} placeholder when prompt length is greater than length', () => { test('titleGenerationTemplate correctly replaces {{prompt:middletruncate:<length>}} placeholder when prompt length is greater than length', () => {
const template = 'Hello {{prompt:middletruncate:4}}!'; const template = 'Hello {{prompt:middletruncate:4}}!';
const prompt = 'world'; const prompt = 'world';
const expected = 'Hello wo...ld!'; const expected = 'Hello wo...ld!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); expect(actual).toBe(expected);
}); });
test('promptTemplate correctly replaces {{prompt:middletruncate:<length>}} placeholder when prompt length is less than or equal to length', () => { test('titleGenerationTemplate correctly replaces {{prompt:middletruncate:<length>}} placeholder when prompt length is less than or equal to length', () => {
const template = 'Hello {{prompt:middletruncate:5}}!'; const template = 'Hello {{prompt:middletruncate:5}}!';
const prompt = 'world'; const prompt = 'world';
const expected = 'Hello world!'; const expected = 'Hello world!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); 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 template = 'Hello world!';
const prompt = 'world'; const prompt = 'world';
const expected = 'Hello world!'; const expected = 'Hello world!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); 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 template = 'Hello {{prompt}}!';
const prompt = 'World, {{prompt}} injection'; const prompt = 'World, {{prompt}} injection';
const expected = 'Hello World, {{prompt}} injection!'; const expected = 'Hello World, {{prompt}} injection!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); 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 template = 'Hello {{prompt}}! This is {{prompt:start:3}}!';
const prompt = 'world'; const prompt = 'world';
const expected = 'Hello world! This is wor!'; const expected = 'Hello world! This is wor!';
const actual = promptTemplate(template, prompt); const actual = titleGenerationTemplate(template, prompt);
expect(actual).toBe(expected); expect(actual).toBe(expected);
}); });

View File

@ -449,6 +449,42 @@ export const blobToFile = (blob, fileName) => {
return file; 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. * This function is used to replace placeholders in a template string with the provided prompt.
* The placeholders can be in the following formats: * 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. * @param {string} prompt - The string to replace the placeholders with.
* @returns {string} The template string with the placeholders replaced by the prompt. * @returns {string} The template string with the placeholders replaced by the prompt.
*/ */
export const promptTemplate = (template: string, prompt: string): string => { export const titleGenerationTemplate = (template: string, prompt: string): string => {
return template.replace( template = template.replace(
/{{prompt}}|{{prompt:start:(\d+)}}|{{prompt:end:(\d+)}}|{{prompt:middletruncate:(\d+)}}/g, /{{prompt}}|{{prompt:start:(\d+)}}|{{prompt:end:(\d+)}}|{{prompt:middletruncate:(\d+)}}/g,
(match, startLength, endLength, middleLength) => { (match, startLength, endLength, middleLength) => {
if (match === '{{prompt}}') { if (match === '{{prompt}}') {
@ -482,6 +518,10 @@ export const promptTemplate = (template: string, prompt: string): string => {
return ''; return '';
} }
); );
template = promptTemplate(template);
return template;
}; };
export const approximateToHumanReadable = (nanoseconds: number) => { export const approximateToHumanReadable = (nanoseconds: number) => {