浏览代码

feat: CURRENT_DATE, USER_NAME prompt variable support

Timothy J. Baek 11 月之前
父节点
当前提交
9dfa334a83
共有 5 个文件被更改,包括 73 次插入27 次删除
  1. 2 2
      src/lib/apis/ollama/index.ts
  2. 2 2
      src/lib/apis/openai/index.ts
  3. 10 4
      src/lib/components/chat/Chat.svelte
  4. 17 17
      src/lib/utils/index.test.ts
  5. 42 2
      src/lib/utils/index.ts

+ 2 - 2
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);
 

+ 2 - 2
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);
 

+ 10 - 4
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')}`
 											: ''

+ 17 - 17
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:<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);
 });

+ 42 - 2
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) => {