Hitesh/autoedits improvements

vscode/package.json

@@ -1441,7 +1441,8 @@
     "uuid": "^9.0.0",
     "vscode-languageserver-textdocument": "^1.0.8",
     "vscode-uri": "^3.0.8",
-    "wink-nlp-utils": "^2.1.0"
+    "wink-nlp-utils": "^2.1.0",
+    "fast-myers-diff": "^3.2.0"
   },
   "devDependencies": {
     "7zip-min": "^1.4.5",

vscode/src/autoedits/autoedits-provider.ts

@@ -1,9 +1,4 @@
-import {
-    type AutoEditsTokenLimit,
-    type DocumentContext,
-    logDebug,
-    tokensToChars,
-} from '@sourcegraph/cody-shared'
+import { type AutoEditsTokenLimit, type DocumentContext, tokensToChars } from '@sourcegraph/cody-shared'
 import { Observable } from 'observable-fns'
 import * as vscode from 'vscode'
 import { ContextMixer } from '../completions/context/context-mixer'
@@ -12,6 +7,7 @@ import { getCurrentDocContext } from '../completions/get-current-doc-context'
 import { getConfiguration } from '../configuration'
 import type { PromptProvider } from './prompt-provider'
 import { DeepSeekPromptProvider } from './providers/deepseek'
+import { FireworksPromptProvider } from './providers/fireworks'
 import { OpenAIPromptProvider } from './providers/openai'
 import { AutoEditsRenderer } from './renderer'
 
@@ -33,11 +29,13 @@ export class AutoeditsProvider implements vscode.Disposable {
     private model: string | undefined
     private apiKey: string | undefined
     private renderer: AutoEditsRenderer = new AutoEditsRenderer()
+    private outputChannel: vscode.OutputChannel
 
     constructor() {
+        this.outputChannel = vscode.window.createOutputChannel('Autoedit Testing')
         const config = getConfiguration().experimentalAutoedits
         if (config === undefined) {
-            logDebug('AutoEdits', 'No Configuration found in the settings')
+            this.logDebug('AutoEdits', 'No Configuration found in the settings')
             return
         }
         this.initizlizePromptProvider(config.provider)
@@ -50,14 +48,19 @@ export class AutoeditsProvider implements vscode.Disposable {
             vscode.commands.registerCommand('cody.experimental.suggest', () => this.getAutoedit())
         )
     }
+    private logDebug(provider: string, ...args: unknown[]): void {
+        this.outputChannel.appendLine(`${provider} █| ${args.join('')}`)
+    }
 
     private initizlizePromptProvider(provider: string) {
         if (provider === 'openai') {
             this.provider = new OpenAIPromptProvider()
         } else if (provider === 'deepseek') {
             this.provider = new DeepSeekPromptProvider()
+        } else if (provider === 'fireworks') {
+            this.provider = new FireworksPromptProvider()
         } else {
-            logDebug('AutoEdits', `provider ${provider} not supported`)
+            this.logDebug('AutoEdits', `provider ${provider} not supported`)
         }
     }
 
@@ -70,7 +73,7 @@ export class AutoeditsProvider implements vscode.Disposable {
 
     public async predictAutoeditAtDocAndPosition(options: AutoEditsProviderOptions) {
         if (!this.provider || !this.autoEditsTokenLimit || !this.model || !this.apiKey) {
-            logDebug('AutoEdits', 'No Provider or Token Limit found in the settings')
+            this.logDebug('AutoEdits', 'No Provider or Token Limit found in the settings')
             return
         }
         const start = Date.now()
@@ -88,9 +91,16 @@ export class AutoeditsProvider implements vscode.Disposable {
             this.autoEditsTokenLimit
         )
         const response = await this.provider.getModelResponse(this.model, this.apiKey, prompt)
+        const postProcessedResponse = this.provider.postProcessResponse(codeToReplace, response)
+        this.logDebug('Autoedits', '========================== Response:\n', postProcessedResponse, '\n')
         const timeToResponse = Date.now() - start
-        logDebug('AutoEdits: (Time LLM Query):', timeToResponse.toString())
-        await this.renderer.render(options, codeToReplace, response)
+        this.logDebug(
+            'Autoedits',
+            '========================== Time Taken For LLM (Msec): ',
+            timeToResponse.toString(),
+            '\n'
+        )
+        await this.renderer.render(options, codeToReplace, postProcessedResponse)
     }
 
     private getDocContext(document: vscode.TextDocument, position: vscode.Position): DocumentContext {

vscode/src/autoedits/prompt-provider.ts

@@ -25,7 +25,7 @@ export interface PromptProvider {
         tokenBudget: AutoEditsTokenLimit
     ): PromptResponseData
 
-    postProcessResponse(completion: string | null): string
+    postProcessResponse(codeToReplace: utils.CodeToReplaceData, completion: string | null): string
 
     getModelResponse(model: string, apiKey: string, prompt: PromptProviderResponse): Promise<string>
 }

vscode/src/autoedits/prompt-utils.ts

@@ -184,7 +184,7 @@ export function getCurrentFilePromptComponents(
     const codeToReplace = {
         codeToRewrite: codeToRewrite,
         startLine: completePrefixLines - prefixContext.codeToRewriteStartLines,
-        endLine: completePrefixLines + suffixContext.codeToRewriteEndLines,
+        endLine: completePrefixLines + suffixContext.codeToRewriteEndLines - 1,
     }
 
     const fileWithMarker = ps`${prefixContext.prefixBeforeArea}

vscode/src/autoedits/providers/deepseek.ts

@@ -6,7 +6,7 @@ import type {
 } from '../../../../lib/shared/src/completions/types'
 import type { PromptProvider, PromptProviderResponse, PromptResponseData } from '../prompt-provider'
 import { getModelResponse } from '../prompt-provider'
-import { SYSTEM_PROMPT, getBaseUserPrompt } from '../prompt-utils'
+import { type CodeToReplaceData, SYSTEM_PROMPT, getBaseUserPrompt } from '../prompt-utils'
 
 export class DeepSeekPromptProvider implements PromptProvider {
     private readonly bosToken: PromptString = ps`<｜begin▁of▁sentence｜>`
@@ -37,7 +37,7 @@ ${this.assistantToken}`
         }
     }
 
-    postProcessResponse(response: string): string {
+    postProcessResponse(codeToReplace: CodeToReplaceData, response: string): string {
         return response
     }
 

vscode/src/autoedits/providers/fireworks.ts

@@ -0,0 +1,84 @@
+import { type AutoEditsTokenLimit, logDebug } from '@sourcegraph/cody-shared'
+import type * as vscode from 'vscode'
+import type {
+    AutocompleteContextSnippet,
+    DocumentContext,
+} from '../../../../lib/shared/src/completions/types'
+import { lines } from '../../completions/text-processing'
+import type {
+    ChatPrompt,
+    PromptProvider,
+    PromptProviderResponse,
+    PromptResponseData,
+} from '../prompt-provider'
+import { getModelResponse } from '../prompt-provider'
+import { type CodeToReplaceData, SYSTEM_PROMPT, getBaseUserPrompt } from '../prompt-utils'
+
+export class FireworksPromptProvider implements PromptProvider {
+    getPrompt(
+        docContext: DocumentContext,
+        document: vscode.TextDocument,
+        context: AutocompleteContextSnippet[],
+        tokenBudget: AutoEditsTokenLimit
+    ): PromptResponseData {
+        const { codeToReplace, promptResponse: userPrompt } = getBaseUserPrompt(
+            docContext,
+            document,
+            context,
+            tokenBudget
+        )
+        const prompt: ChatPrompt = [
+            {
+                role: 'system',
+                content: SYSTEM_PROMPT,
+            },
+            {
+                role: 'user',
+                content: userPrompt,
+            },
+        ]
+        return {
+            codeToReplace: codeToReplace,
+            promptResponse: prompt,
+        }
+    }
+
+    postProcessResponse(codeToReplace: CodeToReplaceData, response: string): string {
+        // todo (hitesh): The finetuned model is messing up the identation of the first line.
+        // todo: correct it manully for now, by checking the first line of the code to rewrite and adding the same indentation to the first line of the completion
+        const codeToRewrite = codeToReplace.codeToRewrite.toString()
+        const codeToRewriteLines = lines(codeToRewrite)
+        const completionLines = lines(response)
+        const firstLineMatch = codeToRewriteLines[0].match(/^(\s*)/)
+        const firstLineIndentation = firstLineMatch ? firstLineMatch[1] : ''
+        completionLines[0] = firstLineIndentation + completionLines[0].trimLeft()
+        const completion = completionLines.join('\n')
+        return completion
+    }
+
+    async getModelResponse(
+        model: string,
+        apiKey: string,
+        prompt: PromptProviderResponse
+    ): Promise<string> {
+        try {
+            const response = await getModelResponse(
+                'https://sourcegraph-6c39ed29.direct.fireworks.ai/v1/chat/completions',
+                JSON.stringify({
+                    model: model,
+                    messages: prompt,
+                    temperature: 0.1,
+                    max_tokens: 256,
+                    response_format: {
+                        type: 'text',
+                    },
+                }),
+                apiKey
+            )
+            return response.choices[0].message.content
+        } catch (error) {
+            logDebug('AutoEdits', 'Error calling OpenAI API:', error)
+            throw error
+        }
+    }
+}

vscode/src/autoedits/providers/openai.ts

@@ -11,7 +11,7 @@ import type {
     PromptResponseData,
 } from '../prompt-provider'
 import { getModelResponse } from '../prompt-provider'
-import { SYSTEM_PROMPT, getBaseUserPrompt } from '../prompt-utils'
+import { type CodeToReplaceData, SYSTEM_PROMPT, getBaseUserPrompt } from '../prompt-utils'
 
 export class OpenAIPromptProvider implements PromptProvider {
     getPrompt(
@@ -42,7 +42,7 @@ export class OpenAIPromptProvider implements PromptProvider {
         }
     }
 
-    postProcessResponse(response: string): string {
+    postProcessResponse(codeToReplace: CodeToReplaceData, response: string): string {
         return response
     }