import os # To allow users to use arrow keys in the REPL. import readline # noqa: F401 import sys import typer from click import UsageError from click.types import Choice from prompt_toolkit import PromptSession from sgpt.config import cfg from sgpt.function import get_openai_schemas from sgpt.handlers.chat_handler import ChatHandler from sgpt.handlers.default_handler import DefaultHandler from sgpt.handlers.repl_handler import ReplHandler from sgpt.llm_functions.init_functions import install_functions as inst_funcs from sgpt.role import DefaultRoles, SystemRole from sgpt.utils import ( get_edited_prompt, get_sgpt_version, install_shell_integration, run_command, ) def main( prompt: str = typer.Argument( "", show_default=False, help="The prompt to generate completions for.", ), model: str = typer.Option( cfg.get("DEFAULT_MODEL"), help="Large language model to use.", ), temperature: float = typer.Option( cfg.get("DEFAULT_TEMPERATURE"), min=0.0, max=2.0, help="Randomness of generated output.", ), top_p: float = typer.Option( 1.0, min=0.0, max=1.0, help="Limits highest probable tokens (words).", ), md: bool = typer.Option( cfg.get("PRETTIFY_MARKDOWN") == "true", help="Prettify markdown output.", ), shell: bool = typer.Option( False, "--shell", "-s", help="Generate and execute shell commands.", rich_help_panel="Assistance Options", ), interaction: bool = typer.Option( cfg.get("SHELL_INTERACTION") == "true", help="Interactive mode for --shell option.", rich_help_panel="Assistance Options", ), describe_shell: bool = typer.Option( False, "--describe-shell", "-d", help="Describe a shell command.", rich_help_panel="Assistance Options", ), code: bool = typer.Option( False, "--code", "-c", help="Generate only code.", rich_help_panel="Assistance Options", ), functions: bool = typer.Option( cfg.get("OPENAI_USE_FUNCTIONS") == "true", help="Allow function calls.", rich_help_panel="Assistance Options", ), editor: bool = typer.Option( False, help="Open $EDITOR to provide a prompt.", ), cache: bool = typer.Option( True, help="Cache completion results.", ), version: bool = typer.Option( False, "--version", help="Show version.", callback=get_sgpt_version, ), chat: str = typer.Option( None, help="Follow conversation with id, " 'use "temp" for quick session.', rich_help_panel="Chat Options", ), repl: str = typer.Option( None, help="Start a REPL (Read-eval-print loop) session.", rich_help_panel="Chat Options", ), show_chat: str = typer.Option( None, help="Show all messages from provided chat id.", rich_help_panel="Chat Options", ), list_chats: bool = typer.Option( False, "--list-chats", "-lc", help="List all existing chat ids.", callback=ChatHandler.list_ids, rich_help_panel="Chat Options", ), role: str = typer.Option( None, help="System role for GPT model.", rich_help_panel="Role Options", ), create_role: str = typer.Option( None, help="Create role.", callback=SystemRole.create, rich_help_panel="Role Options", ), show_role: str = typer.Option( None, help="Show role.", callback=SystemRole.show, rich_help_panel="Role Options", ), list_roles: bool = typer.Option( False, "--list-roles", "-lr", help="List roles.", callback=SystemRole.list, rich_help_panel="Role Options", ), install_integration: bool = typer.Option( False, help="Install shell integration (ZSH and Bash only)", callback=install_shell_integration, hidden=True, # Hiding since should be used only once. ), install_functions: bool = typer.Option( False, help="Install default functions.", callback=inst_funcs, hidden=True, # Hiding since should be used only once. ), ) -> None: stdin_passed = not sys.stdin.isatty() if stdin_passed: stdin = "" # TODO: This is very hacky. # In some cases, we need to pass stdin along with inputs. # When we want part of stdin to be used as a init prompt, # but rest of the stdin to be used as a inputs. For example: # echo "hello\n__sgpt__eof__\nThis is input" | sgpt --repl temp # In this case, "hello" will be used as a init prompt, and # "This is input" will be used as "interactive" input to the REPL. # This is useful to test REPL with some initial context. for line in sys.stdin: if "__sgpt__eof__" in line: break stdin += line prompt = f"{stdin}\n\n{prompt}" if prompt else stdin try: # Switch to stdin for interactive input. if os.name == "posix": sys.stdin = open("/dev/tty", "r") elif os.name == "nt": sys.stdin = open("CON", "r") except OSError: # Non-interactive shell. pass if show_chat: ChatHandler.show_messages(show_chat, md) if sum((shell, describe_shell, code)) > 1: raise UsageError( "Only one of --shell, --describe-shell, and --code options can be used at a time." ) if chat and repl: raise UsageError("--chat and --repl options cannot be used together.") if editor and stdin_passed: raise UsageError("--editor option cannot be used with stdin input.") if editor: prompt = get_edited_prompt() role_class = ( DefaultRoles.check_get(shell, describe_shell, code) if not role else SystemRole.get(role) ) function_schemas = (get_openai_schemas() or None) if functions else None if repl: # Will be in infinite loop here until user exits with Ctrl+C. ReplHandler(repl, role_class, md).handle( init_prompt=prompt, model=model, temperature=temperature, top_p=top_p, caching=cache, functions=function_schemas, ) if chat: full_completion = ChatHandler(chat, role_class, md).handle( prompt=prompt, model=model, temperature=temperature, top_p=top_p, caching=cache, functions=function_schemas, ) else: full_completion = DefaultHandler(role_class, md).handle( prompt=prompt, model=model, temperature=temperature, top_p=top_p, caching=cache, functions=function_schemas, ) session: PromptSession[str] = PromptSession() while shell and interaction: option = typer.prompt( text="[E]xecute, [M]odify, [D]escribe, [A]bort", type=Choice(("e", "m", "d", "a", "y"), case_sensitive=False), default="e" if cfg.get("DEFAULT_EXECUTE_SHELL_CMD") == "true" else "a", show_choices=False, show_default=False, ) if option in ("e", "y"): # "y" option is for keeping compatibility with old version. run_command(full_completion) elif option == "m": full_completion = session.prompt("", default=full_completion) continue elif option == "d": DefaultHandler(DefaultRoles.DESCRIBE_SHELL.get_role(), md).handle( full_completion, model=model, temperature=temperature, top_p=top_p, caching=cache, functions=function_schemas, ) continue break def entry_point() -> None: typer.run(main) if __name__ == "__main__": entry_point()