chore(audoedit): decouple codeToReplaceData from getPromptForModelType

vscode/src/autoedits/analytics-logger/analytics-logger.ts

@@ -631,6 +631,6 @@ export class AutoeditAnalyticsLogger {
 
 export const autoeditAnalyticsLogger = new AutoeditAnalyticsLogger()
 
-function getTimeNowInMillis(): number {
+export function getTimeNowInMillis(): number {
     return Math.floor(performance.now())
 }

vscode/src/autoedits/autoedits-provider.ts

@@ -16,10 +16,11 @@ import { getCurrentDocContext } from '../completions/get-current-doc-context'
 
 import type { AutoeditsModelAdapter, AutoeditsPrompt } from './adapters/base'
 import { createAutoeditsModelAdapter } from './adapters/create-adapter'
+import { getTimeNowInMillis } from './analytics-logger'
 import { autoeditsProviderConfig } from './autoedits-config'
 import { FilterPredictionBasedOnRecentEdits } from './filter-prediction-edits'
 import { autoeditsOutputChannelLogger } from './output-channel-logger'
-import type { CodeToReplaceData } from './prompt/prompt-utils'
+import { type CodeToReplaceData, getCurrentFilePromptComponents } from './prompt/prompt-utils'
 import { ShortTermPromptStrategy } from './prompt/short-term-diff-prompt-strategy'
 import type { DecorationInfo } from './renderer/decorators/base'
 import { DefaultDecorator } from './renderer/decorators/default-decorator'
@@ -139,7 +140,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
         inlineCompletionContext: vscode.InlineCompletionContext,
         token?: vscode.CancellationToken
     ): Promise<vscode.InlineCompletionItem[] | vscode.InlineCompletionList | null> {
-        const start = Date.now()
+        const start = getTimeNowInMillis()
         const controller = new AbortController()
         const abortSignal = controller.signal
         token?.onCancellationRequested(() => controller.abort())
@@ -164,6 +165,13 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
             maxSuffixLength: tokensToChars(autoeditsProviderConfig.tokenLimit.suffixTokens),
         })
 
+        const { fileWithMarkerPrompt, areaPrompt, codeToReplaceData } = getCurrentFilePromptComponents({
+            docContext,
+            document,
+            position,
+            tokenBudget: autoeditsProviderConfig.tokenLimit,
+        })
+
         autoeditsOutputChannelLogger.logDebug(
             'provideInlineCompletionItems',
             'Calculating context from contextMixer...'
@@ -186,10 +194,9 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
             'provideInlineCompletionItems',
             'Calculating prompt from promptStrategy...'
         )
