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 { 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);

View File

@ -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);

View File

@ -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')}`
: ''

View File

@ -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:<length>}} placeholder', () => {
test('titleGenerationTemplate correctly replaces {{prompt:start:<length>}} 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:<length>}} placeholder', () => {
test('titleGenerationTemplate correctly replaces {{prompt:end:<length>}} 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:<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 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:<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 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);
});

View File

@ -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) => {