janus-llm 4.2.0__tar.gz → 4.3.1__tar.gz

{janus_llm-4.2.0 → janus_llm-4.3.1}/PKG-INFO

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

{janus_llm-4.2.0 → janus_llm-4.3.1}/janus/__init__.py

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

{janus_llm-4.2.0 → janus_llm-4.3.1}/janus/cli.py

@@ -19,6 +19,7 @@ from janus.converter.aggregator import Aggregator
 from janus.converter.converter import Converter
 from janus.converter.diagram import DiagramGenerator
 from janus.converter.document import Documenter, MadLibsDocumenter, MultiDocumenter
+from janus.converter.evaluate import InlineCommentEvaluator, RequirementEvaluator
 from janus.converter.partition import Partitioner
 from janus.converter.requirements import RequirementsDocumenter
 from janus.converter.translate import Translator
@@ -127,7 +128,7 @@ embedding = typer.Typer(
 
 def version_callback(value: bool) -> None:
     if value:
-        from janus import __version__ as version
+        from . import __version__ as version
 
         print(f"Janus CLI [blue]v{version}[/blue]")
         raise typer.Exit()
@@ -655,6 +656,16 @@ def partition(
             click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
         ),
     ] = "file",
+    refiner_types: Annotated[
+        list[str],
+        typer.Option(
+            "-r",
+            "--refiner",
+            help="List of refiner types to use. Add -r for each refiner to use in\
+                refinement chain",
+            click_type=click.Choice(list(REFINERS.keys())),
+        ),
+    ] = ["JanusRefiner"],
     max_tokens: Annotated[
         int,
         typer.Option(
@@ -673,6 +684,7 @@ def partition(
         ),
     ] = 8192,
 ):