-        const { codeToReplaceData, prompt } = this.promptStrategy.getPromptForModelType({
-            document,
-            position,
-            docContext,
+        const prompt = this.promptStrategy.getPromptForModelType({
+            fileWithMarkerPrompt,
+            areaPrompt,
             context,
             tokenBudget: autoeditsProviderConfig.tokenLimit,
             isChatModel: autoeditsProviderConfig.isChatModel,
@@ -216,7 +223,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
         autoeditsOutputChannelLogger.logDebug(
             'provideInlineCompletionItems',
             `========================== Response:\n${initialPrediction}\n` +
-                `========================== Time Taken: ${Date.now() - start}ms`
+                `========================== Time Taken: ${getTimeNowInMillis() - start}ms`
         )
 
         const prediction = shrinkPredictionUntilSuffix({

vscode/src/autoedits/prompt/base.ts

@@ -1,59 +1,41 @@
-import type * as vscode from 'vscode'
-
 import type {
     AutoEditsTokenLimit,
     AutocompleteContextSnippet,
-    DocumentContext,
     PromptString,
 } from '@sourcegraph/cody-shared'
 
 import type { AutoeditsPrompt } from '../adapters/base'
 
 import { SYSTEM_PROMPT } from './constants'
-import { type CodeToReplaceData, getCompletionsPromptWithSystemPrompt } from './prompt-utils'
+import { type CurrentFilePromptComponents, getCompletionsPromptWithSystemPrompt } from './prompt-utils'
 
-export interface UserPromptArgs {
-    docContext: DocumentContext
-    document: vscode.TextDocument
-    position: vscode.Position
+export interface UserPromptArgs
+    extends Pick<CurrentFilePromptComponents, 'areaPrompt' | 'fileWithMarkerPrompt'> {
     context: AutocompleteContextSnippet[]
     tokenBudget: AutoEditsTokenLimit
 }
 
-export interface UserPromptResponse {
-    codeToReplaceData: CodeToReplaceData
-    prompt: PromptString
-}
-
 export interface UserPromptForModelArgs extends UserPromptArgs {
     isChatModel: boolean
 }
 
-export interface UserPromptForModelResponse {
-    codeToReplaceData: CodeToReplaceData
-    prompt: AutoeditsPrompt
-}
-
 /**
  * Class for generating user prompts in auto-edits functionality.
  * The major difference between different strategy is the prompt rendering.
  */
 export abstract class AutoeditsUserPromptStrategy {
-    protected abstract getUserPrompt(args: UserPromptArgs): UserPromptResponse
+    protected abstract getUserPrompt(args: UserPromptArgs): PromptString
 
     public getPromptForModelType({
         isChatModel,
         ...userPromptArgs
-    }: UserPromptForModelArgs): UserPromptForModelResponse {
-        const { codeToReplaceData, prompt } = this.getUserPrompt(userPromptArgs)
+    }: UserPromptForModelArgs): AutoeditsPrompt {
+        const prompt = this.getUserPrompt(userPromptArgs)
 
         const adjustedPrompt: AutoeditsPrompt = isChatModel
             ? { systemMessage: SYSTEM_PROMPT, userMessage: prompt }
             : { userMessage: getCompletionsPromptWithSystemPrompt(SYSTEM_PROMPT, prompt) }
 
-        return {
-            codeToReplaceData,
-            prompt: adjustedPrompt,
-        }
+        return adjustedPrompt
     }
 }

vscode/src/autoedits/prompt/default-prompt-strategy.test.ts

@@ -7,6 +7,7 @@ import { getCurrentDocContext } from '../../completions/get-current-doc-context'
 import { documentAndPosition } from '../../completions/test-helpers'
 import type { UserPromptArgs } from './base'
 import { DefaultUserPromptStrategy } from './default-prompt-strategy'
+import { getCurrentFilePromptComponents } from './prompt-utils'
 
 describe('DefaultUserPromptStrategy', () => {
     const promptProvider = new DefaultUserPromptStrategy()
@@ -42,7 +43,6 @@ describe('DefaultUserPromptStrategy', () => {
             maxPrefixLength: 100,
             maxSuffixLength: 100,
         })
-
         const tokenBudget: AutoEditsTokenLimit = {
             prefixTokens: 10,
             suffixTokens: 10,
@@ -58,6 +58,14 @@ describe('DefaultUserPromptStrategy', () => {
                 [RetrieverIdentifier.DiagnosticsRetriever]: 100,
             },
         }
+
+        const { fileWithMarkerPrompt, areaPrompt } = getCurrentFilePromptComponents({
+            docContext,
+            document,
+            position,
+            tokenBudget,
+        })
+
         const context: AutocompleteContextSnippet[] = shouldIncludeContext
             ? [
                   getContextItem(
@@ -164,17 +172,16 @@ describe('DefaultUserPromptStrategy', () => {
             : []
 
         return {
-            docContext,
-            document,
-            position,
             context,
+            fileWithMarkerPrompt,
+            areaPrompt,
             tokenBudget,
         }
     }
 
     it('creates prompt with the context source with context', () => {
         const userPromptData = getUserPromptData({ shouldIncludeContext: true })
-        const { prompt } = promptProvider.getUserPrompt(userPromptData)
+        const prompt = promptProvider.getUserPrompt(userPromptData)
 
         expect(prompt.toString()).toEqual(dedent`
             Help me finish a coding change. In particular, you will see a series of snippets from current open files in my editor, files I have recently viewed, the file I am editing, then a history of my recent codebase changes, then current compiler and linter errors, content I copied from my codebase. You will then rewrite the <code_to_rewrite>, to match what you think I would do next in the codebase. Note: I might have stopped in the middle of typing.
@@ -304,7 +311,7 @@ describe('DefaultUserPromptStrategy', () => {
 
     it('creates a prompt in the correct format', () => {
         const userPromptData = getUserPromptData({ shouldIncludeContext: false })
-        const { prompt } = promptProvider.getUserPrompt(userPromptData)
+        const prompt = promptProvider.getUserPrompt(userPromptData)
 
         const expectedPrompt = dedent`
             Help me finish a coding change. In particular, you will see a series of snippets from current open files in my editor, files I have recently viewed, the file I am editing, then a history of my recent codebase changes, then current compiler and linter errors, content I copied from my codebase. You will then rewrite the <code_to_rewrite>, to match what you think I would do next in the codebase. Note: I might have stopped in the middle of typing.

vscode/src/autoedits/prompt/default-prompt-strategy.ts

@@ -1,44 +1,33 @@
-import { ps } from '@sourcegraph/cody-shared'
+import { type PromptString, ps } from '@sourcegraph/cody-shared'
 
 import { RetrieverIdentifier } from '../../completions/context/utils'
 import { autoeditsOutputChannelLogger } from '../output-channel-logger'
 
-import { AutoeditsUserPromptStrategy, type UserPromptArgs, type UserPromptResponse } from './base'
+import { AutoeditsUserPromptStrategy, type UserPromptArgs } from './base'
 import * as constants from './constants'
 import {
     getContextItemMappingWithTokenLimit,
-    getCurrentFilePromptComponents,
     getJaccardSimilarityPrompt,
     getLintErrorsPrompt,
     getPromptForTheContextSource,
     getPromptWithNewline,
     getRecentCopyPrompt,
     getRecentEditsPrompt,
     getRecentlyViewedSnippetsPrompt,
-    joinPromptsWithNewlineSeperator,
+    joinPromptsWithNewlineSeparator,
 } from './prompt-utils'
 
 export class DefaultUserPromptStrategy extends AutoeditsUserPromptStrategy {
     getUserPrompt({
-        docContext,
-        document,
-        position,
         context,
         tokenBudget,
-    }: UserPromptArgs): UserPromptResponse {
+        fileWithMarkerPrompt,
+        areaPrompt,
+    }: UserPromptArgs): PromptString {
         const contextItemMapping = getContextItemMappingWithTokenLimit(
             context,
             tokenBudget.contextSpecificTokenLimit
         )
-        const { fileWithMarkerPrompt, areaPrompt, codeToReplaceData } = getCurrentFilePromptComponents({
-            docContext,
-            document,
-            position,
-            maxPrefixLinesInArea: tokenBudget.maxPrefixLinesInArea,
-            maxSuffixLinesInArea: tokenBudget.maxSuffixLinesInArea,
-            codeToRewritePrefixLines: tokenBudget.codeToRewritePrefixLines,
-            codeToRewriteSuffixLines: tokenBudget.codeToRewriteSuffixLines,
-        })
         const recentViewsPrompt = getPromptForTheContextSource(
             contextItemMapping.get(RetrieverIdentifier.RecentViewPortRetriever) || [],
             constants.LONG_TERM_SNIPPET_VIEWS_INSTRUCTION,
@@ -71,7 +60,7 @@ export class DefaultUserPromptStrategy extends AutoeditsUserPromptStrategy {
 
         const currentFilePrompt = ps`${constants.CURRENT_FILE_INSTRUCTION}${fileWithMarkerPrompt}`
 
-        const finalPrompt = joinPromptsWithNewlineSeperator(
+        const finalPrompt = joinPromptsWithNewlineSeparator(
             getPromptWithNewline(constants.BASE_USER_PROMPT),
             getPromptWithNewline(jaccardSimilarityPrompt),
             getPromptWithNewline(recentViewsPrompt),
@@ -84,9 +73,6 @@ export class DefaultUserPromptStrategy extends AutoeditsUserPromptStrategy {
         )
 
         autoeditsOutputChannelLogger.logDebug('getUserPrompt', 'Prompt\n', finalPrompt)
-        return {
-            codeToReplaceData,
-            prompt: finalPrompt,
-        }
+        return finalPrompt
     }
 }