Revert "Revert "/unit_tests switch to devchat env""

Author: kagami-l

unit_tests/assistants/directory_structure/base.py

@@ -14,7 +14,6 @@ def __init__(
     ) -> None:
         self._root_path = root_path
 
-        self._client = OpenAI()
         self._chat_language = chat_language
 
     @property

unit_tests/assistants/directory_structure/relevant_file_finder.py

@@ -2,19 +2,17 @@
 from pathlib import Path
 from typing import Callable, List
 
-import tiktoken
 from assistants.directory_structure.base import DirectoryStructureBase
 from assistants.rerank_files import rerank_files
 from openai_util import create_chat_completion_content
-from tools.directory_viewer import (
-    ListViewer,
-)
+from tools.directory_viewer import ListViewer
+from tools.tiktoken_util import get_encoding
 
 
 class RelevantFileFinder(DirectoryStructureBase):
     model_name = "gpt-3.5-turbo-1106"
     dir_token_budget = 16000 * 0.95
-    encoding: tiktoken.Encoding = tiktoken.encoding_for_model("gpt-3.5-turbo-1106")
+    encoding = get_encoding("cl100k_base")
 
     def _paginate_dir_structure(
         self, criteria: Callable[[Path], bool], style: str = "list"
@@ -85,7 +83,6 @@ def _find_relevant_files(self, objective: str, dir_structure_pages: List[str]) -
             user_msg = self._mk_message(objective, dir_structure)
 
             response = create_chat_completion_content(
-                client=self._client,
                 model=self.model_name,
                 messages=[
                     {"role": "user", "content": user_msg},

unit_tests/command.yml

@@ -1,3 +1,3 @@
 description: Generate unit tests.
 steps:
-  - run: $command_python $command_path/main.py "$input"
+  - run: $devchat_python $command_path/main.py "$input"

unit_tests/propose_test.py

@@ -2,13 +2,14 @@
 from functools import partial
 from typing import List
 
-import tiktoken
 from model import FuncToTest, TokenBudgetExceededException
 from openai_util import create_chat_completion_content
 from prompts import PROPOSE_TEST_PROMPT
+from tools.tiktoken_util import get_encoding
 
 MODEL = "gpt-3.5-turbo-1106"
 # MODEL = "gpt-4-1106-preview"
+ENCODING = "cl100k_base"
 TOKEN_BUDGET = int(16000 * 0.9)
 
 
@@ -20,7 +21,7 @@ def _mk_user_msg(
     """
     Create a user message to be sent to the model within the token budget.
     """
-    encoding: tiktoken.Encoding = tiktoken.encoding_for_model(MODEL)
+    encoding = get_encoding(ENCODING)
 
     func_content = f"function code\n```\n{func_to_test.func_content}\n```\n"
     class_content = ""

unit_tests/tools/tiktoken_util.py

@@ -0,0 +1,20 @@
+import tiktoken
+
+
+def get_encoding(encoding_name: str):
+    """
+    Get a tiktoken encoding by name.
+    """
+    try:
+        return tiktoken.get_encoding(encoding_name)
+    except Exception:
+        from tiktoken import registry
+        from tiktoken.core import Encoding
+        from tiktoken.registry import _find_constructors
+
+        def _get_encoding(name: str):
+            _find_constructors()
+            constructor = registry.ENCODING_CONSTRUCTORS[name]
+            return Encoding(**constructor(), use_pure_python=True)
+
+        return _get_encoding(encoding_name)

unit_tests/write_tests.py

@@ -1,13 +1,14 @@
 from functools import partial
 from typing import List, Optional
 
-import tiktoken
 from model import FuncToTest, TokenBudgetExceededException
 from openai_util import create_chat_completion_chunks
 from prompts import WRITE_TESTS_PROMPT
 from tools.file_util import retrieve_file_content
+from tools.tiktoken_util import get_encoding
 
 MODEL = "gpt-4-1106-preview"
+ENCODING = "cl100k_base"
 TOKEN_BUDGET = int(128000 * 0.9)
 
 
@@ -18,7 +19,7 @@ def _mk_write_tests_msg(
     chat_language: str,
     reference_files: Optional[List[str]] = None,
 ) -> Optional[str]:
-    encoding: tiktoken.Encoding = tiktoken.encoding_for_model(MODEL)
+    encoding = get_encoding(ENCODING)
 
     test_cases_str = ""
     for i, test_case in enumerate(test_cases, 1):