diff --git a/samuelcolvin_aicli.py b/samuelcolvin_aicli.py index a26742e..32a9082 100644 --- a/samuelcolvin_aicli.py +++ b/samuelcolvin_aicli.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +import argparse import os import sys from datetime import datetime, timezone @@ -30,8 +31,20 @@ def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderR def cli() -> int: + parser = argparse.ArgumentParser(prog='aicli', description=f'OpenAI powered AI CLI v{__version__}') + parser.add_argument('prompt', nargs='?', help='AI Prompt, if omitted fall into interactive mode') + + # allows you to disable streaming responses if they get annoying or are more expensive. + parser.add_argument('--no-stream', action='store_true', help='Whether to stream responses from OpenAI') + + parser.add_argument('--version', action='store_true', help='Show version and exit') + + args = parser.parse_args() + console = Console() console.print(f'aicli - OpenAI powered AI CLI v{__version__}', style='green bold', highlight=False) + if args.version: + return 0 try: openai.api_key = os.environ['OPENAI_API_KEY'] @@ -42,16 +55,23 @@ def cli() -> int: now_utc = datetime.now(timezone.utc) setup = f"""\ Help the user by responding to their request, the output should be concise and always written in markdown. -The current date and time is {datetime.now()} {now_utc.astimezone().tzinfo.tzname(now_utc)}.""" +The current date and time is {datetime.now()} {now_utc.astimezone().tzinfo.tzname(now_utc)}. +The user is running {sys.platform}.""" + stream = not args.no_stream messages = [{'role': 'system', 'content': setup}] + if args.prompt: + messages.append({'role': 'user', 'content': args.prompt}) + try: + ask_openai(messages, stream, console) + except KeyboardInterrupt: + pass + return 0 + history = Path().home() / '.openai-prompt-history.txt' session = PromptSession(history=FileHistory(str(history))) - # allows you to disable streaming responses if they get annoying or are more expensive. - stream = True - while True: try: text = session.prompt('aicli ➤ ', auto_suggest=AutoSuggestFromHistory()) @@ -67,35 +87,41 @@ def cli() -> int: console.print(Syntax(last_content, lexer='markdown', background_color='default')) continue - status = Status('[dim]Working on it…[/dim]', console=console) - status.start() messages.append({'role': 'user', 'content': text}) try: - response = openai.ChatCompletion.create(model='gpt-4', messages=messages, stream=stream) + content = ask_openai(messages, stream, console) except KeyboardInterrupt: - status.stop() return 0 + messages.append({'role': 'assistant', 'content': content}) + - status.stop() - if stream: - content = '' +def ask_openai(messages: list[dict[str, str]], stream: bool, console: Console) -> str: + with Status('[dim]Working on it…[/dim]', console=console): + response = openai.ChatCompletion.create(model='gpt-4', messages=messages, stream=stream) + + console.print('\nResponse:', style='green') + if stream: + content = '' + interrupted = False + with Live('', refresh_per_second=15, console=console) as live: try: - with Live('', refresh_per_second=15, console=console) as live: - for chunk in response: - if chunk['choices'][0]['finish_reason'] is not None: - break - chunk_text = chunk['choices'][0]['delta'].get('content', '') - content += chunk_text - live.update(Markdown(content)) + for chunk in response: + if chunk['choices'][0]['finish_reason'] is not None: + break + chunk_text = chunk['choices'][0]['delta'].get('content', '') + content += chunk_text + live.update(Markdown(content)) except KeyboardInterrupt: - console.print('[dim]Interrupted[/dim]') - else: - content = response['choices'][0]['message']['content'] - md = Markdown(content) - console.print(md) + interrupted = True - messages.append({'role': 'assistant', 'content': content}) + if interrupted: + console.print('[dim]Interrupted[/dim]') + else: + content = response['choices'][0]['message']['content'] + console.print(Markdown(content)) + + return content if __name__ == '__main__':