120 lines
3.2 KiB
Python
Executable File
120 lines
3.2 KiB
Python
Executable File
import json
|
|
import requests
|
|
import os
|
|
import sys
|
|
import time
|
|
from os.path import expanduser
|
|
import click
|
|
|
|
|
|
# get a new github copilot token
|
|
def get_token():
|
|
|
|
valid = False
|
|
token = {}
|
|
|
|
# check if a token exists in the file
|
|
if os.path.exists("token.json"):
|
|
with open("token.json", "r") as f:
|
|
token = json.load(f)
|
|
|
|
# check if the token is valid
|
|
if "token" in token and "expires_at" in token:
|
|
if token["expires_at"] > time.time():
|
|
valid = True
|
|
|
|
if not valid:
|
|
# read token from file
|
|
|
|
# read file from home directory
|
|
home = expanduser("~")
|
|
with open(f"{home}/.config/github-copilot/hosts.json", "r") as f:
|
|
terms = json.load(f)
|
|
auth_token = terms["github.com"]["oauth_token"]
|
|
|
|
|
|
headers = {
|
|
'Authorization': f"Bearer {auth_token}",
|
|
}
|
|
response = requests.get('https://api.github.com/copilot_internal/v2/token', headers=headers)
|
|
|
|
if response.status_code != 200:
|
|
return None
|
|
|
|
token = response.json()
|
|
|
|
# save the token in a json file with the expiration time
|
|
with open("token.json", "w") as f:
|
|
json.dump(token, f)
|
|
|
|
return token["token"]
|
|
|
|
# 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"
|
|
token = get_token()
|
|
print(get_suggestion(prompt, token)) |