janus-llm 2.0.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

janus/converter/diagram.py

@@ -0,0 +1,124 @@
+import json
+from copy import deepcopy
+
+from janus.converter.converter import run_if_changed
+from janus.converter.document import Documenter
+from janus.language.block import TranslatedCodeBlock
+from janus.llm.models_info import MODEL_PROMPT_ENGINES
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class DiagramGenerator(Documenter):
+    """DiagramGenerator
+
+    A class that translates code from one programming language to a set of diagrams.
+    """
+
+    def __init__(
+        self,
+        diagram_type="Activity",
+        add_documentation=False,
+        **kwargs,
+    ) -> None:
+        """Initialize the DiagramGenerator 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.
+            source_language: The source programming language.
+            max_prompts: The maximum number of prompts to try before giving up.
+            db_path: path to chroma database
+            db_config: database configuraiton
+            diagram_type: type of PLANTUML diagram to generate
+        """
+        super().__init__(**kwargs)
+        self._diagram_type = diagram_type
+        self._add_documentation = add_documentation
+        self._documenter = None
+        if add_documentation:
+            self._diagram_prompt_template_name = "diagram_with_documentation"
+        else:
+            self._diagram_prompt_template_name = "diagram"
+        self._load_diagram_prompt_engine()
+
+    def _add_translation(self, block: TranslatedCodeBlock) -> None:
+        """Given an "empty" `TranslatedCodeBlock`, translate the code represented in
+        `block.original`, setting the relevant fields in the translated block. The
+        `TranslatedCodeBlock` is updated in-pace, nothing is returned. Note that this
+        translates *only* the code for this block, not its children.
+
+        Arguments:
+            block: An empty `TranslatedCodeBlock`
+        """
+        if block.translated:
+            return
+
+        if block.original.text is None:
+            block.translated = True
+            return
+
+        if self._add_documentation:
+            documentation_block = deepcopy(block)
+            super()._add_translation(documentation_block)
+            if not documentation_block.translated:
+                message = "Error: unable to produce documentation for code block"
+                log.info(message)
+                raise ValueError(message)
+            documentation = json.loads(documentation_block.text)["docstring"]
+
+        if self._llm is None:
+            message = (
+                "Model not configured correctly, cannot translate. Try setting "
+                "the model"
+            )
+            log.error(message)
+            raise ValueError(message)
+
+        log.debug(f"[{block.name}] Translating...")
+        log.debug(f"[{block.name}] Input text:\n{block.original.text}")
+
+        self._parser.set_reference(block.original)
+
+        query_and_parse = self.diagram_prompt | self._llm | self._parser
+
+        if self._add_documentation:
+            block.text = query_and_parse.invoke(
+                {
+                    "SOURCE_CODE": block.original.text,
+                    "DIAGRAM_TYPE": self._diagram_type,
+                    "DOCUMENTATION": documentation,
+                }
+            )
+        else:
+            block.text = query_and_parse.invoke(
+                {
+                    "SOURCE_CODE": block.original.text,
+                    "DIAGRAM_TYPE": self._diagram_type,
+                }
+            )
+        block.tokens = self._llm.get_num_tokens(block.text)
+        block.translated = True
+
+        log.debug(f"[{block.name}] Output code:\n{block.text}")
+
+    @run_if_changed(
+        "_diagram_prompt_template_name",
+        "_source_language",
+    )
+    def _load_diagram_prompt_engine(self) -> None:
+        """Load the prompt engine according to this instance's attributes.
+
+        If the relevant fields have not been changed since the last time this method was
+        called, nothing happens.
+        """
+        self._diagram_prompt_engine = MODEL_PROMPT_ENGINES[self._model_name](
+            source_language=self._source_language,
+            target_language="text",
+            target_version=None,
+            prompt_template=self._diagram_prompt_template_name,
+        )
+        self.diagram_prompt = self._diagram_prompt_engine.prompt

janus/converter/document.py

