Optionally use tiktoken - fallback on estimate of tokens

Author: zanjonke

plain2code_console.py

@@ -1,10 +1,16 @@
 import logging
 
-import tiktoken
 from rich.console import Console
 from rich.style import Style
 from rich.tree import Tree
 
+try:
+    import tiktoken
+
+    _tiktoken_available = True
+except ImportError:
+    _tiktoken_available = False
+
 
 class Plain2CodeConsole(Console):
     INFO_STYLE = Style()
@@ -16,7 +22,10 @@ class Plain2CodeConsole(Console):
 
     def __init__(self):
         super().__init__()
-        self.llm_encoding = tiktoken.get_encoding("cl100k_base")
+        if _tiktoken_available:
+            self.llm_encoding = tiktoken.get_encoding("cl100k_base")
+        else:
+            self.llm_encoding = None
 
     def info(self, *args, **kwargs):
         logging.info(" ".join(map(str, args)))
@@ -86,7 +95,7 @@ def _create_tree_from_files(self, root_folder, files):
                             current_level = current_level.add(f"{part} [red]deleted[/red]")
                         else:
                             file_lines = len(content.splitlines())
-                            file_tokens = len(self.llm_encoding.encode(content))
+                            file_tokens = self._count_tokens(content)
                             current_level = current_level.add(f"{part} ({file_lines} lines, {file_tokens} tokens)")
                     else:
                         current_level = current_level.add(part)
@@ -95,6 +104,12 @@ def _create_tree_from_files(self, root_folder, files):
 
         return tree
 
+    def _count_tokens(self, text):
+        """Count tokens using tiktoken if available, otherwise estimate from character count."""
+        if self.llm_encoding is not None:
+            return len(self.llm_encoding.encode(text))
+        return len(text) // 4
+
     def print_resources(self, resources_list, linked_resources):
         if len(resources_list) == 0:
             self.input("No linked resources found.")
@@ -103,7 +118,7 @@ def print_resources(self, resources_list, linked_resources):
         self.input("Linked resources:")
         for resource_name in resources_list:
             if resource_name["target"] in linked_resources:
-                file_tokens = len(self.llm_encoding.encode(linked_resources[resource_name["target"]]))
+                file_tokens = self._count_tokens(linked_resources[resource_name["target"]])
                 self.input(
                     f"- {resource_name['text']} [#4169E1]({resource_name['target']}, {file_tokens} tokens)[/#4169E1]"
                 )