janus-llm 1.0.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

janus/language/splitter.py

@@ -1,4 +1,5 @@
 import re
+from itertools import compress
 from pathlib import Path
 from typing import List
 
@@ -21,6 +22,20 @@ class TokenLimitError(Exception):
     pass
 
 
+class EmptyTreeError(Exception):
+    """An exception raised when the tree is empty or does not exist (can happen
+    when there are no nodes of interest in the tree)
+    """
+
+    pass
+
+
+class FileSizeError(Exception):
+    """An exception raised when the file size is too large for the splitter"""
+
+    pass
+
+
 class Splitter(FileManager):
     """A class for splitting code into functional blocks to prompt with for
     transcoding.
@@ -31,7 +46,10 @@ class Splitter(FileManager):
         language: str,
         model: None | BaseLanguageModel = None,
         max_tokens: int = 4096,
-        use_placeholders: bool = True,
+        skip_merge: bool = False,
+        protected_node_types: tuple[str] = (),
+        prune_node_types: tuple[str] = (),
+        prune_unprotected: bool = False,
     ):
         """
         Arguments:
@@ -39,14 +57,22 @@ class Splitter(FileManager):
             model: The name of the model to use for counting tokens. If the model is None,
                 will use tiktoken's default tokenizer to count tokens.
             max_tokens: The maximum number of tokens to use for each functional block.
-            use_placeholders: Whether to use placeholders when splitting the code.
+            skip_merge: Whether to merge child nodes up to the max_token length.
+                May be used for situations like documentation where function-level
+                documentation is preferred.
+                TODO: Maybe instead support something like a list of node types that
+                      shouldnt be merged (e.g. functions, classes)?
+            prune_unprotected: Whether to prune unprotected nodes from the tree.
         """
         super().__init__(language=language)
         self.model = model
         if self.model is None:
             self._encoding = tiktoken.get_encoding("cl100k_base")
-        self.use_placeholders: bool = use_placeholders
+        self.skip_merge = skip_merge
         self.max_tokens: int = max_tokens
+        self._protected_node_types = set(protected_node_types)
+        self._prune_node_types = set(prune_node_types)
+        self.prune_unprotected = prune_unprotected
 
     def split(self, file: Path | str) -> CodeBlock:
         """Split the given file into functional code blocks.
@@ -59,11 +85,27 @@ class Splitter(FileManager):
         """
         path = Path(file)
         code = path.read_text()
+        return self.split_string(code, path.name)
+
+    def split_string(self, code: str, name: str) -> CodeBlock:
+        """Split the given code into functional code blocks.
+
+        Arguments:
+            code: The code as a string to split into functional blocks.
+            name: The filename of the code block.
+
+        Returns:
+            A `CodeBlock` made up of nested `CodeBlock`s.
+        """
 
         root = self._get_ast(code)
-        self._set_identifiers(root, path)
+        self._prune(root)
+        if self.prune_unprotected:
+            self._prune_unprotected(root)
+        self._set_identifiers(root, name)
         self._segment_leaves(root)
-        self._merge_tree(root)
+        if not self.skip_merge:
+            self._merge_tree(root)
 
         return root
 
@@ -79,12 +121,25 @@ class Splitter(FileManager):
         """
         raise NotImplementedError()
 
-    def _set_identifiers(self, root: CodeBlock, path: Path):
+    def _all_node_types(self, root: CodeBlock) -> set[NodeType]:
+        types = set()
+        stack = [root]
+        while stack:
+            node = stack.pop()
+            types.add(node.node_type)
+            stack.extend(node.children)
+        return types
+
+    def _set_identifiers(self, root: CodeBlock, name: str):
         """Set the IDs and names of each node in the given tree. By default,
         node IDs take the form `child_<i>`, where <i> is an integer counter which
         increments in breadth-first order, and node names take the form
-        `<filename>:<ID>`. Child classes should override this function to use
+        `<name>:<ID>`. Child classes should override this function to use
         more informative names based on the particular programming language.
+
+        Arguments:
+            root: The root of the tree to set identifiers for.
+            name: The name of the file being split.
         """
         seen_ids = 0
         queue = [root]
