Minor codex improvements

slither/detectors/functions/codex.py

@@ -52,6 +52,10 @@ def _run_codex(self, logging_file: str, prompt: str) -> str:
 
         answer = ""
         res = {}
+
+        if self.slither.codex_organization:
+            openai_module.organization = self.slither.codex_organization
+
         try:
             res = openai_module.Completion.create(
                 prompt=prompt,

slither/slither.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Union, List, ValuesView, Type, Dict
+from typing import Union, List, ValuesView, Type, Dict, Optional
 
 from crytic_compile import CryticCompile, InvalidCompilation
 
@@ -90,6 +90,7 @@ def __init__(self, target: Union[str, CryticCompile], **kwargs):
         self.codex_temperature = kwargs.get("codex_temperature", 0)
         self.codex_max_tokens = kwargs.get("codex_max_tokens", 300)
         self.codex_log = kwargs.get("codex_log", False)
+        self.codex_organization: Optional[str] = kwargs.get("codex_organization", None)
 
         self.no_fail = kwargs.get("no_fail", False)
 

slither/tools/documentation/__main__.py

@@ -36,6 +36,13 @@ def parse_args() -> argparse.Namespace:
         default=False,
     )
 
+    parser.add_argument(
+        "--include-tests",
+        help="Include the tests",
+        action="store_true",
+        default=False,
+    )
+
     parser.add_argument(
         "--retry",
         help="Retry failed query (default 1). Each retry increases the temperature by 0.1",
@@ -202,6 +209,7 @@ def _handle_compilation_unit(
     overwrite: bool,
     force: bool,
     retry: int,
+    include_test: bool,
 ) -> None:
     logging_file: Optional[str]
     if slither.codex_log:
@@ -210,9 +218,8 @@ def _handle_compilation_unit(
         logging_file = None
 
     for scope in compilation_unit.scopes.values():
-
         # Dont send tests file
-        if (
+        if not include_test and (
             ".t.sol" in scope.filename.absolute
             or "mock" in scope.filename.absolute.lower()
             or "test" in scope.filename.absolute.lower()
@@ -271,6 +278,7 @@ def main() -> None:
                 args.overwrite,
                 args.force_answer_parsing,
                 int(args.retry),
+                args.include_tests,
             )
     except ImportError:
         pass

slither/utils/codex.py

@@ -64,6 +64,13 @@ def init_parser(parser: ArgumentParser, always_enable_codex: bool = False) -> No
         default=defaults_flag_in_config["codex_max_tokens"],
     )
 
+    group_codex.add_argument(
+        "--codex-organization",
+        help="Codex organization",
+        action="store",
+        default=None,
+    )
+
 
 # TODO: investigate how to set the correct return type
 # So that the other modules can work with openai