Visual Studio CodeからOpen AI ChatGPTを使う拡張機能

以前、テキストエディタのMeryからOpenAI APIをたたくマクロをつくりました。

テキストエディタMeryからOpenAI ChatGPTを使うマクロ - attosci diary

しかし、最近はVSCodeを使うことも多くなりましたので、拡張機能として同様な動作をさせたいと思い、TypeScritバージョンを作りました。環境変数としてOPENAI_API_KEYにOpenAI API Keyを設定する必要があります。

拡張機能の作り方は以下のサイトを参考にしました。使うモデルはコードの中に埋め込んでいるので必要に応じて修正する必要があります。

VSCode Extensions(拡張機能) 自作入門 〜VSCodeにおみくじ機能を追加する〜 #VSCode - Qiita

import * as vscode from 'vscode';
import { OpenAI } from 'openai';

export function activate(context: vscode.ExtensionContext) {

    console.log('Extension "vscode-openai" is now active!');

    const disposable = vscode.commands.registerCommand('vscode-openai.sendToChatGPT', async () => {
        const editor = vscode.window.activeTextEditor;

        if (!editor) {
            vscode.window.showInformationMessage('No active editor!');
            return; // No active editor
        }

        const document = editor.document;
        const selection = editor.selection;
        const text = selection.isEmpty ? document.getText() : document.getText(selection);

        try {
            const response = await sendToChatGPT(text);
            appendTextToDocument(editor, response);
        } catch (error) {
            if (error instanceof Error) {
                // We can now safely read the message property of the Error instance
                vscode.window.showErrorMessage('Error communicating with OpenAI: ' + error.message);
            } else {
                // For other types of thrown values that are not Error instances
                vscode.window.showErrorMessage('An unknown error occurred.');
            }
        }
    });

    context.subscriptions.push(disposable);
}

async function sendToChatGPT(text: string): Promise<string> {
    // Obtain OpenAI API Key from environment variable OPENAI_API_KEY
    const apiKey = process.env.OPENAI_API_KEY;
    if (!apiKey) {
        throw new Error('OpenAI API key is not set in the environment variable OPENAI_API_KEY.');
    }
    const openai = new OpenAI({ apiKey: apiKey });
    const model = vscode.workspace.getConfiguration().get('vscode-openai.model', 'gpt-4-turbo');
    const messages: any[] = makeMessages(text);
    const payload = {
        model: model,
        messages: messages
    };

    const response = await openai.chat.completions.create(payload);
    if (response.choices && response.choices.length > 0 && response.choices[0].message.content) {
        return response.choices[0].message.content;
    } else {
        throw new Error('Invalid response from the API.');
    }
}

function appendTextToDocument(editor: vscode.TextEditor, text: string) {
    editor.edit(editBuilder => {
        const document = editor.document;
        let selection = editor.selection;
        let position: vscode.Position;

        if (!selection.isEmpty) {
            position = selection.end;
        } else {
            const lastLine = document.lineAt(document.lineCount - 1);
            position = new vscode.Position(document.lineCount - 1, lastLine.text.length);
        }

        const assistantMarker = "\n\n### assistant\n\n";
        const userMarker = "\n\n### user\n\n";
        const fullText = assistantMarker + text + userMarker;
        editBuilder.insert(position, fullText);
    });
}

function makeMessages(text: string): any {
    const lines = text.split('\n');

    let messages: any[] = [];
    const roles: string[] = ['system', 'assistant', 'user'];
    let current_role: string = 'user';
    let content: string = '';

    for (const line of lines) {
        let roleHeadline: boolean = false;
        for (const role in roles) {
            if (line.startsWith(`### ${role}`)) {
                if (content.length > 0) {
                    messages.push({ "role": current_role, "content": content});
                }
                current_role = role;
                content = '';
                roleHeadline = true;
                break;
            }
        }
        if (!roleHeadline) {
            content += line + '\n';
        }
    }
    if (content.length > 0) {
        messages.push({ "role": current_role, "content": content});
    }
    return messages;
}

export function deactivate() {}