@@ -92,7 +147,7 @@ class Splitter(FileManager):
             node = queue.pop(0)  # BFS order to keep lower IDs toward the root
             node.id = f"child_{seen_ids}"
             seen_ids += 1
-            node.name = f"{path.name}:{node.id}"
+            node.name = f"{name}:{node.id}"
             queue.extend(node.children)
 
     def _merge_tree(self, root: CodeBlock):
@@ -102,76 +157,90 @@ class Splitter(FileManager):
         the represented code is present in the text of exactly one node in the
         tree.
         """
+        # Simulate recursion with a stack
         stack = [root]
         while stack:
             node = stack.pop()
-            self._merge_children(node)
-            stack.extend(node.children)
 
-    def _merge_children(self, node: CodeBlock):
-        """Given a parent node in an abstract syntax tree, consolidate, merge,
-        and prune its children such that this node's text fits into context,
-        and does not overlap with the text represented by any of its children.
-        After processing, this node's children will have been merged such that
-        they maximally fit into LLM context. If the entire node text can fit
-        into context, all its children will be pruned.
-        """
-        # If the text at the function input is less than the max tokens, then
-        #  we can just return it as a CodeBlock with no children.
-        if node.tokens <= self.max_tokens:
-            node.children = []
-            return
-
-        node.complete = False
+            # If the text of this node can fit in context, then we can just
+            #  prune its children, making it a leaf node.
+            if node.tokens <= self.max_tokens and not self._has_protected_descendents(
+                node
+            ):
+                node.children = []
+                continue
 
-        # Consolidate nodes into groups, and then merge each group into a new node
-        node_groups = self._group_nodes(node.children)
-        node.children = list(map(self.merge_nodes, node_groups))
+            # Otherwise, this is an internal node. Mark it as incomplete, and
+            #  drop its text (which will be represented in its children)
+            node.complete = False
+            node.text = None
+            node.tokens = 0
 
-        # If not using placeholders, simply recurse for every child and delete
-        #  this node's text and tokens
-        if not self.use_placeholders:
+            # If this node has no children but cannot fit into context, then we
+            #  have a problem. Oversized nodes have already been segmented into
+            #  lines, so this node contains a single line too long to send to
+            #  the LLM. If this happens, the source code is probably malformed.
+            # We have no choice but to log an error and simply ignore the node.
             if not node.children:
                 log.error(f"[{node.name}] Childless node too long for context!")
-            node.text = None
-            node.tokens = 0
-            return
+                continue
 
-        text_chunks = [c.complete_placeholder for c in node.children]
-        node.text = "".join(text_chunks)
-        node.tokens = self._count_tokens(node.text)
+            # Consolidate nodes into groups, and then merge each group into a new node
+            node_groups = self._group_nodes(node.children)
+            node.children = list(map(self.merge_nodes, node_groups))
 
-        # If the text is still too long even with every child replaced with
-        #  placeholders, there's no reason to bother with placeholders at all
-        if node.tokens > self.max_tokens:
-            node.text = None
-            node.tokens = 0
-            return
+            # "Recurse" by pushing the children onto the stack
+            stack.extend(node.children)
 
-        sorted_indices: List[int] = sorted(
-            range(len(node.children)),
-            key=lambda idx: node.children[idx].tokens,
-        )
+    def _should_prune(self, node: CodeBlock) -> bool:
+        return node.node_type in self._prune_node_types
 
-        merged_child_indices = set()
-        for idx in sorted_indices:
-            child = node.children[idx]
-            text_chunks[idx] = child.complete_text
-            text = "".join(text_chunks)
-            tokens = self._count_tokens(text)
-            if tokens > self.max_tokens:
-                break
+    def _prune(self, root: CodeBlock) -> None:
+        stack = [root]
+        traversal = []
+        while stack:
+            node = stack.pop()
+            traversal.append(node)
+            node.children = [c for c in node.children if not self._should_prune(c)]
+            stack.extend(node.children)
 
-            node.text = text
-            node.tokens = tokens
-            merged_child_indices.add(idx)
+        for node in traversal[::-1]:
+            node.rebuild_text_from_children()
+            node.tokens = self._count_tokens(node.text)
 
-        # Remove all merged children from the child list
-        node.children = [
-            child
-            for i, child in enumerate(node.children)
-            if i not in merged_child_indices
-        ]
+    def _is_protected(self, node: CodeBlock) -> bool:
+        return node.node_type in self._protected_node_types
+
+    def _has_protected_descendents(self, node: CodeBlock) -> bool:
+        if not self._protected_node_types:
+            return False
+
+        queue = [*node.children]
+        while queue:
+            node = queue.pop(0)
+            if self._is_protected(node):
+                return True
+            queue.extend(node.children)
+        return False
+
+    def _prune_unprotected(self, root: CodeBlock) -> None:
+        if not self._has_protected_descendents(root):
+            if not self._is_protected(root):
+                raise EmptyTreeError("No protected nodes in tree!")
+            root.children = []
+            return
+
+        stack = [root]
+        while stack:
+            node = stack.pop()
+            if self._is_protected(node):
+                node.children = []
+            node.children = [
+                c
+                for c in node.children
+                if self._is_protected(c) or self._has_protected_descendents(c)
+            ]
+            stack.extend(node.children)
 
     def _group_nodes(self, nodes: List[CodeBlock]) -> List[List[CodeBlock]]:
         """Consolidate a list of tree_sitter nodes into groups. Each group should fit
