chat_handler.py

import time
from datetime import datetime
from typing import Optional, cast
import os

import CopilotChat.prompts as system_prompts
from CopilotChat.copilot import Copilot
from CopilotChat.mypynvim.core.buffer import MyBuffer
from CopilotChat.mypynvim.core.nvim import MyNvim


def is_module_installed(name):
    try:
        __import__(name)
        return True
    except ImportError:
        return False


# TODO: Abort request if the user closes the layout
class ChatHandler:
    has_show_extra_info = False

    def __init__(self, nvim: MyNvim, buffer: MyBuffer):
        self.nvim: MyNvim = nvim
        self.copilot: Copilot = None
        self.buffer: MyBuffer = buffer
        self.proxy: str = os.getenv("HTTPS_PROXY") or os.getenv("ALL_PROXY") or ""

    # public

    def chat(
        self,
        prompt: str,
        filetype: str,
        code: str = "",
        winnr: int = 0,
        system_prompt: Optional[str] = None,
        disable_start_separator: bool = False,
        disable_end_separator: bool = False,
        model: str = "gpt-4",
    ):
        disable_separators = (
            self.nvim.eval("g:copilot_chat_disable_separators") == "yes"
        )
        self.proxy = self.nvim.eval("g:copilot_chat_proxy")
        if "://" not in self.proxy:
            self.proxy = None

        if system_prompt is None:
            system_prompt = self._construct_system_prompt(prompt)
        # Start the spinner
        self.nvim.exec_lua('require("CopilotChat.spinner").show()')

        self.nvim.exec_lua(
            'require("CopilotChat.utils").log_info(...)', f"Chatting with {model} model"
        )

        if not disable_start_separator:
            self._add_start_separator(
                system_prompt, prompt, code, filetype, winnr, disable_separators
            )

        self._add_chat_messages(system_prompt, prompt, code, filetype, model)

        # Stop the spinner
        self.nvim.exec_lua('require("CopilotChat.spinner").hide()')

        if not disable_end_separator:
            self._add_end_separator(model, disable_separators)

    # private

    def _construct_system_prompt(self, prompt: str):
        system_prompt = system_prompts.COPILOT_INSTRUCTIONS
        if prompt == system_prompts.FIX_SHORTCUT:
            system_prompt = system_prompts.COPILOT_FIX
        elif prompt == system_prompts.TEST_SHORTCUT:
            system_prompt = system_prompts.COPILOT_TESTS
        elif prompt == system_prompts.EXPLAIN_SHORTCUT:
            system_prompt = system_prompts.COPILOT_EXPLAIN
        return system_prompt

    def _add_start_separator(
        self,
        system_prompt: str,
        prompt: str,
        code: str,
        file_type: str,
        winnr: int,
        no_annoyance: bool = False,
    ):
        if is_module_installed("tiktoken") and not no_annoyance:
            self._add_start_separator_with_token_count(
                system_prompt, prompt, code, file_type, winnr
            )
        else:
            self._add_regular_start_separator(
                system_prompt, prompt, code, file_type, winnr, no_annoyance
            )

    def _add_regular_start_separator(
        self,
        system_prompt: str,
        prompt: str,
        code: str,
        file_type: str,
        winnr: int,
        no_annoyance: bool = False,
    ):
        if code and not no_annoyance:
            code = f"\n        \nCODE:\n```{file_type}\n{code}\n```"

        last_row_before = len(self.buffer.lines())
        system_prompt_height = len(system_prompt.split("\n"))
        code_height = len(code.split("\n"))

        start_separator = (
            f"""### User

SYSTEM PROMPT:
```
{system_prompt}
```
{prompt}{code}

### Copilot

"""
            if not no_annoyance
            else f"### User\n{prompt}\n\n### Copilot\n\n"
        )
        self.buffer.append(start_separator.split("\n"))

        if no_annoyance:
            return
        self._add_folds(code, code_height, last_row_before, system_prompt_height, winnr)

    def _add_start_separator_with_token_count(
        self,
        system_prompt: str,
        prompt: str,
        code: str,
        file_type: str,
        winnr: int,
    ):
        import tiktoken

        encoding = tiktoken.encoding_for_model("gpt-4")

        num_total_tokens = len(encoding.encode(f"{system_prompt}\n{prompt}\n{code}"))
        num_system_tokens = len(encoding.encode(system_prompt))
        num_prompt_tokens = len(encoding.encode(prompt))
        num_code_tokens = len(encoding.encode(code))

        if code:
            code = f"\n        \nCODE: {num_code_tokens} Tokens \n```{file_type}\n{code}\n```"

        last_row_before = len(self.buffer.lines())
        system_prompt_height = len(system_prompt.split("\n"))
        code_height = len(code.split("\n"))

        start_separator = f"""### User

SYSTEM PROMPT: {num_system_tokens} Tokens
```
{system_prompt}
```
{prompt}{code}

### Copilot

"""
        self.buffer.append(start_separator.split("\n"))

        last_row_after = last_row_before + system_prompt_height + 5
        self.buffer.eol(last_row_before, f"{num_total_tokens} Total Tokens", "@float")
        self.buffer.eol(
            last_row_after, f"{num_prompt_tokens} Tokens", "NightflySteelBlue"
        )

        self._add_folds(code, code_height, last_row_before, system_prompt_height, winnr)

    def _add_folds(
        self,
        code: str,
        code_height: int,
        last_row_before: int,
        system_prompt_height: int,
        winnr: int,
    ):
        self.nvim.command("set foldmethod=manual")
        system_fold_start = last_row_before + 2
        system_fold_end = system_fold_start + system_prompt_height + 3
        main_command = f"{system_fold_start}, {system_fold_end} fold | normal! Gzz"
        full_command = f"call win_execute({winnr}, '{main_command}')"
        self.nvim.command(full_command)

        if code != "":
            code_fold_start = system_fold_end + 2
            code_fold_end = code_fold_start + code_height - 1
            main_command = f"{code_fold_start}, {code_fold_end} fold | normal! G"
            full_command = f"call win_execute({winnr}, '{main_command}')"
            self.nvim.command(full_command)

    def _add_chat_messages(
        self, system_prompt: str, prompt: str, code: str, file_type: str, model: str
    ):
        if self.copilot is None:
            self.copilot = Copilot(proxy=self.proxy)
            if self.copilot.github_token is None:
                req = self.copilot.request_auth()
                self.nvim.out_write(
                    f"Please visit {req['verification_uri']} and enter the code {req['user_code']}\n"
                )
                current_time = time.time()
                wait_until = current_time + req["expires_in"]
                while self.copilot.github_token is None:
                    self.copilot.poll_auth(req["device_code"])
                    time.sleep(req["interval"])
                    if time.time() > wait_until:
                        self.nvim.out_write("Timed out waiting for authentication\n")
                        return
                self.nvim.out_write("Successfully authenticated with Copilot\n")
            self.copilot.authenticate()

        last_line_col = 0
        for token in self.copilot.ask(
            system_prompt, prompt, code, language=cast(str, file_type), model=model
        ):
            self.nvim.exec_lua(
                'require("CopilotChat.utils").log_info(...)', f"Token: {token}"
            )
            buffer_lines = cast(list[str], self.buffer.lines())
            last_line_row = len(buffer_lines) - 1
            self.nvim.api.buf_set_text(
                self.buffer.number,
                last_line_row,
                last_line_col,
                last_line_row,
                last_line_col,
                token.split("\n"),
            )
            last_line_col += len(token.encode("utf-8"))
            if "\n" in token:
                last_line_col = 0

    def _add_end_separator(self, model: str, disable_separators: bool = False):
        current_datetime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        model_info = f"\n#### Answer provided by Copilot (Model: `{model}`) on {current_datetime}."
        additional_instructions = (
            "\n> For additional queries, please use the `CopilotChat` command."
        )
        disclaimer = "\n> Please be aware that the AI's output may not always be accurate. Always cross-verify the output."

        end_message = model_info + additional_instructions + disclaimer

        show_extra = disable_separators or ChatHandler.has_show_extra_info

        if show_extra:
            end_message = "\n" + current_datetime + "\n\n---\n"

        ChatHandler.has_show_extra_info = True

        self.buffer.append(end_message.split("\n"))