slither.utils.codex

  1import logging
  2import os
  3from argparse import ArgumentParser
  4from pathlib import Path
  5
  6from slither.utils.command_line import defaults_flag_in_config
  7
  8logger = logging.getLogger("Slither")
  9
 10
 11def init_parser(parser: ArgumentParser, always_enable_codex: bool = False) -> None:
 12    """
 13    Init the cli arg with codex features
 14
 15    Args:
 16        parser:
 17        always_enable_codex (Optional(bool)): if true, --codex is not enabled
 18
 19    Returns:
 20
 21    """
 22    group_codex = parser.add_argument_group("Codex (https://beta.openai.com/docs/guides/code)")
 23
 24    if not always_enable_codex:
 25        group_codex.add_argument(
 26            "--codex",
 27            help="Enable codex (require an OpenAI API Key)",
 28            action="store_true",
 29            default=defaults_flag_in_config["codex"],
 30        )
 31
 32    group_codex.add_argument(
 33        "--codex-log",
 34        help="Log codex queries (in crytic_export/codex/)",
 35        action="store_true",
 36        default=False,
 37    )
 38
 39    group_codex.add_argument(
 40        "--codex-contracts",
 41        help="Comma separated list of contracts to submit to OpenAI Codex",
 42        action="store",
 43        default=defaults_flag_in_config["codex_contracts"],
 44    )
 45
 46    group_codex.add_argument(
 47        "--codex-model",
 48        help="Name of the Codex model to use (affects pricing).  Defaults to 'text-davinci-003'",
 49        action="store",
 50        default=defaults_flag_in_config["codex_model"],
 51    )
 52
 53    group_codex.add_argument(
 54        "--codex-temperature",
 55        help="Temperature to use with Codex.  Lower number indicates a more precise answer while higher numbers return more creative answers.  Defaults to 0",
 56        action="store",
 57        default=defaults_flag_in_config["codex_temperature"],
 58    )
 59
 60    group_codex.add_argument(
 61        "--codex-max-tokens",
 62        help="Maximum amount of tokens to use on the response.  This number plus the size of the prompt can be no larger than the limit (4097 for text-davinci-003)",
 63        action="store",
 64        default=defaults_flag_in_config["codex_max_tokens"],
 65    )
 66
 67    group_codex.add_argument(
 68        "--codex-organization",
 69        help="Codex organization",
 70        action="store",
 71        default=None,
 72    )
 73
 74
 75# TODO: investigate how to set the correct return type
 76# So that the other modules can work with openai
 77def openai_module():  # type: ignore
 78    """
 79    Return the openai module
 80    Consider checking the usage of open (slither.codex_enabled) before using this function
 81
 82    Returns:
 83        Optional[the openai module]
 84    """
 85    try:
 86        # pylint: disable=import-outside-toplevel
 87        import openai
 88
 89        api_key = os.getenv("OPENAI_API_KEY")
 90        if api_key is None:
 91            logger.info(
 92                "Please provide an Open API Key in OPENAI_API_KEY (https://beta.openai.com/account/api-keys)"
 93            )
 94            return None
 95        openai.api_key = api_key
 96    except ImportError:
 97        logger.info("OpenAI was not installed")  # type: ignore
 98        logger.info('run "pip install openai"')
 99        return None
100    return openai
101
102
103def log_codex(filename: str, prompt: str) -> None:
104    """
105    Log the prompt in crytic/export/codex/filename
106    Append to the file
107
108    Args:
109        filename: filename to write to
110        prompt: prompt to write
111
112    Returns:
113        None
114    """
115
116    Path("crytic_export/codex").mkdir(parents=True, exist_ok=True)
117
118    with open(Path("crytic_export/codex", filename), "a", encoding="utf8") as file:
119        file.write(prompt)
120        file.write("\n")