Commit 21527124 authored by Fred de Gier's avatar Fred de Gier 🛑 Committed by Stanislav Lashmanov
Browse files

feat: ai assist stop sequences

parent 7c001b8b

src/completion/gitlab_code_completion_provider.ts

+7 −0
Original line number Diff line number Diff line
@@ -4,6 +4,7 @@ import { log } from '../log'; 
import { getAiAssistConfiguration, AiAssistConfiguration } from '../utils/extension_configuration';

import { getActiveProject } from '../commands/run_with_valid_project';

import { tokenExchangeService } from '../gitlab/token_exchange_service';

import { getStopSequences } from '../utils/ai_assist/get_stop_sequences';



export class GitLabCodeCompletionProvider implements vscode.InlineCompletionItemProvider {

  engine: string;
@@ -58,7 +59,12 @@ export class GitLabCodeCompletionProvider implements vscode.InlineCompletionItem 
  async provideInlineCompletionItems(

    document: vscode.TextDocument,

    position: vscode.Position,

    context: vscode.InlineCompletionContext,

  ): Promise<vscode.InlineCompletionItem[]> {

    if (context.triggerKind === 0) {

      //  TODO: In case of a hover, this will be triggered which is not desired as it calls for a new prediction

    }



    const prompt = document.getText(new vscode.Range(0, 0, position.line, position.character));



    // TODO: Sanitize prompt to prevent exposing sensitive information
@@ -78,6 +84,7 @@ export class GitLabCodeCompletionProvider implements vscode.InlineCompletionItem 
    const response = await oa.createCompletion({

      model: this.model,

      prompt: prompt as openai.CreateCompletionRequestPrompt,

      stop: getStopSequences(position.line, document),

    });



    const completions =

src/utils/ai_assist/get_stop_sequences.test.ts

0 → 100644
+36 −0
Original line number Diff line number Diff line 
import * as vscode from 'vscode';

import { getStopSequences } from './get_stop_sequences';



const TextDocument = class {

  constructor(text: string) {

    this.text = text;

    this.lineCount = text.split('\n').length;

  }



  text: any;



  lineCount: number;



  public lineAt(position: number) {

    return {

      text: this.text.split('\n')[position],

    };

  }

};



describe('getStopSequences', () => {

  it.each([

    ['import pandas as pd\nimport numpy as np', 0, 2, ['import numpy as np']],

    ['import pandas as pd\nimport numpy as np', 1, 2, []],

    ['import pandas as pd\nimport numpy as np', 2, 2, []],

    ['import pandas as pd\nimport numpy as np\nimport os', 1, 3, ['import os']],

  ])(

    'gets stop sequences',

    (text: string, position: number, lineCount: number, stopSequence: string[]) => {

      const doc = new TextDocument(text);

      const stopSequences = getStopSequences(position, doc as unknown as vscode.TextDocument);

      expect(doc.lineCount).toBe(lineCount);

      expect(stopSequences).toStrictEqual(stopSequence);

    },

  );

});

src/utils/ai_assist/get_stop_sequences.ts

0 → 100644
+18 −0
Original line number Diff line number Diff line 
import * as vscode from 'vscode';

import { log } from '../../log';



export function getStopSequences(position: number, doc: vscode.TextDocument) {

  const stopSequences: string[] = [];



  // If the position is the last line, we don't use stop sequences

  if (doc.lineCount === position + 1) {

    return stopSequences;

  }

  // If the nextLine is not empty, it will be used as a stop sequence

  const nextLine = doc.lineAt(position + 1).text;

  if (nextLine) {

    stopSequences.push(nextLine.toString());

  }

  log.debug(`Stop sequences: ${stopSequences.toString()}`);

  return stopSequences;

}