@@ -197,11 +266,19 @@ class Splitter(FileManager):
         # Estimate the length of each adjacent pair were they merged
         adj_sums = [lengths[i] + lengths[i + 1] for i in range(len(lengths) - 1)]
 
+        # Create list of booleans parallel with adj_sums indicating whether that
+        #  merge is allowed (according to the protected node types list)
+        protected = list(map(self._is_protected, nodes))
+        merge_allowed = [
+            not (protected[i] or protected[i + 1]) for i in range(len(protected) - 1)
+        ]
+
         groups = [[n] for n in nodes]
-        while len(groups) > 1 and min(adj_sums) <= self.max_tokens:
+        while len(groups) > 1 and min(adj_sums) <= self.max_tokens and any(merge_allowed):
             # Get the indices of the adjacent nodes that would result in the
-            #  smallest possible merged snippet
-            i0 = int(min(range(len(adj_sums)), key=adj_sums.__getitem__))
+            #  smallest possible merged snippet. Ignore protected nodes.
+            mergeable_indices = compress(range(len(adj_sums)), merge_allowed)
+            i0 = int(min(mergeable_indices, key=adj_sums.__getitem__))
             i1 = i0 + 1
 
             # Recalculate the length. We can't simply use the adj_sum, because
@@ -222,8 +299,13 @@ class Splitter(FileManager):
             if i1 < len(adj_sums) - 1:
                 adj_sums[i1 + 1] += merged_text_length
 
+            if i0 > 0 and i1 < len(merge_allowed) - 1:
+                if not (merge_allowed[i0 - 1] and merge_allowed[i1 + 1]):
+                    merge_allowed[i0 - 1] = merge_allowed[i1 + 1] = False
+
             # The potential merge length for this pair is removed
             adj_sums.pop(i0)
+            merge_allowed.pop(i0)
 
             # Merge the pair of node groups
             groups[i0 : i1 + 1] = [groups[i0] + groups[i1]]
@@ -269,7 +351,7 @@ class Splitter(FileManager):
             start_byte=nodes[0].start_byte,
             end_byte=nodes[-1].end_byte,
             affixes=(prefix, suffix),
-            type=NodeType("merge"),
+            node_type=NodeType("merge"),
             children=sorted(sum([node.children for node in nodes], [])),
             language=language,
             tokens=tokens,
@@ -301,6 +383,9 @@ class Splitter(FileManager):
         if node.tokens <= self.max_tokens:
             return
 
+        if self._is_protected(node):
+            raise TokenLimitError(r"Irreducible node too large for context!")
+
         if node.children:
             for child in node.children:
                 self._segment_leaves(child)
