From 1d2afef9fe70f504c79fdbdef8a63d1b75669349 Mon Sep 17 00:00:00 2001
From: irene-chou <irene.chou0710@gmail.com>
Date: Tue, 7 Mar 2023 13:00:53 +0800
Subject: [PATCH] feat: add prompt param to use customized prompt template

---
 src/commands/aicommits.ts | 25 ++++++++++++++++++++-----
 src/utils/openai.ts       |  3 ++-
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/commands/aicommits.ts b/src/commands/aicommits.ts
index 2ed074c5..4d05f1ed 100644
--- a/src/commands/aicommits.ts
+++ b/src/commands/aicommits.ts
@@ -16,6 +16,7 @@ import { KnownError, handleCliError } from '../utils/error.js';
 
 export default async (
 	generate: number,
+	prompt: string,
 	rawArgv: string[],
 ) => (async () => {
 	intro(bgCyan(black(' aicommits ')));
@@ -46,9 +47,22 @@ export default async (
 		OPENAI_KEY,
 		staged.diff,
 		generate,
+		prompt,
 	);
 	s.stop('Changes analyzed');
 
+	if (!prompt) {
+		commitMessageProcess(messages, rawArgv);
+	} else {
+		customizedProcess(messages);
+	}
+})().catch((error) => {
+	outro(`${red('✖')} ${error.message}`);
+	handleCliError(error);
+	process.exit(1);
+});
+
+const commitMessageProcess = async (messages: string[], rawArgv: string[]) => {
 	let message: string;
 	if (messages.length === 1) {
 		[message] = messages;
@@ -77,8 +91,9 @@ export default async (
 	await execa('git', ['commit', '-m', message, ...rawArgv]);
 
 	outro(`${green('✔')} Successfully committed!`);
-})().catch((error) => {
-	outro(`${red('✖')} ${error.message}`);
-	handleCliError(error);
-	process.exit(1);
-});
+};
+
+const customizedProcess = async (messages: string[]) => {
+	const message = messages.length === 1 ? messages[0] : messages.join('\n');
+	outro(`Response message:\n\n   ${message}\n`);
+};
diff --git a/src/utils/openai.ts b/src/utils/openai.ts
index 1c3191f3..e6dbc668 100644
--- a/src/utils/openai.ts
+++ b/src/utils/openai.ts
@@ -58,7 +58,7 @@ const sanitizeMessage = (message: string) => message.trim().replace(/[\n\r]/g, '
 
 const deduplicateMessages = (array: string[]) => Array.from(new Set(array));
 
-const promptTemplate = 'Write an insightful but concise Git commit message in a complete sentence in present tense for the following diff without prefacing it with anything:';
+const promptDefaultTemplate = 'Write an insightful but concise Git commit message in a complete sentence in present tense for the following diff without prefacing it with anything:';
 
 const model = 'text-davinci-003';
 const encoder = encodingForModel(model);
@@ -67,6 +67,7 @@ export const generateCommitMessage = async (
 	apiKey: string,
 	diff: string,
 	completions: number,
+	promptTemplate: string = promptDefaultTemplate,
 ) => {
 	const prompt = `${promptTemplate}\n${diff}`;