copilot_suggestion.py: GitHub Copilot LSP Integration in Python

copilot_suggestion.py

import subprocess
import os
import time
import pathlib

import pylspclient
from pylspclient import lsp_structs


def get_uri_from_path(path: str) -> str:
    return 'file://' + path


def get_current_directory_uri() -> str:
    current_directory = pathlib.Path().absolute()
    return get_uri_from_path(str(current_directory))


def get_absolute_path_from_relative(relative_path: str) -> str:
    return os.path.join(os.path.dirname(__file__), relative_path)


def start_copilot_agent() -> subprocess.Popen:
    node_executable = 'node'
    copilot_base_dir = '.'
    copilot_agent_command = [node_executable, os.path.join(copilot_base_dir, "dist/agent.js")]

    # Starting the Copilot agent process
    return subprocess.Popen(copilot_agent_command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=False)


class Doc(object):
    def __init__(self, uri: str, languageId: str, position: lsp_structs.Position, version: int):
        self.uri = uri
        self.languageId = languageId
        self.position = position
        self.version = version


class Completion(object):
    def __init__(self, uuid: str, text: str, range: dict, displayText: str, position: dict, docVersion: int):
        self.uuid = uuid
        self.text = text
        self.range = lsp_structs.Range(**range),
        self.displayText = displayText,
        self.position = lsp_structs.Position(**position)
        self.docVersion = docVersion

    def __str__(self) -> str:
        return self.text


def get_suggestion(client: pylspclient.LspClient, doc: Doc) -> list:
    try:
        result_dict = client.lsp_endpoint.call_method(
            "getCompletions", 
            doc=doc
        )
    except ValueError as e:
        print("Error:", e)
        return []

    if  "cancellationReason" in result_dict:
        print("Cancellation reason:", result_dict["cancellationReason"], "\n")

    if "completions" in result_dict:
        return [Completion(**completion) for completion in result_dict["completions"]]
    
    return []


# copied from https://github.com/copilot-emacs/copilot.el
DEFAULT_CAPABILITIES = {
    'textDocument': {
        'completion': {
            'completionItem': {
                'commitCharactersSupport': True,
                'documentationFormat': ['markdown', 'plaintext'],
                'snippetSupport': True
            }
        }
    }
}


copilot_agent_process = start_copilot_agent()
time.sleep(1) # Wait for the agent to start
json_rpc_endpoint = pylspclient.JsonRpcEndpoint(copilot_agent_process.stdin, copilot_agent_process.stdout)
lsp_endpoint = pylspclient.LspEndpoint(
    json_rpc_endpoint,
    notify_callbacks={
        "featureFlagsNotification": lambda params: print("Feature flags:", params, "\n"),
        "LogMessage": lambda params: print("Log message:", params, "\n"),
        "statusNotification": lambda params: print("Status:", params, "\n"),
    }
)
lsp_client = pylspclient.LspClient(lsp_endpoint=lsp_endpoint)
initialize_response = lsp_client.initialize(
    processId=copilot_agent_process.pid,
    rootPath='.',
    rootUri=get_current_directory_uri(),
    initializationOptions=None,
    capabilities=DEFAULT_CAPABILITIES,
    trace="off",
    workspaceFolders=None
)
print("Initialization response:", initialize_response, "\n")
lsp_client.initialized()

opened_document_path = get_absolute_path_from_relative('fibonacci.swift')
opened_document_text = open(opened_document_path).read()
opened_document = lsp_structs.TextDocumentItem(
    uri=get_uri_from_path(opened_document_path),
    languageId=lsp_structs.LANGUAGE_IDENTIFIER.SWIFT,
    version=0,
    text=opened_document_text
)
did_open_response = lsp_client.didOpen(opened_document)
print("Did open response:", did_open_response, "\n")
last_line_index = len(opened_document_text.splitlines()) - 1
last_line = opened_document_text.splitlines()[-1]
last_column_index = len(last_line)
doc = Doc(
    uri=opened_document.uri,
    languageId=lsp_structs.LANGUAGE_IDENTIFIER.PYTHON,
    position=lsp_structs.Position(line=last_line_index, character=last_column_index),
    version=0
)
suggestion_response = get_suggestion(
    client=lsp_client,
    doc=doc
)
print("Suggestion response:", ", ".join(map(str, suggestion_response)), "\n")

# to-do: didChange, didFocus, didClose

lsp_client.shutdown()
lsp_client.exit()

fibonacci.swift

func fibonacci(n: Int) -> Int

output.txt

Log message: {'level': 0, 'message': '[DEBUG] [agent] [2024-01-26T13:30:10.383Z] Agent service starting', 'metadataStr': '[DEBUG] [agent] [2024-01-26T13:30:10.383Z]', 'extra': ['Agent service starting']} 

Initialization response: {'capabilities': {'textDocumentSync': {'openClose': True, 'change': 2}, 'workspace': {'workspaceFolders': {'supported': False, 'changeNotifications': False}}}} 

Did open response: None 

Feature flags: {'ssc': False, 'rt': True, 'sn': False, 'chat': True, 'chatjb': False} 

Log message: {'level': 0, 'message': '[DEBUG] [getCompletions] [2024-01-26T13:30:11.061Z] Requesting completion at position 0:29, between "func fibonacci(n: Int) -> Int" and "".', 'metadataStr': '[DEBUG] [getCompletions] [2024-01-26T13:30:11.061Z]', 'extra': ['Requesting completion at position 0:29, between "func fibonacci(n: Int) -> Int" and "".']} 

Status: {'status': 'InProgress', 'message': ''} 

Log message: {'level': 1, 'message': '[INFO] [default] [2024-01-26T13:30:11.489Z] [fetchCompletions] engine https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex', 'metadataStr': '[INFO] [default] [2024-01-26T13:30:11.489Z]', 'extra': ['[fetchCompletions] engine https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex']}

Log message: {'level': 1, 'message': '[INFO] [default] [2024-01-26T13:30:12.166Z] request.response: [https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex/completions] took 672 ms', 'metadataStr': '[INFO] [default] [2024-01-26T13:30:12.166Z]', 'extra': ['request.response: [https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex/completions] took 672 ms']}

Log message: {'level': 1, 'message': '[INFO] [streamChoices] [2024-01-26T13:30:12.175Z] solution 0 returned. finish reason: [stop]', 'metadataStr': '[INFO] [streamChoices] [2024-01-26T13:30:12.175Z]', 'extra': ['solution 0 returned. finish reason: [stop]']}

Status: {'status': 'Normal', 'message': ''}

Log message: {'level': 1, 'message': '[INFO] [streamChoices] [2024-01-26T13:30:12.180Z] request done: headerRequestId: [885cbe1b-d00f-4f5d-9eba-039eb1472cee] model deployment ID: [w9bda9c6b4b91]', 'metadataStr': '[INFO] [streamChoices] [2024-01-26T13:30:12.180Z]', 'extra': ['request done: headerRequestId: [885cbe1b-d00f-4f5d-9eba-039eb1472cee] model deployment ID: [w9bda9c6b4b91]']}

Suggestion response: func fibonacci(n: Int) -> Int {
    if n == 0 || n == 1 {
        return n
    }
    return fibonacci(n: n - 1) + fibonacci(n: n - 2)
}