@@ -309,9 +394,19 @@ class Splitter(FileManager):
         if node.start_point is None or node.end_point is None:
             raise ValueError("Node has no start or end point")
 
+        self._split_into_lines(node)
+
+    def _split_into_lines(self, node: CodeBlock):
         split_text = re.split(r"(\n+)", node.text)
-        betweens = split_text[1::2]
-        lines = split_text[::2]
+
+        # If the string didn't start/end with newlines, make sure to include
+        #  empty strings for the prefix/suffixes
+        if split_text[0].strip("\n"):
+            split_text = [""] + split_text
+        if split_text[-1].strip("\n"):
+            split_text.append("")
+        betweens = split_text[::2]
+        lines = split_text[1::2]
 
         start_byte = node.start_byte
         node_line = 0
@@ -322,7 +417,7 @@ class Splitter(FileManager):
             end_byte = start_byte + len(bytes(line, "utf-8"))
             end_char = len(line)
 
-            name = f"{node.name}L#{node_line}"
+            name = f"{node.name}-L#{node_line}"
             tokens = self._count_tokens(line)
             if tokens > self.max_tokens:
                 raise TokenLimitError(r"Irreducible node too large for context!")
@@ -337,14 +432,13 @@ class Splitter(FileManager):
                     start_byte=start_byte,
                     end_byte=end_byte,
                     affixes=(prefix, suffix),
-                    type=NodeType("segment"),
+                    node_type=NodeType(f"{node.node_type}__segment"),
                     children=[],
                     language=self.language,
                     tokens=tokens,
                 )
             )
-            start_byte = end_byte + len(suffix)
-            node_line += len(suffix)
+            start_byte = end_byte
 
         # Keep the first child's prefix
         node.children[0].omit_prefix = False

janus/language/treesitter/_tests/test_treesitter.py

@@ -11,7 +11,7 @@ class TestTreeSitterSplitter(unittest.TestCase):
 
     def setUp(self):
         """Set up the tests."""
-        model_name = "gpt-3.5-turbo"
+        model_name = "gpt-3.5-turbo-0125"
         self.maxDiff = None
         self.llm, _, _ = load_model(model_name)
 
@@ -31,7 +31,13 @@ class TestTreeSitterSplitter(unittest.TestCase):
         self.test_file = Path("janus/language/treesitter/_tests/languages/fortran.f90")
         self._split()
 
-        self.splitter.use_placeholders = False
+    def test_split_ibmhlasm(self):
+        """Test the split method."""
+        self.splitter = TreeSitterSplitter(
+            language="ibmhlasm", model=self.llm, max_tokens=100
+        )
+        self.combiner = Combiner(language="ibmhlasm")
+        self.test_file = Path("janus/language/treesitter/_tests/languages/ibmhlasm.asm")
         self._split()
 
     def test_split_matlab(self):