@@ -0,0 +1,131 @@
+import json
+import re
+from copy import deepcopy
+
+from janus.converter.converter import Converter
+from janus.language.block import TranslatedCodeBlock
+from janus.language.combine import JsonCombiner
+from janus.parsers.doc_parser import (
+    MadlibsDocumentationParser,
+    MultiDocumentationParser,
+)
+from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class Documenter(Converter):
+    def __init__(
+        self, source_language: str = "fortran", drop_comments: bool = True, **kwargs
+    ):
+        kwargs.update(source_language=source_language)
+        super().__init__(**kwargs)
+        self.set_prompt("document")
+
+        if drop_comments:
+            comment_node_type = LANGUAGES[source_language].get(
+                "comment_node_type", "comment"
+            )
+            self.set_prune_node_types((comment_node_type,))
+
+        self._load_parameters()
+
+
+class MultiDocumenter(Documenter):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.set_prompt("multidocument")
+        self._combiner = JsonCombiner()
+        self._parser = MultiDocumentationParser()
+
+
+class MadLibsDocumenter(Documenter):
+    def __init__(
+        self,
+        comments_per_request: int | None = None,
+        **kwargs,
+    ) -> None:
+        kwargs.update(drop_comments=False)
+        super().__init__(**kwargs)
+        self.set_prompt("document_madlibs")
+        self._combiner = JsonCombiner()
+        self._parser = MadlibsDocumentationParser()
+
+        self.comments_per_request = comments_per_request
+
+    def _add_translation(self, block: TranslatedCodeBlock):
+        if block.translated:
+            return
+
+        if block.original.text is None:
+            block.translated = True
+            return
+
+        if self.comments_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,
+            )
+        )
+
+        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.comments_per_request:
+            return super()._add_translation(block)
+
+        comment_group_indices = list(range(0, len(comments), self.comments_per_request))
+        log.debug(
+            f"[{block.name}] Block contains more than {self.comments_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.comments_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.comments_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)
+            suffix = re.sub(comment_pattern, "", suffix)
+
+            # 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
+            out_text = self._parser.parse(working_block.text)
+            obj.update(json.loads(out_text))
+
+        self._parser.set_reference(block.original)
+        block.text = self._parser.parse(json.dumps(obj))
+        block.tokens = self._llm.get_num_tokens(block.text)
+        block.translated = True

janus/converter/evaluate.py

@@ -0,0 +1,15 @@
+from janus.converter.converter import Converter
+from janus.language.combine import JsonCombiner
+from janus.parsers.eval_parser import EvaluationParser
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class Evaluator(Converter):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.set_prompt("evaluate")
+        self._combiner = JsonCombiner()
+        self._parser = EvaluationParser()
+        self._load_parameters()

janus/converter/requirements.py

@@ -0,0 +1,50 @@
+import json
+from pathlib import Path
+
+from janus.converter.document import Documenter
+from janus.language.block import TranslatedCodeBlock
+from janus.language.combine import ChunkCombiner
+from janus.parsers.reqs_parser import RequirementsParser
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class RequirementsDocumenter(Documenter):
+    """RequirementsGenerator
+
+    A class that translates code from one programming language to its requirements.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.set_prompt("requirements")
+        self._combiner = ChunkCombiner()
+        self._parser = RequirementsParser()
+
+    def _save_to_file(self, block: TranslatedCodeBlock, out_path: Path) -> None:
+        """Save a file to disk.
+
+        Arguments:
+            block: The `CodeBlock` to save to a file.
+        """
+        output_list = list()
+        # For each chunk of code, get generation metadata, the text of the code,
+        # and the LLM generated requirements
+        for child in block.children:
+            code = child.original.text
+            requirements = self._parser.parse_combined_output(child.complete_text)
+            metadata = dict(
+                retries=child.total_retries,
+                cost=child.total_cost,
+                processing_time=child.processing_time,
+            )
+            # Put them all in a top level 'output' key
+            output_list.append(
+                dict(metadata=metadata, code=code, requirements=requirements)
+            )
+        obj = dict(
+            output=output_list,
+        )
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_text(json.dumps(obj, indent=2), encoding="utf-8")

janus/converter/translate.py

@@ -0,0 +1,108 @@
+from janus.converter.converter import Converter, run_if_changed
+from janus.llm.models_info import MODEL_PROMPT_ENGINES
+from janus.parsers.code_parser import CodeParser
+from janus.prompts.prompt import SAME_OUTPUT
+from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class Translator(Converter):
+    """A class that translates code from one programming language to another."""
+
+    def __init__(
+        self,
+        target_language: str = "python",
+        target_version: str | None = "3.10",
+        **kwargs,
+    ) -> None:
+        """Initialize a Translator instance.
+
+        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.
+            source_language: The source programming language.
+            target_language: The target programming language.
+            target_version: The target version of the target programming language.
+            max_prompts: The maximum number of prompts to try before giving up.
+            max_tokens: The maximum number of tokens the model will take in.
+                If unspecificed, model's default max will be used.
+            prompt_template: name of prompt template directory
+                (see janus/prompts/templates) or path to a directory.
+        """
+        super().__init__(**kwargs)
+
+        self._target_version: str | None
+
+        self.set_target_language(
+            target_language=target_language,
+            target_version=target_version,
+        )
+
+        self._load_parameters()
+
+    def _load_parameters(self) -> None:
+        self._load_parser()
+        super()._load_parameters()
+
+    def set_target_language(
+        self, target_language: str, target_version: str | None
+    ) -> None:
+        """Validate and set the target language.
+
+        The affected objects will not be updated until translate() is called.
+
+        Arguments:
+            target_language: The target programming language.
+            target_version: The target version of the target programming language.
+        """
+        target_language = target_language.lower()
+        if target_language not in LANGUAGES:
+            raise ValueError(
+                f"Invalid target language: {target_language}. "
+                "Valid target languages are found in `janus.utils.enums.LANGUAGES`."
+            )
+        self._target_language = target_language
+        self._target_version = target_version
+        self._target_suffix = f".{LANGUAGES[target_language]['suffix']}"
+
+    @run_if_changed(
+        "_prompt_template_name",
+        "_source_language",
+        "_target_language",
+        "_target_version",
+        "_model_name",
+    )
+    def _load_prompt(self) -> None:
+        """Load the prompt according to this instance's attributes.
+
+        If the relevant fields have not been changed since the last time this
+        method was called, nothing happens.
+        """
+        if self._prompt_template_name in SAME_OUTPUT:
+            if self._target_language != self._source_language:
+                raise ValueError(
+                    f"Prompt template ({self._prompt_template_name}) suggests "
+                    f"source and target languages should match, but do not "
+                    f"({self._source_language} != {self._target_language})"
+                )
+
+        prompt_engine = MODEL_PROMPT_ENGINES[self._model_name](
+            source_language=self._source_language,
+            target_language=self._target_language,
+            target_version=self._target_version,
+            prompt_template=self._prompt_template_name,
+        )
+        self._prompt = prompt_engine.prompt
+
+    @run_if_changed("_target_language")
+    def _load_parser(self) -> None:
+        """Load the parser according to this instance's attributes.
+
+        If the relevant fields have not been changed since the last time this
+        method was called, nothing happens.
+        """
+        self._parser = CodeParser(language=self._target_language)

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/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/language/alc/__init__.py

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

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/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/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],
             ]
         )
@@ -214,7 +216,7 @@ class TranslatedCodeBlock(CodeBlock):
         self.translated = False
         self.cost = 0.0
         self.retries = 0
-        self.processing_time = 0
+        self.processing_time = 0.0
 
     @property
     def total_cost(self) -> float:

janus/language/combine.py

@@ -44,7 +44,6 @@ class Combiner(FileManager):
         if block.text is None:
             children = sorted(block.children)
             block.text = "".join([c.complete_text for c in children])
-            block.children = []
             block.complete = children_complete
             return
 

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/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/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/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/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()