feat(smart-apply): Add feature flag for internal dogfooding (#7342)

Add the feature flags for the smart apply custom model

## Test plan
Unit test

vscode/src/edit/edit-context-logging.ts

@@ -1,9 +1,11 @@
 import {
     type EditModel,
     FeatureFlag,
+    currentAuthStatus,
     displayPathWithoutWorkspaceFolderPrefix,
     featureFlagProvider,
     isDotComAuthed,
+    isS2,
     storeLastValue,
     telemetryRecorder,
 } from '@sourcegraph/cody-shared'
@@ -235,8 +237,9 @@ export function getEditLoggingContext(param: {
 }
 
 function shouldLogEditContextItem<T>(payload: T, isFeatureFlagEnabledForLogging: boolean): boolean {
-    // 🚨 SECURITY: included only for DotCom users and for users in the feature flag.
-    if (isDotComAuthed() && isFeatureFlagEnabledForLogging) {
+    // 🚨 SECURITY: included only for DotCom or S2 users and for users in the feature flag.
+    const authStatus = currentAuthStatus()
+    if ((isDotComAuthed() || isS2(authStatus)) && isFeatureFlagEnabledForLogging) {
         const payloadSize = calculatePayloadSizeInBytes(payload)
         return payloadSize !== undefined && payloadSize < MAX_LOGGING_PAYLOAD_SIZE_BYTES
     }

vscode/src/edit/prompt/smart-apply/selection.ts

@@ -1,7 +1,9 @@
 import {
     BotResponseMultiplexer,
     type ChatClient,
+    type CompletionParameters,
     type EditModel,
+    type Message,
     type PromptString,
     modelsService,
 } from '@sourcegraph/cody-shared'
@@ -15,14 +17,11 @@ import { SMART_APPLY_TOPICS, type SmartApplySelectionProvider } from './selectio
 import { CustomModelSelectionProvider } from './selection/custom-model'
 import { DefaultSelectionProvider } from './selection/default'
 
-async function promptModelForOriginalCode(
-    selectionProvider: SmartApplySelectionProvider,
-    instruction: PromptString,
-    replacement: PromptString,
-    document: vscode.TextDocument,
-    model: EditModel,
+export async function getSelectionFromModel(
     client: ChatClient,
-    codyApiVersion: number
+    prefix: string,
+    messages: Message[],
+    params: CompletionParameters
 ): Promise<string> {
     const multiplexer = new BotResponseMultiplexer()
 
@@ -37,14 +36,6 @@ async function promptModelForOriginalCode(
     })
 
     const abortController = new AbortController()
-    const { prefix, messages } = await selectionProvider.getPrompt({
-        instruction,
-        replacement,
-        document,
-        model,
-        codyApiVersion,
-    })
-    const params = selectionProvider.getLLMCompletionsParameters()
     const stream = await client.chat(messages, params, abortController.signal)
 
     let textConsumed = 0
@@ -86,15 +77,11 @@ interface SmartApplySelection {
     range: vscode.Range
 }
 
-function getSmartApplySelectionProvider(
-    model: string,
-    replacement: PromptString
-): SmartApplySelectionProvider {
-    const contextWindow = modelsService.getContextWindowByID(model)
+function getSmartApplySelectionProvider(model: EditModel): SmartApplySelectionProvider {
     if (Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(model)) {
-        return new CustomModelSelectionProvider(model, contextWindow, replacement.toString())
+        return new CustomModelSelectionProvider()
     }
-    return new DefaultSelectionProvider(model, contextWindow)
+    return new DefaultSelectionProvider()
 }
 
 export async function getSmartApplySelection({
@@ -115,17 +102,18 @@ export async function getSmartApplySelection({
     codyApiVersion: number
 }): Promise<SmartApplySelection | null> {
     let originalCode: string
-    const selectionProvider = getSmartApplySelectionProvider(model, replacement)
+    const selectionProvider = getSmartApplySelectionProvider(model)
+    const contextWindow = modelsService.getContextWindowByID(model)
     try {
-        originalCode = await promptModelForOriginalCode(
-            selectionProvider,
+        originalCode = await selectionProvider.getSelectedText({
             instruction,
             replacement,
             document,
             model,
+            contextWindow,
             chatClient,
-            codyApiVersion
-        )
+            codyApiVersion,
+        })
     } catch (error: unknown) {
         // We erred when asking the LLM to produce the original code.
         // Surface this error back to the user

vscode/src/edit/prompt/smart-apply/selection/base.ts

@@ -1,7 +1,8 @@
 import {
-    type CompletionParameters,
+    type ChatClient,
     type EditModel,
     type Message,
+    type ModelContextWindow,
     type PromptString,
     ps,
 } from '@sourcegraph/cody-shared'
@@ -12,7 +13,9 @@ export interface SelectionPromptProviderArgs {
     replacement: PromptString
     document: vscode.TextDocument
     model: EditModel
+    contextWindow: ModelContextWindow
     codyApiVersion: number
+    chatClient: ChatClient
 }
 
 export interface SelectionPromptProviderResult {
@@ -21,9 +24,7 @@ export interface SelectionPromptProviderResult {
 }
 
 export interface SmartApplySelectionProvider {
-    getPrompt(args: SelectionPromptProviderArgs): Promise<SelectionPromptProviderResult>
-
-    getLLMCompletionsParameters(): CompletionParameters
+    getSelectedText(args: SelectionPromptProviderArgs): Promise<string>
 }
 
 export const SMART_APPLY_TOPICS = {

vscode/src/edit/prompt/smart-apply/selection/custom-model.ts

@@ -2,13 +2,12 @@ import {
     type CompletionParameters,
     type EditModel,
     type Message,
-    type ModelContextWindow,
     PromptString,
     TokenCounterUtils,
-    modelsService,
     psDedent,
 } from '@sourcegraph/cody-shared'
 import * as vscode from 'vscode'
+import { getSelectionFromModel } from '../selection'
 import { getInstructionPromptWithCharLimit } from '../utils'
 import {
     LLM_PARAMETERS,
@@ -18,7 +17,8 @@ import {
     type SmartApplySelectionProvider,
 } from './base'
 
-export const SMART_APPLY_INSTRUCTION_TOKEN_LIMIT = 500
+const SMART_APPLY_INSTRUCTION_TOKEN_LIMIT = 500
+const FULL_FILE_REWRITE_TOKEN_TOKEN_LIMIT = 12000
 
 const DEFAULT_SELECTION_PROMPT = {
     system: psDedent`
@@ -44,31 +44,50 @@ const DEFAULT_SELECTION_PROMPT = {
 }
 
 export class CustomModelSelectionProvider implements SmartApplySelectionProvider {
-    private model: EditModel
-    private contextWindow: ModelContextWindow
-    private replacement: string
-
-    constructor(model: EditModel, contextWindow: ModelContextWindow, replacement: string) {
-        this.model = model
-        this.contextWindow = contextWindow
-        this.replacement = replacement
-    }
-
-    async getPrompt({
+    public async getSelectedText({
         instruction,
         replacement,
         document,
         model,
-    }: SelectionPromptProviderArgs): Promise<SelectionPromptProviderResult> {
+        chatClient,
+        contextWindow,
+    }: SelectionPromptProviderArgs): Promise<string> {
         const documentRange = new vscode.Range(0, 0, document.lineCount - 1, 0)
         const documentText = PromptString.fromDocumentText(document, documentRange)
         const tokenCount = await TokenCounterUtils.countPromptString(documentText)
 
-        const contextWindow = modelsService.getContextWindowByID(model)
         if (tokenCount > contextWindow.input) {
             throw new Error("The amount of text in this document exceeds Cody's current capacity.")
         }
+        if (tokenCount < FULL_FILE_REWRITE_TOKEN_TOKEN_LIMIT) {
+            return 'ENTIRE_FILE'
+        }
+        const { prefix, messages } = await this.getPrompt(
+            instruction,
+            replacement,
+            document,
+            documentText
+        )
+        const completionParameters = this.getLLMCompletionsParameters(
+            model,
+            contextWindow.output,
+            replacement.toString()
+        )
+        const selectedText = await getSelectionFromModel(
+            chatClient,
+            prefix,
+            messages,
+            completionParameters
+        )
+        return selectedText
+    }
 
+    private async getPrompt(
+        instruction: PromptString,
+        replacement: PromptString,
+        document: vscode.TextDocument,
+        fileContent: PromptString
+    ): Promise<SelectionPromptProviderResult> {
         const instructionPromptWithLimit = getInstructionPromptWithCharLimit(
             instruction,
             SMART_APPLY_INSTRUCTION_TOKEN_LIMIT
@@ -78,7 +97,7 @@ export class CustomModelSelectionProvider implements SmartApplySelectionProvider
         const userPrompt = DEFAULT_SELECTION_PROMPT.instruction
             .replaceAll('{instruction}', instructionPromptWithLimit)
             .replaceAll('{incomingText}', replacement)
-            .replaceAll('{fileContents}', documentText)
+            .replaceAll('{fileContents}', fileContent)
             .replaceAll('{filePath}', PromptString.fromDisplayPath(document.uri))
 
         const prompt: Message[] = [
@@ -91,16 +110,20 @@ export class CustomModelSelectionProvider implements SmartApplySelectionProvider
         }
     }
 
-    getLLMCompletionsParameters(): CompletionParameters {
+    getLLMCompletionsParameters(
+        model: EditModel,
+        outputTokens: number,
+        replacement: string
+    ): CompletionParameters {
         return {
-            model: this.model,
+            model,
             stopSequences: LLM_PARAMETERS.stopSequences,
-            maxTokensToSample: this.contextWindow.output,
+            maxTokensToSample: outputTokens,
             temperature: 0.1,
             stream: true,
             prediction: {
                 type: 'content',
-                content: this.replacement,
+                content: replacement,
             },
             rewriteSpeculation: true,
             adaptiveSpeculation: true,

vscode/src/edit/prompt/smart-apply/selection/default.ts

@@ -6,11 +6,11 @@ import {
     PromptString,
     TokenCounterUtils,
     getSimplePreamble,
-    modelsService,
     psDedent,
 } from '@sourcegraph/cody-shared'
 import * as vscode from 'vscode'
 import { PromptBuilder } from '../../../../prompt-builder'
+import { getSelectionFromModel } from '../selection'
 import {
     LLM_PARAMETERS,
     SMART_APPLY_TOPICS,
@@ -52,29 +52,51 @@ const DEFAULT_SELECTION_PROMPT = {
 }
 
 export class DefaultSelectionProvider implements SmartApplySelectionProvider {
-    private model: EditModel
-    private contextWindow: ModelContextWindow
-
-    constructor(model: EditModel, contextWindow: ModelContextWindow) {
-        this.model = model
-        this.contextWindow = contextWindow
-    }
-
-    async getPrompt({
+    public async getSelectedText({
         instruction,
         replacement,
         document,
         model,
+        chatClient,
         codyApiVersion,
-    }: SelectionPromptProviderArgs): Promise<SelectionPromptProviderResult> {
+        contextWindow,
+    }: SelectionPromptProviderArgs): Promise<string> {
         const documentRange = new vscode.Range(0, 0, document.lineCount - 1, 0)
         const documentText = PromptString.fromDocumentText(document, documentRange)
         const tokenCount = await TokenCounterUtils.countPromptString(documentText)
-        const contextWindow = modelsService.getContextWindowByID(model)
+
         if (tokenCount > contextWindow.input) {
             throw new Error("The amount of text in this document exceeds Cody's current capacity.")
         }
 
+        const { prefix, messages } = await this.getPrompt(
+            instruction,
+            replacement,
+            document,
+            model,
+            codyApiVersion,
+            documentText,
+            contextWindow
+        )
+        const completionParameters = this.getLLMCompletionsParameters(model, contextWindow.output)
+        const selectedText = await getSelectionFromModel(
+            chatClient,
+            prefix,
+            messages,
+            completionParameters
+        )
+        return selectedText
+    }
+
+    private async getPrompt(
+        instruction: PromptString,
+        replacement: PromptString,
+        document: vscode.TextDocument,
+        model: EditModel,
+        codyApiVersion: number,
+        fileContent: PromptString,
+        contextWindow: ModelContextWindow
+    ): Promise<SelectionPromptProviderResult> {
         const promptBuilder = await PromptBuilder.create(contextWindow)
         const preamble = getSimplePreamble(
             model,
@@ -87,7 +109,7 @@ export class DefaultSelectionProvider implements SmartApplySelectionProvider {
         const text = DEFAULT_SELECTION_PROMPT.instruction
             .replaceAll('{instruction}', instruction)
             .replaceAll('{incomingText}', replacement)
-            .replaceAll('{fileContents}', documentText)
+            .replaceAll('{fileContents}', fileContent)
             .replaceAll('{filePath}', PromptString.fromDisplayPath(document.uri))
 
         const transcript: ChatMessage[] = [{ speaker: 'human', text }]
@@ -101,11 +123,11 @@ export class DefaultSelectionProvider implements SmartApplySelectionProvider {
         }
     }
 
-    getLLMCompletionsParameters(): CompletionParameters {
+    private getLLMCompletionsParameters(model: EditModel, outputTokens: number): CompletionParameters {
         return {
-            model: this.model,
+            model,
             stopSequences: LLM_PARAMETERS.stopSequences,
-            maxTokensToSample: this.contextWindow.output,
+            maxTokensToSample: outputTokens,
         } as CompletionParameters
     }
 }