@@ -39,13 +45,10 @@ class TestTreeSitterSplitter(unittest.TestCase):
         self.splitter = TreeSitterSplitter(
             language="matlab",
             model=self.llm,
-            max_tokens=(4096 // 3)
+            max_tokens=(4096 // 3),
             # max_tokens used to be / 3 always in TreeSitterSplitter to leave just as
             # much space for the prompt as for the translated code.
         )
         self.combiner = Combiner(language="matlab")
         self.test_file = Path("janus/language/treesitter/_tests/languages/matlab.m")
         self._split()
-
-        self.splitter.use_placeholders = False
-        self._split()

janus/language/treesitter/treesitter.py

@@ -2,6 +2,7 @@ import os
 import platform
 from collections import defaultdict
 from pathlib import Path
+from typing import Optional
 
 import tree_sitter
 from git import Repo
@@ -9,7 +10,7 @@ from langchain.schema.language_model import BaseLanguageModel
 
 from ...utils.enums import LANGUAGES
 from ...utils.logger import create_logger
-from ..block import CodeBlock
+from ..block import CodeBlock, NodeType
 from ..splitter import Splitter
 
 log = create_logger(__name__)
@@ -25,7 +26,9 @@ class TreeSitterSplitter(Splitter):
         language: str,
         model: None | BaseLanguageModel = None,
         max_tokens: int = 4096,
-        use_placeholders: bool = False,
+        protected_node_types: tuple[str] = (),
+        prune_node_types: tuple[str] = (),
+        prune_unprotected: bool = False,
     ) -> None:
         """Initialize a TreeSitterSplitter instance.
 
@@ -38,26 +41,39 @@ class TreeSitterSplitter(Splitter):
             language=language,
             model=model,
             max_tokens=max_tokens,
-            use_placeholders=use_placeholders,
+            protected_node_types=protected_node_types,
+            prune_node_types=prune_node_types,
+            prune_unprotected=prune_unprotected,
         )
         self._load_parser()
 
     def _get_ast(self, code: str) -> CodeBlock:
         code = bytes(code, "utf-8")
-
         tree = self.parser.parse(code)
         root = tree.walk().node
         root = self._node_to_block(root, code)
         return root
 
-    def _set_identifiers(self, root: CodeBlock, path: Path):
+    # Recursively print tree to view parsed output (dev helper function)
+    # Example call: self._print_tree(tree.walk(), "")
+    def _print_tree(self, cursor: tree_sitter.TreeCursor, indent: str) -> None:
+        node = cursor.node
+        print(f"{indent}{node.type} {node.start_point}-{node.end_point}")
+        if cursor.goto_first_child():
+            while True:
+                self._print_tree(cursor, indent + "    ")
+                if not cursor.goto_next_sibling():
+                    break
+            cursor.goto_parent()
+
+    def _set_identifiers(self, root: CodeBlock, name: str):
         seen_types = defaultdict(int)
         queue = [root]
         while queue:
             node = queue.pop(0)  # BFS order to keep lower IDs toward the root
-            node.id = f"{node.type}[{seen_types[node.type]}]"
-            seen_types[node.type] += 1
-            node.name = f"{path.name}:{node.id}"
+            node.id = f"{node.node_type}[{seen_types[node.node_type]}]"
+            seen_types[node.node_type] += 1
+            node.name = f"{name}:{node.id}"
             queue.extend(node.children)
 
     def _node_to_block(self, node: tree_sitter.Node, original_text: bytes) -> CodeBlock:
@@ -84,14 +100,14 @@ class TreeSitterSplitter(Splitter):
         children = [self._node_to_block(child, original_text) for child in node.children]
         node = CodeBlock(
             id=node.id,
-            name=node.id,
+            name=str(node.id),
             text=text,
             affixes=(prefix, suffix),
             start_point=node.start_point,
             end_point=node.end_point,
             start_byte=node.start_byte,
             end_byte=node.end_byte,
-            type=node.type,
+            node_type=NodeType(node.type),
             children=children,
             language=self.language,
             tokens=self._count_tokens(text),
@@ -142,14 +158,22 @@ class TreeSitterSplitter(Splitter):
                 message = f"Tree-sitter does not support {self.language} yet."
                 log.error(message)
                 raise ValueError(message)
-            self._git_clone(github_url, lang_dir)
+            if LANGUAGES[self.language].get("branch"):
+                self._git_clone(github_url, lang_dir, LANGUAGES[self.language]["branch"])
+            else:
+                self._git_clone(github_url, lang_dir)
 
         tree_sitter.Language.build_library(str(so_file), [str(lang_dir)])
 
     @staticmethod
-    def _git_clone(repository_url: str, destination_folder: Path | str) -> None:
+    def _git_clone(
+        repository_url: str, destination_folder: Path | str, branch: Optional[str] = None
+    ) -> None:
         try:
-            Repo.clone_from(repository_url, destination_folder)
+            if branch:
+                Repo.clone_from(repository_url, destination_folder, branch=branch)
+            else:
+                Repo.clone_from(repository_url, destination_folder)
             log.debug(f"{repository_url} cloned to {destination_folder}")
         except Exception as e:
             log.error(f"Error: {e}")