+    refiner_types = [REFINERS[r] for r in refiner_types]
     model_arguments = dict(temperature=temperature)
     kwargs = dict(
         model=llm_name,
@@ -681,6 +693,7 @@ def partition(
         max_prompts=max_prompts,
         max_tokens=max_tokens,
         splitter_type=splitter_type,
+        refiner_types=refiner_types,
         partition_token_limit=partition_token_limit,
     )
     partitioner = Partitioner(**kwargs)
@@ -815,6 +828,139 @@ def diagram(
     diagram_generator.translate(input_dir, output_dir, overwrite, collection)
 
 
+@app.command(
+    help="LLM self evaluation",
+    no_args_is_help=True,
+)
+def llm_self_eval(
+    input_dir: Annotated[
+        Path,
+        typer.Option(
+            "--input",
+            "-i",
+            help="The directory containing the source code to be evaluated. "
+            "The files should all be in one flat directory.",
+        ),
+    ],
+    language: Annotated[
+        str,
+        typer.Option(
+            "--language",
+            "-l",
+            help="The language of the source code.",
+            click_type=click.Choice(sorted(LANGUAGES)),
+        ),
+    ],
+    output_dir: Annotated[
+        Path,
+        typer.Option(
+            "--output-dir", "-o", help="The directory to store the evaluations in."
+        ),
+    ],
+    llm_name: Annotated[
+        str,
+        typer.Option(
+            "--llm",
+            "-L",
+            help="The custom name of the model set with 'janus llm add'.",
+        ),
+    ] = "gpt-4o",
+    evaluation_type: Annotated[
+        str,
+        typer.Option(
+            "--evaluation-type",
+            "-e",
+            help="Type of output to evaluate.",
+            click_type=click.Choice(["incose", "comments"]),
+        ),
+    ] = "incose",
+    max_prompts: Annotated[
+        int,
+        typer.Option(
+            "--max-prompts",
+            "-m",
+            help="The maximum number of times to prompt a model on one functional block "
+            "before exiting the application. This is to prevent wasting too much money.",
+        ),
+    ] = 10,
+    overwrite: Annotated[
+        bool,
+        typer.Option(
+            "--overwrite/--preserve",
+            help="Whether to overwrite existing files in the output directory",
+        ),
+    ] = False,
+    temperature: Annotated[
+        float,
+        typer.Option("--temperature", "-t", help="Sampling temperature.", min=0, max=2),
+    ] = 0.7,
+    collection: Annotated[
+        str,
+        typer.Option(
+            "--collection",
+            "-c",
+            help="If set, will put the translated result into a Chroma DB "
+            "collection with the name provided.",
+        ),
+    ] = None,
+    splitter_type: Annotated[
+        str,
+        typer.Option(
+            "-S",
+            "--splitter",
+            help="Name of custom splitter to use",
+            click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
+        ),
+    ] = "file",
+    refiner_types: Annotated[
+        list[str],
+        typer.Option(
+            "-r",
+            "--refiner",
+            help="List of refiner types to use. Add -r for each refiner to use in\
+                refinement chain",
+            click_type=click.Choice(list(REFINERS.keys())),
+        ),
+    ] = ["JanusRefiner"],
+    eval_items_per_request: Annotated[
+        int,
+        typer.Option(
+            "--eval-items-per-request",
+            "-rc",
+            help="The maximum number of evaluation items per request",
+        ),
+    ] = None,
+    max_tokens: Annotated[
+        int,
+        typer.Option(
+            "--max-tokens",
+            "-M",
+            help="The maximum number of tokens the model will take in. "
+            "If unspecificed, model's default max will be used.",
+        ),
+    ] = None,
+):
+    model_arguments = dict(temperature=temperature)
+    refiner_types = [REFINERS[r] for r in refiner_types]
+    kwargs = dict(
+        eval_items_per_request=eval_items_per_request,
+        model=llm_name,
+        model_arguments=model_arguments,
+        source_language=language,
+        max_prompts=max_prompts,
+        max_tokens=max_tokens,
+        splitter_type=splitter_type,
+        refiner_types=refiner_types,
+    )
+    # Setting parser type here
+    if evaluation_type == "incose":
+        evaluator = RequirementEvaluator(**kwargs)
+    elif evaluation_type == "comments":
+        evaluator = InlineCommentEvaluator(**kwargs)
+
+    evaluator.translate(input_dir, output_dir, overwrite, collection)
+
+
 @db.command("init", help="Connect to or create a database.")
 def db_init(
     path: Annotated[
@@ -1116,13 +1262,12 @@ def llm_add(
             show_choices=False,
         )
         params = dict(
-            # OpenAI uses the "model_name" key for what we're calling "long_model_id"
-            model_name=MODEL_ID_TO_LONG_ID[model_id],
+            model_name=model_name,
             temperature=0.7,
             n=1,
         )
-        max_tokens = TOKEN_LIMITS[MODEL_ID_TO_LONG_ID[model_id]]
-        model_cost = COST_PER_1K_TOKENS[MODEL_ID_TO_LONG_ID[model_id]]
+        max_tokens = TOKEN_LIMITS[model_name]
+        model_cost = COST_PER_1K_TOKENS[model_name]
         cfg = {
             "model_type": model_type,
             "model_id": model_id,

{janus_llm-4.2.0 → janus_llm-4.3.1}/janus/converter/converter.py

@@ -464,6 +464,7 @@ class Converter:
         for in_path, out_path in in_out_pairs:
             # Translate the file, skip it if there's a rate limit error
             try:
+                log.info(f"Processing {in_path.relative_to(input_directory)}")
                 out_block = self.translate_file(in_path)
                 total_cost += out_block.total_cost
             except RateLimitError:

janus_llm-4.3.1/janus/converter/evaluate.py

@@ -0,0 +1,241 @@
+import json
+import re
+from copy import deepcopy
+
+from langchain_core.runnables import Runnable, RunnableLambda, RunnableParallel
+
+from janus.converter.converter import Converter
+from janus.language.block import TranslatedCodeBlock
+from janus.language.combine import JsonCombiner
+from janus.parsers.eval_parsers.incose_parser import IncoseParser
+from janus.parsers.eval_parsers.inline_comment_parser import InlineCommentParser
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class Evaluator(Converter):
+    """Evaluator
+
+    A class that performs an LLM self evaluation"
+    "on an input target, with an associated prompt.
+
+    Current valid evaluation types:
+    ['incose', 'comments']
+
+    """
+
+    def __init__(self, **kwargs) -> None:
+        """Initialize the Evaluator class
+
+        Arguments:
+            model: The LLM to use for translation. If an OpenAI model, the
+                `OPENAI_API_KEY` environment variable must be set and the
+                `OPENAI_ORG_ID` environment variable should be set if needed.
+            model_arguments: Additional arguments to pass to the LLM constructor.
+            max_prompts: The maximum number of prompts to try before giving up.
+        """
+        super().__init__(**kwargs)
+        self._combiner = JsonCombiner()
+        self._load_parameters()
+
+
+class RequirementEvaluator(Evaluator):
+    """INCOSE Requirement Evaluator
+
+    A class that performs an LLM self evaluation on an input target,
+    with an associated prompt.
+
+    The evaluation prompts are for Incose Evaluations
+
+    """
+
+    def __init__(self, eval_items_per_request: int | None = None, **kwargs) -> None:
+        """Initialize the Evaluator class
+
+        Arguments:
+            model: The LLM to use for translation. If an OpenAI model, the
+                `OPENAI_API_KEY` environment variable must be set and the
+                `OPENAI_ORG_ID` environment variable should be set if needed.
+            model_arguments: Additional arguments to pass to the LLM constructor.
+            max_prompts: The maximum number of prompts to try before giving up.
+        """
+        super().__init__(**kwargs)
+        self.eval_items_per_request = eval_items_per_request
+        self._parser = IncoseParser()
+        self.set_prompt("eval_prompts/incose")
+
+    def _input_runnable(self) -> Runnable:
+        def _get_code(json_text: str) -> str:
+            return json.loads(json_text)["code"]
+
+        def _get_reqs(json_text: str) -> str:
+            return json.dumps(json.loads(json_text)["requirements"])
+
+        return RunnableLambda(self._parser.parse_input) | RunnableParallel(
+            SOURCE_CODE=_get_code,
+            REQUIREMENTS=_get_reqs,
+            context=self._retriever,
+        )
+
+    def _add_translation(self, block: TranslatedCodeBlock):
+        if block.translated:
+            return
+
+        if block.original.text is None:
+            block.translated = True
+            return
+
+        if self.eval_items_per_request is None:
+            return super()._add_translation(block)
+
+        input_obj = json.loads(block.original.text)
+        requirements = input_obj.get("requirements", [])
+
+        if not requirements:
+            log.debug(f"[{block.name}] Skipping empty block")
+            block.translated = True
+            block.text = None
+            block.complete = True
+            return
+
+        # For some reason requirements objects are in nested lists?
+        while isinstance(requirements[0], list):
+            requirements = [r for lst in requirements for r in lst]
+
+        if len(requirements) <= self.eval_items_per_request:
+            input_obj["requirements"] = requirements
+            block.original.text = json.dumps(input_obj)
+            return super()._add_translation(block)
+
+        block.processing_time = 0
+        block.cost = 0
+        block.retries = 0
+        obj = {}
+        for i in range(0, len(requirements), self.eval_items_per_request):
+            # Build a new TranslatedBlock using the new working text
+            working_requirements = requirements[i : i + self.eval_items_per_request]
+            working_copy = deepcopy(block.original)
+            working_obj = json.loads(working_copy.text)  # type: ignore
+            working_obj["requirements"] = working_requirements
+            working_copy.text = json.dumps(working_obj)
+            working_block = TranslatedCodeBlock(working_copy, self._target_language)
+
+            # Run the LLM on the working text
+            super()._add_translation(working_block)
+
+            # Update metadata to include for all runs
+            block.retries += working_block.retries
+            block.cost += working_block.cost
+            block.processing_time += working_block.processing_time
+
+            # Update the output text to merge this section's output in
+            obj.update(json.loads(working_block.text))
+
+        block.text = json.dumps(obj)
+        block.tokens = self._llm.get_num_tokens(block.text)
+        block.translated = True
+
+        log.debug(
+            f"[{block.name}] Output code:\n{json.dumps(json.loads(block.text), indent=2)}"
+        )
+
+
+class InlineCommentEvaluator(Evaluator):
+    """Inline Comment Evaluator
+
+    A class that performs an LLM self evaluation on inline comments,
+    with an associated prompt.
+    """
+
+    def __init__(self, eval_items_per_request: int | None = None, **kwargs) -> None:
+        """Initialize the Evaluator class
+
+        Arguments:
+            model: The LLM to use for translation. If an OpenAI model, the
+                `OPENAI_API_KEY` environment variable must be set and the
+                `OPENAI_ORG_ID` environment variable should be set if needed.
+            model_arguments: Additional arguments to pass to the LLM constructor.
+            max_prompts: The maximum number of prompts to try before giving up.
+        """
+        super().__init__(**kwargs)
+        self._combiner = JsonCombiner()
+        self._load_parameters()
+        self._parser = InlineCommentParser()
+        self.set_prompt("eval_prompts/inline_comments")
+        self.eval_items_per_request = eval_items_per_request
+
+    def _add_translation(self, block: TranslatedCodeBlock):
+        if block.translated:
+            return
+
+        if block.original.text is None:
+            block.translated = True
+            return
+
+        if self.eval_items_per_request is None:
+            return super()._add_translation(block)
+
+        comment_pattern = r"<(?:INLINE|BLOCK)_COMMENT \w{8}>.*$"
+        comments = list(
+            re.finditer(comment_pattern, block.original.text, flags=re.MULTILINE)
+        )
+
+        if not comments:
+            log.info(f"[{block.name}] Skipping commentless block")
+            block.translated = True
+            block.text = None
+            block.complete = True
+            return
+
+        if len(comments) <= self.eval_items_per_request:
+            return super()._add_translation(block)
+
+        comment_group_indices = list(range(0, len(comments), self.eval_items_per_request))
+        log.debug(
+            f"[{block.name}] Block contains more than {self.eval_items_per_request}"
+            f" comments, splitting {len(comments)} comments into"
+            f" {len(comment_group_indices)} groups"
+        )
+
+        block.processing_time = 0
+        block.cost = 0
+        block.retries = 0
+        obj = {}
+        for i in range(0, len(comments), self.eval_items_per_request):
+            # Split the text into the section containing comments of interest,
+            #  all the text prior to those comments, and all the text after them
+            working_comments = comments[i : i + self.eval_items_per_request]
+            start_idx = working_comments[0].start()
+            end_idx = working_comments[-1].end()
+            prefix = block.original.text[:start_idx]
+            keeper = block.original.text[start_idx:end_idx]
+            suffix = block.original.text[end_idx:]
+
+            # Strip all comment placeholders outside of the section of interest
+            prefix = re.sub(comment_pattern, "", prefix, flags=re.MULTILINE)
+            suffix = re.sub(comment_pattern, "", suffix, flags=re.MULTILINE)
+
+            # Build a new TranslatedBlock using the new working text
+            working_copy = deepcopy(block.original)
+            working_copy.text = prefix + keeper + suffix
+            working_block = TranslatedCodeBlock(working_copy, self._target_language)
+
+            # Run the LLM on the working text
+            super()._add_translation(working_block)
+
+            # Update metadata to include for all runs
+            block.retries += working_block.retries
+            block.cost += working_block.cost
+            block.processing_time += working_block.processing_time
+
+            # Update the output text to merge this section's output in
+            obj.update(json.loads(working_block.text))
+
+        block.text = json.dumps(obj)
+        block.tokens = self._llm.get_num_tokens(block.text)
+        block.translated = True
+
+        log.debug(
+            f"[{block.name}] Output code:\n{json.dumps(json.loads(block.text), indent=2)}"
+        )

{janus_llm-4.2.0 → janus_llm-4.3.1}/janus/language/alc/_tests/test_alc.py

@@ -20,7 +20,7 @@ class TestAlcSplitter(unittest.TestCase):
     def test_split(self):
         """Test the split method."""
         tree_root = self.splitter.split(self.test_file)
-        self.assertAlmostEqual(tree_root.n_descendents, 32, delta=5)
+        self.assertAlmostEqual(tree_root.n_descendents, 16, delta=2)
         self.assertLessEqual(tree_root.max_tokens, self.splitter.max_tokens)
         self.assertFalse(tree_root.complete)
         self.combiner.combine_children(tree_root)

{janus_llm-4.2.0 → janus_llm-4.3.1}/janus/language/alc/alc.py

@@ -79,10 +79,15 @@ class AlcSplitter(TreeSitterSplitter):
             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])
+                    node_type = sect[0].node_type
+                    if node_type in sect_types:
+                        if len(sect) == 1:
+                            # Don't make a node its own child
+                            sect_node = sect[0]
+                        else:
+                            sect_node = self.merge_nodes(sect)
+                            sect_node.children = sect
+                        sect_node.node_type = NodeType(str(node_type)[:5])
                         block.children.append(sect_node)
                     else:
                         block.children.extend(sect)

{janus_llm-4.2.0 → janus_llm-4.3.1}/janus/language/splitter.py

@@ -275,42 +275,50 @@ class Splitter(FileManager):
 
         groups = [[n] for n in nodes]
         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. Ignore protected nodes.
+            # Get the index of the node that would result in the smallest
+            #  merged snippet when merged with the node that follows it.
+            #  Ignore protected nodes.
             mergeable_indices = compress(range(len(adj_sums)), merge_allowed)
-            i0 = int(min(mergeable_indices, key=adj_sums.__getitem__))
-            i1 = i0 + 1
+            C = int(min(mergeable_indices, key=adj_sums.__getitem__))
+
+            # C: Central index
+            # L: Index to the left
+            # R: Index to the right (to be merged in to C)
+            # N: Next index (to the right of R, the "new R")
+            L, R, N = C - 1, C + 1, C + 2
 
             # Recalculate the length. We can't simply use the adj_sum, because
             #  it is an underestimate due to the adjoining suffix/prefix.
-            central_node = groups[i0][-1]
-            merged_text = "".join([text_chunks[i0], central_node.suffix, text_chunks[i1]])
+            central_node = groups[C][-1]
+            merged_text = "".join([text_chunks[C], central_node.suffix, text_chunks[R]])
             merged_text_length = self._count_tokens(merged_text)
 
             # If the true length of the merged pair is too long, don't merge them
             #  Instead, correct the estimate, since shorter pairs may yet exist
             if merged_text_length > self.max_tokens:
-                adj_sums[i0] = merged_text_length
+                adj_sums[C] = merged_text_length
                 continue
 
             # Update adjacent sum estimates
-            if i0 > 0:
-                adj_sums[i0 - 1] += merged_text_length
-            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
+            if L >= 0:
+                adj_sums[L] = lengths[L] + merged_text_length
+            if N < len(adj_sums):
+                adj_sums[R] = lengths[N] + merged_text_length
 
             # The potential merge length for this pair is removed
-            adj_sums.pop(i0)
-            merge_allowed.pop(i0)
+            adj_sums.pop(C)
+
+            # The merged-in node is removed from the protected list
+            #  The merge_allowed list need not be updated - if the node now to
+            #  its right is protected, the merge_allowed element corresponding
+            #  to the merged neighbor will have been True, and now corresponds
+            #  to the merged node.
+            merge_allowed.pop(C)
 
             # Merge the pair of node groups
-            groups[i0 : i1 + 1] = [groups[i0] + groups[i1]]
-            text_chunks[i0 : i1 + 1] = [merged_text]
-            lengths[i0 : i1 + 1] = [merged_text_length]
+            groups[C:N] = [groups[C] + groups[R]]
+            text_chunks[C:N] = [merged_text]
+            lengths[C:N] = [merged_text_length]
 
         return groups
 
@@ -403,13 +411,13 @@ class Splitter(FileManager):
         self._split_into_lines(node)
 
     def _split_into_lines(self, node: CodeBlock):
-        split_text = re.split(r"(\n+)", node.text)
+        split_text = list(re.split(r"(\n+)", node.text))
 
         # 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"):
+        if not re.match(r"^\n+$", split_text[0]):
             split_text = [""] + split_text
-        if split_text[-1].strip("\n"):
+        if not re.match(r"^\n+$", split_text[-1]):
             split_text.append("")
         betweens = split_text[::2]
         lines = split_text[1::2]

{janus_llm-4.2.0 → janus_llm-4.3.1}/janus/language/treesitter/treesitter.py

@@ -154,7 +154,15 @@ class TreeSitterSplitter(Splitter):
             The pointer to the language.
         """
         lib = cdll.LoadLibrary(os.fspath(so_file))
-        language_function = getattr(lib, f"tree_sitter_{self.language}")
+        # Added this try-except block to handle the case where the language is not
+        # supported in lowercase by the creator of the grammar. Ex: COBOL
+        # https://github.com/yutaro-sakamoto/tree-sitter-cobol/blob/main/grammar.js#L13
+        try:
+            language_function = getattr(lib, f"tree_sitter_{self.language}")
+        except AttributeError:
+            language = self.language.upper()
+            language_function = getattr(lib, f"tree_sitter_{language}")
+
         language_function.restype = c_void_p
         pointer = language_function()
         return pointer