115 lines
3.0 KiB
Python
Executable File
115 lines
3.0 KiB
Python
Executable File
import json
|
|
import requests
|
|
import os
|
|
from AuthToken import AuthToken
|
|
import click
|
|
from rich import print
|
|
from rich.panel import Panel
|
|
from rich.text import Text
|
|
import pyperclip
|
|
import readline
|
|
|
|
# get a suggestion using github copilot
|
|
def get_suggestion(prompt, token):
|
|
headers = {
|
|
'Content-Type': 'application/json',
|
|
'Openai-Organization': 'copilot-ghost',
|
|
'OpenAI-Intent': 'github-copilot',
|
|
'Authorization': f'Bearer {token}',
|
|
}
|
|
|
|
data = {
|
|
'prompt': prompt,
|
|
'stream': True,
|
|
'temperature': 0.3,
|
|
'max_tokens': 50,
|
|
'top_p': 0.5,
|
|
'n': 1,
|
|
'logprobs': 2,
|
|
}
|
|
|
|
response = requests.post('https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex/completions', headers=headers, data=json.dumps(data))
|
|
# get the first completion from the response
|
|
data = response.text
|
|
|
|
# split by \n
|
|
data = data.split("\n")
|
|
|
|
# filter where start with data: and remove it
|
|
data = [x.replace("data: ", "") for x in data if x.startswith("data: ")]
|
|
|
|
# filter [Done]
|
|
data = [x for x in data if x != "[DONE]"]
|
|
|
|
choices: dict(int, str) = {}
|
|
for x in data:
|
|
# x to json
|
|
x = json.loads(x)
|
|
c = x["choices"][0]
|
|
|
|
if c["index"] not in choices:
|
|
choices[c["index"]] = c["text"]
|
|
else:
|
|
choices[c["index"]] += c["text"]
|
|
|
|
# get the first choice else return "Command not found"
|
|
if 0 in choices:
|
|
choice = choices[0]
|
|
|
|
# remove unwanted characters from front (#, $, \n)
|
|
while choice.startswith("#") or choice.startswith("$") or choice.startswith("\n") or choice.startswith(" "):
|
|
choice = choice[1:]
|
|
|
|
# strip \n
|
|
choice = choice.strip()
|
|
|
|
return choice.split("\n")[0]
|
|
|
|
return "Command not found"
|
|
|
|
@click.group()
|
|
def main():
|
|
pass
|
|
|
|
@main.command()
|
|
@click.argument('prompt', nargs=-1)
|
|
def generate_suggestion(prompt):
|
|
prompt = " ".join(prompt)
|
|
prompt = "!/bin/bash\n\n" + prompt + ":\n"
|
|
|
|
authToken = AuthToken()
|
|
token = authToken.get_github_api_token()
|
|
|
|
suggestion = get_suggestion(prompt, token)
|
|
panel = Panel(Text(suggestion, justify="center"), title_align="center", title="Execute? \[y/n/e/c]", expand=False)
|
|
print(panel)
|
|
# read a character from stdin without displaying it
|
|
c = click.getchar()
|
|
if c == "y":
|
|
# execute the suggestion
|
|
os.system(suggestion)
|
|
elif c == "e":
|
|
# prompt the suggestion by using redisplay
|
|
def pre_input_hook():
|
|
readline.insert_text(suggestion)
|
|
readline.redisplay()
|
|
|
|
|
|
readline.set_pre_input_hook(pre_input_hook)
|
|
line = input()
|
|
|
|
# execute line
|
|
os.system(line)
|
|
# remove pre input hook
|
|
readline.set_pre_input_hook()
|
|
|
|
# set history
|
|
readline.clear_history()
|
|
readline.add_history(line)
|
|
elif c == "c":
|
|
# copy the suggestion to clipboard
|
|
pyperclip.copy(suggestion)
|
|
print("Copied to clipboard")
|
|
|
|
|