janus-llm 2.0.2__tar.gz → 2.1.0__tar.gz

{janus_llm-2.0.2 → janus_llm-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: janus-llm
-Version: 2.0.2
+Version: 2.1.0
 Summary: A transcoding library using LLMs.
 Home-page: https://github.com/janus-llm/janus-llm
 License: Apache 2.0

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/__init__.py

@@ -5,7 +5,7 @@ from langchain_core._api.deprecation import LangChainDeprecationWarning
 from .metrics import *  # noqa: F403
 from .translate import Translator
 
-__version__ = "2.0.2"
+__version__ = "2.1.0"
 
 # Ignoring a deprecation warning from langchain_core that I can't seem to hunt down
 warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/cli.py

@@ -12,8 +12,6 @@ from rich.console import Console
 from rich.prompt import Confirm
 from typing_extensions import Annotated
 
-from janus.language.naive.registry import CUSTOM_SPLITTERS
-
 from .embedding.collections import Collections
 from .embedding.database import ChromaEmbeddingDatabase
 from .embedding.embedding_models_info import (
@@ -25,6 +23,7 @@ from .embedding.embedding_models_info import (
 from .embedding.vectorize import ChromaDBVectorizer
 from .language.binary import BinarySplitter
 from .language.mumps import MumpsSplitter
+from .language.naive.registry import CUSTOM_SPLITTERS
 from .language.treesitter import TreeSitterSplitter
 from .llm.model_callbacks import COST_PER_1K_TOKENS
 from .llm.models_info import MODEL_CONFIG_DIR, MODEL_TYPE_CONSTRUCTORS, TOKEN_LIMITS

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/converter.py

@@ -3,6 +3,7 @@ from typing import Any
 
 from langchain.schema.language_model import BaseLanguageModel
 
+from .language.alc.alc import AlcSplitter
 from .language.binary import BinarySplitter
 from .language.mumps import MumpsSplitter
 from .language.splitter import Splitter
@@ -152,6 +153,8 @@ class Converter:
         if self._source_language in CUSTOM_SPLITTERS:
             if self._source_language == "mumps":
                 self._splitter = MumpsSplitter(**kwargs)
+            elif self._source_language == "ibmhlasm":
+                self._splitter = AlcSplitter(**kwargs)
             elif self._source_language == "binary":
                 self._splitter = BinarySplitter(**kwargs)
         else:

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/embedding/_tests/test_collections.py

@@ -4,8 +4,8 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from janus.embedding.collections import Collections
-from janus.utils.enums import EmbeddingType
+from ...utils.enums import EmbeddingType
+from ..collections import Collections
 
 
 class TestCollections(unittest.TestCase):

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/_tests/test_splitter.py

@@ -1,6 +1,6 @@
 import unittest
 
-from janus.language.splitter import Splitter
+from ..splitter import Splitter
 
 
 class TestSplitter(unittest.TestCase):

janus_llm-2.1.0/janus/language/alc/__init__.py

@@ -0,0 +1 @@
+from .alc import AlcCombiner, AlcSplitter

janus_llm-2.1.0/janus/language/alc/_tests/test_alc.py

@@ -0,0 +1,28 @@
+import unittest
+from pathlib import Path
+
+from ....llm import load_model
+from ...combine import Combiner
+from ..alc import AlcSplitter
+
+
+class TestAlcSplitter(unittest.TestCase):
+    """Tests for the Splitter class."""
+
+    def setUp(self):
+        """Set up the tests."""
+        model_name = "gpt-3.5-turbo-0125"
+        llm, _, _ = load_model(model_name)
+        self.splitter = AlcSplitter(model=llm)
+        self.combiner = Combiner(language="ibmhlasm")
+        self.test_file = Path("janus/language/alc/_tests/alc.asm")
+
+    def test_split(self):
+        """Test the split method."""
+        tree_root = self.splitter.split(self.test_file)
+        self.assertEqual(tree_root.n_descendents, 34)
+        self.assertLessEqual(tree_root.max_tokens, self.splitter.max_tokens)
+        self.assertFalse(tree_root.complete)
+        self.combiner.combine_children(tree_root)
+        self.assertTrue(tree_root.complete)
+        self.assertEqual(tree_root.complete_text, self.test_file.read_text())

janus_llm-2.1.0/janus/language/alc/alc.py

@@ -0,0 +1,87 @@
+from langchain.schema.language_model import BaseLanguageModel
+
+from ...utils.logger import create_logger
+from ..block import CodeBlock
+from ..combine import Combiner
+from ..node import NodeType
+from ..treesitter import TreeSitterSplitter
+
+log = create_logger(__name__)
+
+
+class AlcCombiner(Combiner):
+    """A class that combines code blocks into ALC files."""
+
+    def __init__(self) -> None:
+        """Initialize a AlcCombiner instance."""
+        super().__init__("ibmhlasm")
+
+
+class AlcSplitter(TreeSitterSplitter):
+    """A class for splitting ALC code into functional blocks to prompt
+    with for transcoding.
+    """
+
+    def __init__(
+        self,
+        model: None | BaseLanguageModel = None,
+        max_tokens: int = 4096,
+        protected_node_types: tuple[str, ...] = (),
+        prune_node_types: tuple[str, ...] = (),
+        prune_unprotected: bool = False,
+    ):
+        """Initialize a AlcSplitter instance.
+
+        Arguments:
+            max_tokens: The maximum number of tokens supported by the model
+        """
+        super().__init__(
+            language="ibmhlasm",
+            model=model,
+            max_tokens=max_tokens,
+            protected_node_types=protected_node_types,
+            prune_node_types=prune_node_types,
+            prune_unprotected=prune_unprotected,
+        )
+
+    def _get_ast(self, code: str) -> CodeBlock:
+        root = super()._get_ast(code)
+
+        # Current treesitter implementation does not nest csects and dsects
+        # The loop below nests nodes following csect/dsect instructions into
+        #  the children of that instruction
+        sect_types = {"csect_instruction", "dsect_instruction"}
+        queue: list[CodeBlock] = [root]
+        while queue:
+            block = queue.pop(0)
+
+            # Search this children for csects and dsects. Create a list of groups
+            #  where each group is a csect or dsect, starting with the csect/dsect
+            #  instruction and containing all the subsequent nodes up until the
+            #  next csect or dsect instruction
+            sects: list[list[CodeBlock]] = [[]]
+            for c in block.children:
+                if c.node_type in sect_types:
+                    sects.append([c])
+                else:
+                    sects[-1].append(c)
+
+            sects = [s for s in sects if s]
+
+            # Restructure the tree, making the head of each group the parent
+            #  of all the remaining nodes in that group
+            if len(sects) > 1:
+                block.children = []
+                for sect in sects:
+                    if sect[0].node_type in sect_types:
+                        sect_node = self.merge_nodes(sect)
+                        sect_node.children = sect
+                        sect_node.node_type = NodeType(str(sect[0].node_type)[:5])
+                        block.children.append(sect_node)
+                    else:
+                        block.children.extend(sect)
+
+            # Push the children onto the queue
+            queue.extend(block.children)
+
+        return root

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/block.py

@@ -152,9 +152,11 @@ class CodeBlock:
         Returns:
             A string representation of the tree with this block as the root
         """
+        tokens = self.tokens
         identifier = self.id
         if self.text is None:
             identifier = f"({identifier})"
+            tokens = self.total_tokens
         elif not self.complete:
             identifier += "*"
         if self.start_point is not None and self.end_point is not None:
@@ -165,7 +167,7 @@ class CodeBlock:
             seg = ""
         return "\n".join(
             [
-                f"{'| '*depth}{identifier}{seg}",
+                f"{'| '*depth}{identifier}{seg}  ({tokens:,d} tokens)",
                 *[c.tree_str(depth + 1) for c in self.children],
             ]
         )

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/mumps/mumps.py

@@ -48,6 +48,7 @@ class MumpsSplitter(Splitter):
         max_tokens: int = 4096,
         protected_node_types: tuple[str] = ("routine_definition",),
         prune_node_types: tuple[str] = (),
+        prune_unprotected: bool = False,
     ):
         """Initialize a MumpsSplitter instance.
 
@@ -60,11 +61,9 @@ class MumpsSplitter(Splitter):
             max_tokens=max_tokens,
             protected_node_types=protected_node_types,
             prune_node_types=prune_node_types,
+            prune_unprotected=prune_unprotected,
         )
 
-        # MUMPS code tends to take about 2/3 the space of Python
-        self.max_tokens: int = int(max_tokens * 2 / 5)
-
     def _set_identifiers(self, root: CodeBlock, name: str):
         stack = [root]
         while stack:

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/naive/__init__.py

@@ -1,4 +1,4 @@
 from .basic_splitter import FileSplitter
 from .chunk_splitter import ChunkSplitter
-from .simple_ast import FlexibleTreeSitterSplitter, StrictTreeSitterSplitter
+from .simple_ast import get_flexible_ast, get_strict_ast
 from .tag_splitter import TagSplitter

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/naive/basic_splitter.py

@@ -1,7 +1,7 @@
-from janus.language.block import CodeBlock
-from janus.language.naive.chunk_splitter import ChunkSplitter
-from janus.language.naive.registry import register_splitter
-from janus.language.splitter import FileSizeError
+from ..block import CodeBlock
+from ..naive.chunk_splitter import ChunkSplitter
+from ..naive.registry import register_splitter
+from ..splitter import FileSizeError
 
 
 @register_splitter("file")

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/naive/chunk_splitter.py

@@ -1,7 +1,7 @@
-from janus.language.block import CodeBlock
-from janus.language.naive.registry import register_splitter
-from janus.language.node import NodeType
-from janus.language.splitter import Splitter
+from ..block import CodeBlock
+from ..node import NodeType
+from ..splitter import Splitter
+from .registry import register_splitter
 
 
 @register_splitter("chunk")

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/naive/registry.py

@@ -1,6 +1,6 @@
 from typing import Callable, Dict
 
-from janus.language.splitter import Splitter
+from ..splitter import Splitter
 
 CUSTOM_SPLITTERS: Dict[str, Callable[..., Splitter]] = dict()
 

janus_llm-2.1.0/janus/language/naive/simple_ast.py

@@ -0,0 +1,29 @@
+from ...utils.enums import LANGUAGES
+from ..alc.alc import AlcSplitter
+from ..mumps.mumps import MumpsSplitter
+from ..treesitter import TreeSitterSplitter
+from .registry import register_splitter
+
+
+@register_splitter("ast-flex")
+def get_flexible_ast(language: str, **kwargs):
+    if language == "ibmhlasm":
+        return AlcSplitter(**kwargs)
+    elif language == "mumps":
+        return MumpsSplitter(**kwargs)
+    else:
+        return TreeSitterSplitter(language=language, **kwargs)
+
+
+@register_splitter("ast-strict")
+def get_strict_ast(language: str, **kwargs):
+    kwargs.update(
+        protected_node_types=LANGUAGES[language]["functional_node_types"],
+        prune_unprotected=True,
+    )
+    if language == "ibmhlasm":
+        return AlcSplitter(**kwargs)
+    elif language == "mumps":
+        return MumpsSplitter(**kwargs)
+    else:
+        return TreeSitterSplitter(language=language, **kwargs)

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/naive/tag_splitter.py

@@ -1,7 +1,7 @@
-from janus.language.block import CodeBlock
-from janus.language.naive.registry import register_splitter
-from janus.language.node import NodeType
-from janus.language.splitter import Splitter
+from ..block import CodeBlock
+from ..node import NodeType
+from ..splitter import Splitter
+from .registry import register_splitter
 
 
 @register_splitter("tag")

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/splitter.py

@@ -47,8 +47,8 @@ class Splitter(FileManager):
         model: None | BaseLanguageModel = None,
         max_tokens: int = 4096,
         skip_merge: bool = False,
-        protected_node_types: tuple[str] = (),
-        prune_node_types: tuple[str] = (),
+        protected_node_types: tuple[str, ...] = (),
+        prune_node_types: tuple[str, ...] = (),
         prune_unprotected: bool = False,
     ):
         """
@@ -340,7 +340,10 @@ class Splitter(FileManager):
         # Double check length (in theory this should never be an issue)
         tokens = self._count_tokens(text)
         if tokens > self.max_tokens:
-            log.error(f"Merged node ({name}) too long for context!")
+            log.error(
+                f"Merged node ({name}) too long for context!"
+                f" ({tokens} > {self.max_tokens})"
+            )
 
         return CodeBlock(
             text=text,
@@ -420,7 +423,10 @@ class Splitter(FileManager):
             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!")
+                raise TokenLimitError(
+                    "Irreducible node too large for context!"
+                    f" ({tokens} > {self.max_tokens})"
+                )
 
             node.children.append(
                 CodeBlock(

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/language/treesitter/treesitter.py

@@ -26,8 +26,8 @@ class TreeSitterSplitter(Splitter):
         language: str,
         model: None | BaseLanguageModel = None,
         max_tokens: int = 4096,
-        protected_node_types: tuple[str] = (),
-        prune_node_types: tuple[str] = (),
+        protected_node_types: tuple[str, ...] = (),
+        prune_node_types: tuple[str, ...] = (),
         prune_unprotected: bool = False,
     ) -> None:
         """Initialize a TreeSitterSplitter instance.
@@ -48,10 +48,10 @@ class TreeSitterSplitter(Splitter):
         self._load_parser()
 
     def _get_ast(self, code: str) -> CodeBlock:
-        code = bytes(code, "utf-8")
-        tree = self.parser.parse(code)
+        code_bytes = bytes(code, "utf-8")
+        tree = self.parser.parse(code_bytes)
         root = tree.walk().node
-        root = self._node_to_block(root, code)
+        root = self._node_to_block(root, code_bytes)
         return root
 
     # Recursively print tree to view parsed output (dev helper function)
@@ -98,7 +98,7 @@ class TreeSitterSplitter(Splitter):
 
         text = node.text.decode()
         children = [self._node_to_block(child, original_text) for child in node.children]
-        node = CodeBlock(
+        return CodeBlock(
             id=node.id,
             name=str(node.id),
             text=text,
@@ -112,7 +112,6 @@ class TreeSitterSplitter(Splitter):
             language=self.language,
             tokens=self._count_tokens(text),
         )
-        return node
 
     def _load_parser(self) -> None:
         """Load the parser for the given language.

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/llm/model_callbacks.py

@@ -8,7 +8,7 @@ from langchain_core.messages import AIMessage
 from langchain_core.outputs import ChatGeneration, LLMResult
 from langchain_core.tracers.context import register_configure_hook
 
-from janus.utils.logger import create_logger
+from ..utils.logger import create_logger
 
 log = create_logger(__name__)
 

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/llm/models_info.py

@@ -8,8 +8,7 @@ from langchain_community.llms import HuggingFaceTextGenInference
 from langchain_core.language_models import BaseLanguageModel
 from langchain_openai import ChatOpenAI
 
-from janus.llm.model_callbacks import COST_PER_1K_TOKENS
-from janus.prompts.prompt import (
+from ..prompts.prompt import (
     ChatGptPromptEngine,
     ClaudePromptEngine,
     CoherePromptEngine,
@@ -18,8 +17,8 @@ from janus.prompts.prompt import (
     PromptEngine,
     TitanPromptEngine,
 )
-
 from ..utils.logger import create_logger
+from .model_callbacks import COST_PER_1K_TOKENS
 
 log = create_logger(__name__)
 

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/metrics/_tests/test_llm.py

@@ -3,8 +3,7 @@ from unittest.mock import patch
 
 import pytest
 
-from janus.llm.models_info import load_model
-
+from ...llm.models_info import load_model
 from ..llm_metrics import llm_evaluate_option, llm_evaluate_ref_option
 
 
@@ -40,7 +39,7 @@ class TestLLMMetrics(unittest.TestCase):
         print("'Hello, world!")
         """
 
-    @patch("janus.llm.models_info.load_model")
+    @patch(".llm.models_info.load_model")
     @patch("janus.metrics.llm_metrics.llm_evaluate")
     @pytest.mark.llm_eval
     def test_llm_self_eval_quality(self, mock_llm_evaluate, mock_load_model):

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/metrics/_tests/test_rouge_score.py

@@ -1,6 +1,6 @@
 import unittest
 
-from janus.metrics.rouge_score import rouge
+from ..rouge_score import rouge
 
 
 class TestRouge(unittest.TestCase):

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/metrics/_tests/test_similarity_score.py

@@ -1,6 +1,6 @@
 import unittest
 
-from janus.metrics.similarity import similarity_score
+from ..similarity import similarity_score
 
 
 class TestSimilarityScore(unittest.TestCase):

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/metrics/complexity_metrics.py

@@ -1,10 +1,9 @@
 import math
 from typing import List, Optional
 
-from janus.language.block import CodeBlock
-from janus.language.treesitter.treesitter import TreeSitterSplitter
-from janus.utils.enums import LANGUAGES
-
+from ..language.block import CodeBlock
+from ..language.treesitter.treesitter import TreeSitterSplitter
+from ..utils.enums import LANGUAGES
 from .metric import metric
 
 

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/metrics/metric.py

@@ -7,10 +7,9 @@ import click
 import typer
 from typing_extensions import Annotated
 
-from janus.llm import load_model
-from janus.utils.enums import LANGUAGES
-from janus.utils.logger import create_logger
-
+from ..llm import load_model
+from ..utils.enums import LANGUAGES
+from ..utils.logger import create_logger
 from ..utils.progress import track
 from .cli import evaluate
 from .file_pairing import FILE_PAIRING_METHODS

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/prompts/prompt.py

@@ -34,6 +34,40 @@ HUMAN_PROMPT_TEMPLATE_FILENAME = "human.txt"
 PROMPT_VARIABLES_FILENAME = "variables.json"
 
 
+retry_with_output_prompt_text = """Instructions:
+--------------
+{instructions}
+--------------
+Completion:
+--------------
+{input}
+--------------
+
+Above, the Completion did not satisfy the constraints given in the Instructions.
+Error:
+--------------
+{error}
+--------------
+
+Please try again. Please only respond with an answer that satisfies the
+constraints laid out in the Instructions:"""
+
+
+retry_with_error_and_output_prompt_text = """Prompt:
+{prompt}
+Completion:
+{input}
+
+Above, the Completion did not satisfy the constraints given in the Prompt.
+Details: {error}
+Please try again:"""
+
+retry_with_output_prompt = PromptTemplate.from_template(retry_with_output_prompt_text)
+retry_with_error_and_output_prompt = PromptTemplate.from_template(
+    retry_with_error_and_output_prompt_text
+)
+
+
 class PromptEngine(ABC):
     """A class defining prompting schemes for the LLM."""
 

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/translate.py

@@ -16,12 +16,11 @@ from langchain_core.runnables import RunnableLambda, RunnableParallel
 from openai import BadRequestError, RateLimitError
 from text_generation.errors import ValidationError
 
-from janus.language.naive.registry import CUSTOM_SPLITTERS
-
 from .converter import Converter, run_if_changed
 from .embedding.vectorize import ChromaDBVectorizer
 from .language.block import CodeBlock, TranslatedCodeBlock
 from .language.combine import ChunkCombiner, Combiner, JsonCombiner
+from .language.naive.registry import CUSTOM_SPLITTERS
 from .language.splitter import EmptyTreeError, FileSizeError, TokenLimitError
 from .llm import load_model
 from .llm.model_callbacks import get_model_callback
@@ -30,7 +29,12 @@ from .parsers.code_parser import CodeParser, GenericParser
 from .parsers.doc_parser import MadlibsDocumentationParser, MultiDocumentationParser
 from .parsers.eval_parser import EvaluationParser
 from .parsers.reqs_parser import RequirementsParser
-from .prompts.prompt import SAME_OUTPUT, TEXT_OUTPUT
+from .prompts.prompt import (
+    SAME_OUTPUT,
+    TEXT_OUTPUT,
+    retry_with_error_and_output_prompt,
+    retry_with_output_prompt,
+)
 from .utils.enums import LANGUAGES
 from .utils.logger import create_logger
 
@@ -407,10 +411,10 @@ class Translator(Converter):
         """
         self._parser.set_reference(block.original)
 
-        # Retries with just the output and the error
+        # Retries with just the format instructions, the output, and the error
         n1 = round(self.max_prompts ** (1 / 3))
 
-        # Retries with the input, output, and error
+        # Retries with the input, the output, and the error
         n2 = round((self.max_prompts // n1) ** (1 / 2))
 
         # Retries with just the input
@@ -420,11 +424,13 @@ class Translator(Converter):
             llm=self._llm,
             parser=self._parser,
             max_retries=n1,
+            prompt=retry_with_output_prompt,
         )
         retry = RetryWithErrorOutputParser.from_llm(
             llm=self._llm,
             parser=fix_format,
             max_retries=n2,
+            prompt=retry_with_error_and_output_prompt,
         )
 
         completion_chain = self._prompt | self._llm

{janus_llm-2.0.2 → janus_llm-2.1.0}/janus/utils/enums.py

@@ -10,7 +10,7 @@ class EmbeddingType(Enum):
     TARGET = 5  # placeholder embeddings, are these useful for analysis?
 
 
-CUSTOM_SPLITTERS: Set[str] = {"mumps", "binary"}
+CUSTOM_SPLITTERS: Set[str] = {"mumps", "binary", "ibmhlasm"}
 
 LANGUAGES: Dict[str, Dict[str, Any]] = {
     "ada": {
@@ -63,7 +63,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
             '#include <stdio.h>\n\nint main() {\n    printf("Hello, World!\\n");\n'
             "    return 0;\n}\n"
         ),
-        "functional_node_type": "function_definition",
+        "functional_node_types": ["function_definition"],
         "comment_node_type": "comment",
     },
     "capnp": {
@@ -206,7 +206,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
         "example": (
             "program HelloWorld\n  print *, 'Hello, World!'\nend program HelloWorld\n"
         ),
-        "functional_node_type": "function",
+        "functional_node_types": ["function"],
         "comment_node_type": "comment",
     },
     "gitattributes": {
@@ -300,6 +300,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
                      END   HELLO
         """
         ),
+        "functional_node_types": ["csect", "dsect"],
         "branch_node_types": ["branch_instruction"],
         "operation_node_types": ["operation", "branch_operation"],
         "operand_node_types": ["operands"],
@@ -420,7 +421,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
         "suffix": "m",
         "url": "https://github.com/janus-llm/tree-sitter-mumps",
         "example": 'WRITE "Hello, World!"',
-        "functional_node_type": "routine_definition",
+        "functional_node_types": ["routine_definition"],
         "comment_node_type": "comment",
         "branch_node_types": ["if_statement"],
         "operation_node_types": [
@@ -512,7 +513,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
         "suffix": "py",
         "url": "https://github.com/tree-sitter/tree-sitter-python",
         "example": "# Hello, World!\nprint('Hello, World!')\n",
-        "functional_node_type": "function_definition",
+        "functional_node_types": ["function_definition"],
         "comment_node_type": "comment",
     },
     "qmljs": {

{janus_llm-2.0.2 → janus_llm-2.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "janus-llm"
-version = "2.0.2"
+version = "2.1.0"
 description = "A transcoding library using LLMs."
 authors = ["Michael Doyle <mdoyle@mitre.org>", "Chris Glasz <cglasz@mitre.org>",
            "Chris Tohline <ctohline@mitre.org>", "William Macke <wmacke@mitre.org>",

janus_llm-2.0.2/janus/language/naive/simple_ast.py

@@ -1,18 +0,0 @@
-from janus.language.naive.registry import register_splitter
-from janus.language.treesitter import TreeSitterSplitter
-from janus.utils.enums import LANGUAGES
-
-
-@register_splitter("ast-flex")
-class FlexibleTreeSitterSplitter(TreeSitterSplitter):
-    pass
-
-
-@register_splitter("ast-strict")
-class StrictTreeSitterSplitter(TreeSitterSplitter):
-    def __init__(self, language: str, **kwargs):
-        kwargs.update(
-            protected_node_types=(LANGUAGES[language]["functional_node_type"],),
-            prune_unprotected=True,
-        )
-        super().__init__(language=language, **kwargs)