Updates generate rebase into a single-pass prompt

Co-authored-by: Eric Amodio <[email protected]>
Co-authored-by: Ramin Tadayon <[email protected]>

Author: 3 people

src/commands/generateRebase.ts

@@ -83,7 +83,7 @@ export class GenerateCommitsCommand extends GlCommandBase {
 				createReference(uncommitted, svc.path, { refType: 'revision' }),
 				createReference('HEAD', svc.path, { refType: 'revision' }),
 				args?.source ?? { source: 'commandPalette' },
-				{ progress: { location: ProgressLocation.Notification } },
+				{ title: 'Generate Commits', progress: { location: ProgressLocation.Notification } },
 			);
 		} catch (ex) {
 			Logger.error(ex, 'GenerateCommitsCommand', 'execute');
@@ -163,18 +163,20 @@ export async function generateRebase(
 	head: GitReference,
 	base: GitReference,
 	source: Source,
-	options?: { cancellation?: CancellationToken; progress?: ProgressOptions },
+	options?: { title?: string; cancellation?: CancellationToken; progress?: ProgressOptions },
 ): Promise<void> {
+	const { title, ...aiOptions } = options ?? {};
+
 	const repo = svc.getRepository()!;
-	const result = await container.ai.generateRebase(repo, base.ref, head.ref, source, options);
+	const result = await container.ai.generateRebase(repo, base.ref, head.ref, source, aiOptions);
 	if (result == null) return;
 
 	try {
 		// Extract the diff information from the reorganized commits
 		const diffInfo = extractRebaseDiffInfo(result.commits, result.diff, result.hunkMap);
 
 		// Generate the markdown content that shows each commit and its diffs
-		const markdownContent = generateRebaseMarkdown(result);
+		const markdownContent = generateRebaseMarkdown(result, title);
 
 		const shas = await repo.git.patch?.createUnreachableCommitsFromPatches(base.ref, diffInfo);
 		if (shas?.length) {
@@ -254,19 +256,32 @@ export function extractRebaseDiffInfo(
 /**
  * Formats the reorganized commits into a readable markdown document with proper git diff format
  */
-function generateRebaseMarkdown(result: AIRebaseResult): string {
-	let markdown = `# Rebase Commits\n\n`;
+function generateRebaseMarkdown(result: AIRebaseResult, title = 'Rebase Commits'): string {
+	let markdown = `# ${title}\n\n> Generated by ${result.model.name}\n\n`;
 
-	const { commits, diff: originalDiff, hunkMap, explanation } = result;
+	const { commits, diff: originalDiff, hunkMap } = result;
+
+	if (commits.length === 0) {
+		markdown += `No commits generated\n\n`;
+		return markdown;
+	}
 
-	markdown += `## Explanation\n${explanation}\n\n----\n\n`;
+	let explanations =
+		"## Explanation\n\nOkay, here's the breakdown of the commits I'd create from the provided diff, along with explanations for each:\n\n";
 
+	let changes = '## Commits\n\n';
 	for (let i = 0; i < commits.length; i++) {
 		const commit = commits[i];
 
-		markdown += `## Commit ${i + 1}: ${commit.message}\n\n`;
-		// markdown += `### Explanation\n${commit.explanation}\n\n`;
-		markdown += `### Changes\n`;
+		const commitTitle = `### Commit ${i + 1}: ${commit.message}`;
+
+		if (commit.explanation) {
+			explanations += `${commitTitle}\n\n${commit.explanation}\n\n`;
+		} else {
+			explanations += `${commitTitle}\n\nNo explanation provided.\n\n`;
+		}
+
+		changes += `${commitTitle}\n\n`;
 
 		// Group hunks by file (diff header)
 		const fileHunks = new Map<string, string[]>();
@@ -290,25 +305,27 @@ function generateRebaseMarkdown(result: AIRebaseResult): string {
 
 		// Output each file with its hunks in git patch format
 		for (const [diffHeader, hunkHeaders] of fileHunks.entries()) {
-			markdown += '```diff\n';
-			markdown += `${escapeMarkdownCodeBlocks(diffHeader)}\n`;
+			changes += '```diff\n';
+			changes += `${escapeMarkdownCodeBlocks(diffHeader)}\n`;
 
 			// Extract and include the actual content for each hunk from the original diff
 			for (const hunkHeader of hunkHeaders) {
-				// markdown += `${hunkHeader}\n`;
 				// Find the hunk content in the original diff
 				const hunkContent = extractHunkContent(originalDiff, diffHeader, hunkHeader);
 				if (hunkContent) {
-					markdown += `${escapeMarkdownCodeBlocks(hunkContent)}\n`;
+					changes += `${escapeMarkdownCodeBlocks(hunkContent)}\n`;
 				} else {
-					markdown += `Unable to extract hunk content for ${hunkHeader}\n`;
+					changes += `Unable to extract hunk content for ${hunkHeader}\n`;
 				}
 			}
 
-			markdown += '```\n\n';
+			changes += '```\n\n';
 		}
 	}
 
+	markdown += explanations;
+	markdown += changes;
+
 	// markdown += `\n\n----\n\n## Raw commits\n\n\`\`\`${escapeMarkdownCodeBlocks(JSON.stringify(commits))}\`\`\``;
 	// markdown += `\n\n----\n\n## Original Diff\n\n\`\`\`${escapeMarkdownCodeBlocks(originalDiff)}\`\`\`\n`;
 

src/plus/ai/aiProviderService.ts

@@ -93,10 +93,10 @@ export interface AISummarizeResult extends AIResult {
 
 export interface AIRebaseResult extends AIResult {
 	readonly diff: string;
-	readonly explanation: string;
 	readonly hunkMap: { index: number; hunkHeader: string }[];
 	readonly commits: {
 		readonly message: string;
+		readonly explanation: string;
 		readonly hunks: { hunk: number }[];
 	}[];
 }
@@ -871,82 +871,57 @@ export class AIProviderService implements Disposable {
 			commits: [],
 		} as unknown as AIRebaseResult;
 
-		const lazyDiff = lazy(() => repo.git.diff.getDiff?.(headRef, baseRef, { notation: '...' }));
-
-		const conversation: AIChatMessage[] = [];
-
-		const req1 = await this.sendRequest(
+		const rq = await this.sendRequest(
 			'generate-rebase',
 			async (model, reporting, cancellation, maxInputTokens, retries) => {
-				const diff = await lazyDiff.value;
-				if (!diff?.contents) {
-					throw new AINoRequestDataError('No changes found to generate a rebase from.');
-				}
-				if (cancellation.isCancellationRequested) throw new CancellationError();
+				const [diffResult, logResult] = await Promise.allSettled([
+					repo.git.diff.getDiff?.(headRef, baseRef, { notation: '...' }),
+					repo.git.commits.getLog(`${baseRef}..${headRef}`),
+				]);
 
-				result.diff = diff.contents;
+				const diff = getSettledValue(diffResult);
+				const log = getSettledValue(logResult);
 
-				const { prompt } = await this.getPrompt(
-					'generate-rebase-explanation',
-					model,
-					{
-						diff: diff.contents,
-						context: options?.context,
-						// instructions: configuration.get('ai.generateRebase.customInstructions'),
-					},
-					maxInputTokens,
-					retries,
-					reporting,
-				);
-				if (cancellation.isCancellationRequested) throw new CancellationError();
-
-				const messages: AIChatMessage[] = [{ role: 'user', content: prompt }];
-
-				conversation.push(...messages);
-				return messages;
-			},
-			m => `Generating rebase (examining changes) with ${m.name}...`,
-			source,
-			m => ({
-				key: 'ai/generate',
-				data: {
-					type: 'rebase',
-					'model.id': m.id,
-					'model.provider.id': m.provider.id,
-					'model.provider.name': m.provider.name,
-					'retry.count': 0,
-				},
-			}),
-			options,
-		);
-
-		conversation.push({
-			role: 'assistant',
-			content: req1!.content,
-		});
-		result.explanation = req1!.content;
-
-		const req2 = await this.sendRequest(
-			'generate-rebase',
-			async (model, reporting, cancellation, maxInputTokens, retries) => {
-				const diff = await lazyDiff.value;
-				if (!diff?.contents) {
+				if (!diff?.contents || !log?.commits?.size) {
 					throw new AINoRequestDataError('No changes found to generate a rebase from.');
 				}
 				if (cancellation.isCancellationRequested) throw new CancellationError();
 
+				result.diff = diff.contents;
+
 				const hunkMap: { index: number; hunkHeader: string }[] = [];
 				let counter = 0;
+				//const filesDiffs = await repo.git.diff().getDiffFiles!(diff.contents)!;
+				//for (const f of filesDiffs!.files)
+				//for (const hunk of parsedDiff.hunks) {
+				//	hunkMap.push({ index: ++counter, hunkHeader: hunk.contents.split('\n', 1)[0] });
+				//}
+
+				// let hunksByNumber= '';
+
 				for (const hunkHeader of diff.contents.matchAll(/@@ -\d+,\d+ \+\d+,\d+ @@(.*)$/gm)) {
 					hunkMap.push({ index: ++counter, hunkHeader: hunkHeader[0] });
 				}
 
 				result.hunkMap = hunkMap;
+				// 	const hunkNumber = `hunk-${counter++}`;
+				// 	hunksByNumber += `${hunkNumber}: ${hunk[0]}\n`;
+				// }
+
+				// const commits: { diff: string; message: string }[] = [];
+				// for (const commit of [...log.commits.values()].sort((a, b) => a.date.getTime() - b.date.getTime())) {
+				// 	const diff = await repo.git.diff().getDiff?.(commit.ref);
+				// 	commits.push({ message: commit.message ?? commit.summary, diff: diff?.contents ?? '' });
+
+				// 	if (cancellation.isCancellationRequested) throw new CancellationError();
+				// }
 
 				const { prompt } = await this.getPrompt(
-					'generate-rebase-commits',
+					'generate-rebase',
 					model,
 					{
+						diff: diff.contents,
+						// commits: JSON.stringify(commits),
 						data: JSON.stringify(hunkMap),
 						context: options?.context,
 						// instructions: configuration.get('ai.generateRebase.customInstructions'),
@@ -958,7 +933,7 @@ export class AIProviderService implements Disposable {
 				if (cancellation.isCancellationRequested) throw new CancellationError();
 
 				const messages: AIChatMessage[] = [{ role: 'user', content: prompt }];
-				return [...conversation, ...messages];
+				return messages;
 			},
 			m => `Generating rebase with ${m.name}...`,
 			source,
@@ -972,21 +947,23 @@ export class AIProviderService implements Disposable {
 					'retry.count': 0,
 				},
 			}),
-			{ ...options, modelOptions: { temperature: 0.2 } },
+			options,
 		);
 
-		// if it is wrapped in markdown, we need to strip it
-		const content = req2!.content.replace(/^\s*```json\s*/, '').replace(/\s*```$/, '');
-
 		try {
+			// if it is wrapped in markdown, we need to strip it
+			const content = rq!.content.replace(/^\s*```json\s*/, '').replace(/\s*```$/, '');
 			// Parse the JSON content from the result
 			result.commits = JSON.parse(content) as AIRebaseResult['commits'];
 		} catch {
 			debugger;
 			throw new Error('Unable to parse rebase result');
 		}
 
-		return result;
+		return {
+			...rq,
+			...result,
+		};
 	}
 
 	private async sendRequest<T extends AIActionType>(

src/plus/ai/models/promptTemplates.ts

@@ -28,19 +28,14 @@ interface CreatePullRequestPromptTemplateContext {
 	instructions?: string;
 }
 
-interface RebaseExplanationPromptTemplateContext {
+interface RebasePromptTemplateContext {
 	diff: string;
+	data?: string;
 	commits?: string;
 	context?: string;
 	instructions?: string;
 }
 
-interface RebaseCommitsPromptTemplateContext {
-	data: string;
-	context?: string;
-	instructions?: string;
-}
-
 interface ExplainChangesPromptTemplateContext {
 	diff: string;
 	message: string;
@@ -58,7 +53,7 @@ export type PromptTemplateType =
 	| 'generate-stashMessage'
 	| 'generate-changelog'
 	| `generate-create-${'cloudPatch' | 'codeSuggestion' | 'pullRequest'}`
-	| `generate-rebase-${'explanation' | 'commits'}`
+	| 'generate-rebase'
 	| 'explain-changes';
 
 type PromptTemplateVersions = '' | '_v2';
@@ -76,10 +71,8 @@ export type PromptTemplateContext<T extends PromptTemplateType> = T extends 'gen
 	? CreatePullRequestPromptTemplateContext
 	: T extends 'generate-changelog'
 	? ChangelogPromptTemplateContext
-	: T extends 'generate-rebase-explanation'
-	? RebaseExplanationPromptTemplateContext
-	: T extends 'generate-rebase-commits'
-	? RebaseCommitsPromptTemplateContext
+	: T extends 'generate-rebase'
+	? RebasePromptTemplateContext
 	: T extends 'explain-changes'
 	? ExplainChangesPromptTemplateContext
 	: never;