diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..8ba6df7 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,29 @@ +# EditorConfig is awesome: http://EditorConfig.org + +# top-most EditorConfig file +root = true + +# Unix-style newlines with a newline ending every file +[*] +end_of_line = lf +insert_final_newline = true +trim_trailing_whitespace = true +charset = utf-8 + +# 4 space indentation +[*.{py,java,r,R}] +indent_style = space +indent_size = 4 + +# 2 space indentation +[*.{js,json,y{a,}ml,html,cwl}] +indent_style = space +indent_size = 2 + +[*.{md,Rmd,rst}] +trim_trailing_whitespace = false +indent_style = space +indent_size = 2 + +[*.py] +skip = build,.tox,.venv diff --git a/.gitignore b/.gitignore index 31aa444..10151c7 100644 --- a/.gitignore +++ b/.gitignore @@ -161,4 +161,7 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ \ No newline at end of file +#.idea/ + + +*.tmp diff --git a/GPTtrace.py b/GPTtrace.py index ec1a560..b22824d 100755 --- a/GPTtrace.py +++ b/GPTtrace.py @@ -1,43 +1,76 @@ #! /bin/env python3 -import os import argparse - -from revChatGPT.V1 import Chatbot +import os from typing import List -from marko.parser import Parser + +import pygments from marko.block import FencedCode from marko.inline import RawText +from marko.parser import Parser +from prompt_toolkit import print_formatted_text +from prompt_toolkit.formatted_text import PygmentsTokens +from pygments.token import Token +from pygments_markdown_lexer import MarkdownLexer +from revChatGPT.V1 import Chatbot ENV_UUID = "GPTTRACE_CONV_UUID" ENV_ACCESS_TOKEN = "GPTTRACE_ACCESS_TOKEN" +def pretty_print(input, lexer=MarkdownLexer): + tokens = list(pygments.lex(input, lexer=lexer())) + print_formatted_text(PygmentsTokens(tokens)) + + def main(): parser = argparse.ArgumentParser( - prog='GPTtrace', - description='Use ChatGPT to write eBPF programs (bpftrace, etc.)') + prog="GPTtrace", + description="Use ChatGPT to write eBPF programs (bpftrace, etc.)", + ) group = parser.add_mutually_exclusive_group() group.add_argument( - "-i", "--info", help="Let ChatGPT explain what's eBPF", action="store_true") + "-i", "--info", help="Let ChatGPT explain what's eBPF", action="store_true" + ) group.add_argument( - "-e", "--execute", help="Generate commands using your input with ChatGPT, and run it", action="store", metavar="TEXT") + "-e", + "--execute", + help="Generate commands using your input with ChatGPT, and run it", + action="store", + metavar="TEXT", + ) group.add_argument( - "-g", "--generate", help="Generate eBPF programs using your input with ChatGPT", action="store", metavar="TEXT") - + "-g", + "--generate", + help="Generate eBPF programs using your input with ChatGPT", + action="store", + metavar="TEXT", + ) + + parser.add_argument( + "-v", + "--verbose", + help="Print the prompt and receive message", + action="store_true", + ) parser.add_argument( - "-v", "--verbose", help="Print the prompt and receive message", action="store_true") + "-u", + "--uuid", + help=f"Conversion UUID to use, or passed through environment variable `{ENV_UUID}`", + ) parser.add_argument( - "-u", "--uuid", help=f"Conversion UUID to use, or passed through environment variable `{ENV_UUID}`") - parser.add_argument("-t", "--access-token", - help=f"ChatGPT access token, see `https://chat.openai.com/api/auth/session` or passed through `{ENV_ACCESS_TOKEN}`") + "-t", + "--access-token", + help=f"ChatGPT access token, see `https://chat.openai.com/api/auth/session` or passed through `{ENV_ACCESS_TOKEN}`", + ) args = parser.parse_args() access_token = args.access_token or os.environ.get(ENV_ACCESS_TOKEN, None) conv_uuid = args.uuid or os.environ.get(ENV_UUID, None) if access_token is None: print( - f"Either provide your access token through `-t` or through environment variable {ENV_ACCESS_TOKEN}") + f"Either provide your access token through `-t` or through environment variable {ENV_ACCESS_TOKEN}" + ) return chatbot = Chatbot(config={"access_token": access_token}) if args.info: @@ -46,7 +79,9 @@ def main(): desc: str = args.execute print("Sending query to ChatGPT: " + desc) ret_val = generate_result( - chatbot, construct_running_prompt(desc), conv_uuid, args.verbose) + chatbot, construct_running_prompt(desc), conv_uuid, args.verbose + ) + pretty_print(ret_val) # print(ret_val) parsed = make_executable_command(ret_val) # print(f"Command to run: {parsed}") @@ -56,7 +91,8 @@ def main(): desc: str = args.generate print("Sending query to ChatGPT: " + desc) ret_val = generate_result( - chatbot, construct_generate_prompt(desc), conv_uuid, True) + chatbot, construct_generate_prompt(desc), conv_uuid, True + ) # print(ret_val) parsed = extract_code_blocks(ret_val) # print(f"Command to run: {parsed}") @@ -66,44 +102,47 @@ def main(): else: parser.print_help() + def construct_generate_prompt(text: str) -> str: - return f'''You are now a translater from human language to {os.uname()[0]} eBPF programs. + return f"""You are now a translater from human language to {os.uname()[0]} eBPF programs. Please write eBPF programs for me. No explanation required, no instruction required, don't tell me how to compile and run. -What I want is a eBPF program for: {text}.''' +What I want is a eBPF program for: {text}.""" + def construct_running_prompt(text: str) -> str: - return f'''You are now a translater from human language to {os.uname()[0]} shell bpftrace command. + return f"""You are now a translater from human language to {os.uname()[0]} shell bpftrace command. No explanation required. -respond with only the raw shell bpftrace command. -It should be start with `bpftrace`. +respond with only the raw shell bpftrace command. +It should be start with `bpftrace`. Your response should be able to put into a shell and run directly. Just output in one line, without any description, or any other text that cannot be run in shell. -What should I type to shell to trace using bpftrace for: {text}, in one line.''' +What should I type to shell to trace using bpftrace for: {text}, in one line.""" def make_executable_command(command: str) -> str: - if command.startswith('\n'): + if command.startswith("\n"): command = command[1:] - if command.endswith('\n'): + if command.endswith("\n"): command = command[:-1] - if command.startswith('`'): + if command.startswith("`"): command = command[1:] - if command.endswith('`'): + if command.endswith("`"): command = command[:-1] command = command.strip() - command = command.split('User: ')[0] + command = command.split("User: ")[0] return command -def generate_result(bot: Chatbot, text: str, session: str = None, print_out: bool = False) -> str: +def generate_result( + bot: Chatbot, text: str, session: str = None, print_out: bool = False +) -> str: from io import StringIO + prev_text = "" buf = StringIO() - for data in bot.ask( - text, conversation_id=session - ): - message = data["message"][len(prev_text):] + for data in bot.ask(text, conversation_id=session): + message = data["message"][len(prev_text) :] if print_out: print(message, end="", flush=True) buf.write(message) diff --git a/requirements.txt b/requirements.txt index 911747d..afdcddb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,9 +2,11 @@ anyio==3.6.2 astroid==2.14.2 asyncio==3.4.3 autopep8==2.0.1 +black==23.1.0 blobfile==2.0.1 certifi==2022.12.7 charset-normalizer==3.0.1 +click==8.1.3 dill==0.3.6 filelock==3.9.0 flake8==6.0.0 @@ -15,12 +17,20 @@ idna==3.4 isort==5.12.0 lazy-object-proxy==1.9.0 lxml==4.9.2 +marko==1.3.0 mccabe==0.7.0 +mypy-extensions==1.0.0 OpenAIAuth==0.2.0 +packaging==23.0 +pathspec==0.11.0 +pkg_resources==0.0.0 platformdirs==3.0.0 +prompt-toolkit==3.0.36 pycodestyle==2.10.0 pycryptodomex==3.17 pyflakes==3.0.1 +Pygments==2.14.0 +pygments-markdown-lexer==0.1.0.dev39 pylint==2.16.2 regex==2022.10.31 requests==2.28.2 @@ -34,5 +44,5 @@ tomli==2.0.1 tomlkit==0.11.6 typing_extensions==4.5.0 urllib3==1.26.14 +wcwidth==0.2.6 wrapt==1.14.1 -marko==1.3.0 \ No newline at end of file