From 2102725241e56cfc28c2bad1072c1c9f7f86553d Mon Sep 17 00:00:00 2001
From: Feist Josselin <josselin.feist@gmail.com>
Date: Fri, 13 Jan 2023 10:38:18 +0100
Subject: [PATCH] Minor codex improvements

---
 slither/detectors/functions/codex.py    |  4 ++++
 slither/slither.py                      |  3 ++-
 slither/tools/documentation/__main__.py | 12 ++++++++++--
 slither/utils/codex.py                  |  7 +++++++
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/slither/detectors/functions/codex.py b/slither/detectors/functions/codex.py
index fb00f64c04..6b694155e0 100644
--- a/slither/detectors/functions/codex.py
+++ b/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,
diff --git a/slither/slither.py b/slither/slither.py
index 45d99906fd..227a37365e 100644
--- a/slither/slither.py
+++ b/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)
 
diff --git a/slither/tools/documentation/__main__.py b/slither/tools/documentation/__main__.py
index 8e545fb095..39b1eacdac 100644
--- a/slither/tools/documentation/__main__.py
+++ b/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
diff --git a/slither/utils/codex.py b/slither/utils/codex.py
index 3b06efe6f5..8b15656708 100644
--- a/slither/utils/codex.py
+++ b/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