janus-llm 3.3.2__py3-none-any.whl → 3.4.0__py3-none-any.whl

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__ = "3.3.2"
+__version__ = "3.4.0"
 
 # Ignoring a deprecation warning from langchain_core that I can't seem to hunt down
 warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

janus/cli.py

@@ -200,6 +200,14 @@ def translate(
             help="Whether to overwrite existing files in the output directory",
         ),
     ] = False,
+    skip_context: Annotated[
+        bool,
+        typer.Option(
+            "--skip-context",
+            help="Prompts will include any context information associated with source"
+            " code blocks, unless this option is specified",
+        ),
+    ] = False,
     temp: Annotated[
         float,
         typer.Option("--temperature", "-T", help="Sampling temperature.", min=0, max=2),
@@ -265,6 +273,7 @@ def translate(
         db_path=db_loc,
         db_config=collections_config,
         splitter_type=splitter_type,
+        skip_context=skip_context,
     )
     translator.translate(input_dir, output_dir, overwrite, collection)
 
@@ -322,6 +331,14 @@ def document(
             help="Whether to overwrite existing files in the output directory",
         ),
     ] = False,
+    skip_context: Annotated[
+        bool,
+        typer.Option(
+            "--skip-context",
+            help="Prompts will include any context information associated with source"
+            " code blocks, unless this option is specified",
+        ),
+    ] = False,
     doc_mode: Annotated[
         str,
         typer.Option(
@@ -390,6 +407,7 @@ def document(
         db_path=db_loc,
         db_config=collections_config,
         splitter_type=splitter_type,
+        skip_context=skip_context,
     )
     if doc_mode == "madlibs":
         documenter = MadLibsDocumenter(
@@ -458,6 +476,14 @@ def diagram(
             help="Whether to overwrite existing files in the output directory",
         ),
     ] = False,
+    skip_context: Annotated[
+        bool,
+        typer.Option(
+            "--skip-context",
+            help="Prompts will include any context information associated with source"
+            " code blocks, unless this option is specified",
+        ),
+    ] = False,
     temperature: Annotated[
         float,
         typer.Option("--temperature", "-t", help="Sampling temperature.", min=0, max=2),
@@ -507,6 +533,7 @@ def diagram(
         diagram_type=diagram_type,
         add_documentation=add_documentation,
         splitter_type=splitter_type,
+        skip_context=skip_context,
     )
     diagram_generator.translate(input_dir, output_dir, overwrite, collection)
 

janus/converter/converter.py

@@ -3,13 +3,13 @@ import json
 import math
 import time
 from pathlib import Path
-from typing import Any
+from typing import Any, List, Optional, Tuple
 
 from langchain.output_parsers import RetryWithErrorOutputParser
 from langchain_core.exceptions import OutputParserException
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.output_parsers import BaseOutputParser
-from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate
 from langchain_core.runnables import RunnableLambda, RunnableParallel
 from openai import BadRequestError, RateLimitError
 from pydantic import ValidationError
@@ -77,6 +77,7 @@ class Converter:
         prune_node_types: tuple[str, ...] = (),
         splitter_type: str = "file",
         refiner_type: str = "basic",
+        skip_context: bool = False,
     ) -> None:
         """Initialize a Converter instance.
 
@@ -142,6 +143,8 @@ class Converter:
         self.set_db_path(db_path=db_path)
         self.set_db_config(db_config=db_config)
 
+        self.skip_context = skip_context
+
         # Child class must call this. Should we enforce somehow?
         # self._load_parameters()
 
@@ -602,6 +605,9 @@ class Converter:
 
         # Retries with just the input
         n3 = math.ceil(self.max_prompts / (n1 * n2))
+        # Make replacements in the prompt
+        if not self.skip_context:
+            self._make_prompt_additions(block)
 
         refine_output = RefinerParser(
             parser=self._parser,
@@ -648,6 +654,35 @@ class Converter:
             output=output,
         )
 
+    @staticmethod
+    def _get_prompt_additions(block) -> Optional[List[Tuple[str, str]]]:
+        """Get a list of strings to append to the prompt.
+
+        Arguments:
+            block: The `TranslatedCodeBlock` to save to a file.
+        """
+        return [(key, item) for key, item in block.context_tags.items()]
+
+    def _make_prompt_additions(self, block: CodeBlock):
+        # Prepare the additional context to prepend
+        additional_context = "".join(
+            [
+                f"{context_tag}: {context}\n"
+                for context_tag, context in self._get_prompt_additions(block)
+            ]
+        )
+
+        # Iterate through existing messages to find and update the system message
+        for i, message in enumerate(self._prompt.messages):
+            if isinstance(message, SystemMessagePromptTemplate):
+                # Prepend the additional context to the system message
+                updated_system_message = SystemMessagePromptTemplate.from_template(
+                    additional_context + message.prompt.template
+                )
+                # Directly modify the message in the list
+                self._prompt.messages[i] = updated_system_message
+                break  # Assuming there's only one system message to update
+
     def _save_to_file(self, block: TranslatedCodeBlock, out_path: Path) -> None:
         """Save a file to disk.
 

janus/converter/requirements.py

@@ -22,6 +22,11 @@ class RequirementsDocumenter(Documenter):
         self._combiner = ChunkCombiner()
         self._parser = RequirementsParser()
 
+    @staticmethod
+    def get_prompt_replacements(block) -> dict[str, str]:
+        prompt_replacements: dict[str, str] = {"SOURCE_CODE": block.original.text}
+        return prompt_replacements
+
     def _save_to_file(self, block: TranslatedCodeBlock, out_path: Path) -> None:
         """Save a file to disk.
 

janus/language/alc/alc.py

@@ -1,3 +1,6 @@
+import re
+from typing import Optional
+
 from langchain.schema.language_model import BaseLanguageModel
 
 from janus.language.block import CodeBlock
@@ -61,7 +64,11 @@ class AlcSplitter(TreeSitterSplitter):
             #  next csect or dsect instruction
             sects: list[list[CodeBlock]] = [[]]
             for c in block.children:
-                if c.node_type in sect_types:
+                if c.node_type == "csect_instruction":
+                    c.context_tags["alc_section"] = "CSECT"
+                    sects.append([c])
+                elif c.node_type == "dsect_instruction":
+                    c.context_tags["alc_section"] = "DSECT"
                     sects.append([c])
                 else:
                     sects[-1].append(c)
@@ -85,3 +92,94 @@ class AlcSplitter(TreeSitterSplitter):
             queue.extend(block.children)
 
         return root
+
+
+class AlcListingSplitter(AlcSplitter):
+    """A class for splitting ALC listing 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
+        """
+        # The string to mark the end of the listing header
+        self.header_indicator_str: str = (
+            "Loc  Object Code    Addr1 Addr2  Stmt   Source Statement"
+        )
+        # How many characters to trim from the right side to remove the address column
+        self.address_column_chars: int = 10
+        # The string to mark the end of the left margin
+        self.left_margin_indicator_str: str = "Stmt"
+        super().__init__(
+            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:
+        active_usings = self.get_active_usings(code)
+        code = self.preproccess_assembly(code)
+        ast: CodeBlock = super()._get_ast(code)
+        ast.context_tags["active_usings"] = active_usings
+        return ast
+
+    def preproccess_assembly(self, code: str) -> str:
+        """Remove non-essential lines from an assembly snippet"""
+
+        lines = code.splitlines()
+        lines = self.strip_header_and_left(lines)
+        lines = self.strip_addresses(lines)
+        return "".join(str(line) for line in lines)
+
+    def get_active_usings(self, code: str) -> Optional[str]:
+        """Look for 'active usings' in the ALC listing header"""
+        lines = code.splitlines()
+        for line in lines:
+            if "Active Usings:" in line:
+                return line.split("Active Usings:")[1]
+        return None
+
+    def strip_header_and_left(
+        self,
+        lines: list[str],
+    ) -> list[str]:
+        """Remove the header and the left panel from the assembly sample"""
+
+        esd_regex = re.compile(f".*{self.header_indicator_str}.*")
+
+        header_end_index: int = [
+            i for i, item in enumerate(lines) if re.search(esd_regex, item)
+        ][0]
+
+        left_content_end_column = lines[header_end_index].find(
+            self.left_margin_indicator_str
+        )
+        hori_output_lines = lines[(header_end_index + 1) :]
+
+        left_output_lines = [
+            line[left_content_end_column + 5 :] for line in hori_output_lines
+        ]
+        return left_output_lines
+
+    def strip_addresses(self, lines: list[str]) -> list[str]:
+        """Strip the addresses which run down the right side of the assembly snippet"""
+
+        stripped_lines = [line[: -self.address_column_chars] for line in lines]
+        return stripped_lines
+
+    def strip_footer(self, lines: list[str]):
+        """Strip the footer from the assembly snippet"""
+        return NotImplementedError

janus/language/block.py

@@ -45,6 +45,7 @@ class CodeBlock:
         children: list[ForwardRef("CodeBlock")],
         embedding_id: Optional[str] = None,
         affixes: Tuple[str, str] = ("", ""),
+        context_tags: dict[str, str] = {},
     ) -> None:
         self.id: Hashable = id
         self.name: Optional[str] = name
@@ -59,6 +60,7 @@ class CodeBlock:
         self.children: list[ForwardRef("CodeBlock")] = sorted(children)
         self.embedding_id: Optional[str] = embedding_id
         self.affixes: Tuple[str, str] = affixes
+        self.context_tags: dict[str, str] = context_tags
 
         self.complete = True
         self.omit_prefix = True

janus/language/naive/simple_ast.py

@@ -1,12 +1,24 @@
-from janus.language.alc.alc import AlcSplitter
+from janus.language.alc.alc import AlcListingSplitter, AlcSplitter
 from janus.language.mumps.mumps import MumpsSplitter
 from janus.language.naive.registry import register_splitter
+from janus.language.splitter import Splitter
 from janus.language.treesitter import TreeSitterSplitter
 from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
 
 
 @register_splitter("ast-flex")
-def get_flexible_ast(language: str, **kwargs):
+def get_flexible_ast(language: str, **kwargs) -> Splitter:
+    """Get a flexible AST splitter for the given language.
+
+    Arguments:
+        language: The language to get the splitter for.
+
+    Returns:
+        A flexible AST splitter for the given language.
+    """
     if language == "ibmhlasm":
         return AlcSplitter(**kwargs)
     elif language == "mumps":
@@ -16,7 +28,17 @@ def get_flexible_ast(language: str, **kwargs):
 
 
 @register_splitter("ast-strict")
-def get_strict_ast(language: str, **kwargs):
+def get_strict_ast(language: str, **kwargs) -> Splitter:
+    """Get a strict AST splitter for the given language.
+
+    The strict splitter will only return nodes that are of a functional type.
+
+    Arguments:
+        language: The language to get the splitter for.
+
+    Returns:
+        A strict AST splitter for the given language.
+    """
     kwargs.update(
         protected_node_types=LANGUAGES[language]["functional_node_types"],
         prune_unprotected=True,
@@ -27,3 +49,45 @@ def get_strict_ast(language: str, **kwargs):
         return MumpsSplitter(**kwargs)
     else:
         return TreeSitterSplitter(language=language, **kwargs)
+
+
+@register_splitter("ast-strict-listing")
+def get_strict_listing_ast(language: str, **kwargs) -> Splitter:
+    """Get a strict AST splitter for the given language. This splitter is intended for
+    use with IBM HLASM.
+
+    The strict splitter will only return nodes that are of a functional type.
+
+    Arguments:
+        language: The language to get the splitter for.
+
+    Returns:
+        A strict AST splitter for the given language.
+    """
+    kwargs.update(
+        protected_node_types=LANGUAGES[language]["functional_node_types"],
+        prune_unprotected=True,
+    )
+    if language == "ibmhlasm":
+        return AlcListingSplitter(**kwargs)
+    else:
+        log.warning("Listing splitter is only intended for use with IBMHLASM!")
+        return TreeSitterSplitter(language=language, **kwargs)
+
+
+@register_splitter("ast-flex-listing")
+def get_flexible_listing_ast(language: str, **kwargs) -> Splitter:
+    """Get a flexible AST splitter for the given language. This splitter is intended for
+    use with IBM HLASM.
+
+    Arguments:
+        language: The language to get the splitter for.
+
+    Returns:
+        A flexible AST splitter for the given language.
+    """
+    if language == "ibmhlasm":
+        return AlcListingSplitter(**kwargs)
+    else:
+        log.warning("Listing splitter is only intended for use with IBMHLASM!")
+        return TreeSitterSplitter(language=language, **kwargs)

{janus_llm-3.3.2.dist-info → janus_llm-3.4.0.dist-info}/METADATA

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

{janus_llm-3.3.2.dist-info → janus_llm-3.4.0.dist-info}/RECORD

@@ -1,17 +1,17 @@
-janus/__init__.py,sha256=MyAIBSlNInMKo5GQ3bDD7TxPwy1U-UHo7831cggUHlE,361
+janus/__init__.py,sha256=DDs8SRs9v96slIS8XcqcLZrUAfuwXiT9vFa6OlwYHEY,361
 janus/__main__.py,sha256=lEkpNtLVPtFo8ySDZeXJ_NXDHb0GVdZFPWB4gD4RPS8,64
 janus/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/_tests/conftest.py,sha256=V7uW-oq3YbFiRPvrq15YoVVrA1n_83pjgiyTZ-IUGW8,963
 janus/_tests/test_cli.py,sha256=oYJsUGWfpBJWEGRG5NGxdJedU5DU_m6fwJ7xEbJVYl0,4244
-janus/cli.py,sha256=S6IaQyUG55xSB166xARn6TQc_cOCIQ_ZMdIi7vj6BMk,31626
+janus/cli.py,sha256=Uvg6xPnJK7dJcDw1J58s3I7qwiNgPfj2qyXwpC6UGbI,32538
 janus/converter/__init__.py,sha256=U2EOMcCykiC0ZqhorNefOP_04hOF18qhYoPKrVp1Vrk,345
 janus/converter/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/converter/_tests/test_translate.py,sha256=yzcFEGc_z8QmBBBmC9dZnfL9tT8CD1rkpc8Hz44Jp4c,5631
-janus/converter/converter.py,sha256=m0eGrybrPWw2B0nXrd1Fx6o0BO1JvY4PfC0rFUtpLxY,25633
+janus/converter/converter.py,sha256=tBUX9rsu7filfqm6B_knafQUfNb4x4ls8so7KbEq3gs,27160
 janus/converter/diagram.py,sha256=5mo1H3Y1uIBPYdIsWz9kxluN5DNyuUMZrtcJmGF2Uw0,5335
 janus/converter/document.py,sha256=hsW512veNjFWbdl5WriuUdNmMEqZy8ktRvqn9rRmA6E,4566
 janus/converter/evaluate.py,sha256=APWQUY3gjAXqkJkPzvj0UA4wPK3Cv9QSJLM-YK9t-ng,476
-janus/converter/requirements.py,sha256=6YvrJRVH9BuPCOPxnXmaJQFYmoLYYvCu3zTntDLHeNg,1832
+janus/converter/requirements.py,sha256=9tvQ40FZJtG8niIFn45gPQCgKKHVPPoFLinBv6RAqO4,2027
 janus/converter/translate.py,sha256=0brQTlSfBYmXtoM8QYIOiyr0LrTr0S1n68Du-BR7_WQ,4236
 janus/embedding/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/embedding/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -29,13 +29,13 @@ janus/language/_tests/test_splitter.py,sha256=Hqexa39LLEXlK3ZUw7Zot4PUIACvye2vkq
 janus/language/alc/__init__.py,sha256=j7vOMGhT1Vri6p8dsjSaY-fkO5uFn0sJ0nrNGGvcizM,42
 janus/language/alc/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/language/alc/_tests/test_alc.py,sha256=NgVeOctm9zf-S328DdUNn9et_-lK1t5O0O2FKElb91Q,1027
-janus/language/alc/alc.py,sha256=CIyY3LneYToItadUGg54AeMcZL-dwx8buXB71e6EAag,3116
+janus/language/alc/alc.py,sha256=l1p6zwyE7ZzY9rnjsUZGuQW41hijFrzxnXnFOpfGq8k,6590
 janus/language/binary/__init__.py,sha256=AlNAe12ZA366kcGSrQ1FJyOdbwxFqGBFkYR2K6yL818,51
 janus/language/binary/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/language/binary/_tests/test_binary.py,sha256=SDdI6tsQj9yXle7wBsksHuKULLMHv7mNgUkDx1nCvpw,1733
 janus/language/binary/binary.py,sha256=jcc-LZx8Ss-g4j0a691U2yVJV6JGV_zpE_6y_aDq4Cw,6579
 janus/language/binary/reveng/decompile_script.py,sha256=veW51oJzuO-4UD3Er062jXZ_FYtTFo9OCkl82Z2xr6A,2182
-janus/language/block.py,sha256=viMU3MBltiHsJXLAiIfjnlXCR-Iw2Sb-ENaK_7hUipw,9298
+janus/language/block.py,sha256=2rjAYUosHFfWRgLnzf50uAgTMST4Md9Kx6JrlUfEfX4,9398
 janus/language/combine.py,sha256=Wtve06fa-_Wjv_V5RIf1Nfmg0UxcOEtFNj4vVHpSNbo,2940
 janus/language/file.py,sha256=jy-cReAoI6F97TXR5bbhPyt8XyUZCdFYnVboubDA_y4,571
 janus/language/mumps/__init__.py,sha256=-Ou_wJ-JgHezfp1dub2_qCYNiK9wO-zo2MlqxM9qiwE,48
@@ -47,7 +47,7 @@ janus/language/naive/__init__.py,sha256=_Gq4inONyVYxe8WLB59d_69kqGbtF40BGKoJPnK4
 janus/language/naive/basic_splitter.py,sha256=RM9pJK2YkHfb6_EFEV-dh_rLqkjS6v0cn3ASPf8A6Fg,459
 janus/language/naive/chunk_splitter.py,sha256=ebRSbaJhDW-Hyr5__ukbdmAl6kQ1WWFqrq_SfCgHo6k,772
 janus/language/naive/registry.py,sha256=8YQX1q0IdAm7t69-oC_00I-vfkdRnHuX-OD3KEjEIuU,294
-janus/language/naive/simple_ast.py,sha256=_--5mshYFt7BaZc_gAG4sSj-CTt9C1_fedP9fi2ffz4,969
+janus/language/naive/simple_ast.py,sha256=T53UwAZyRfePXzpiNUhe4FyDev6YcX1dVDxYkTcDRPE,3032
 janus/language/naive/tag_splitter.py,sha256=IXWMn9tBVUGAtzvQi89GhoZ6g7fPXk5MzO0kMCr2mb0,2045
 janus/language/node.py,sha256=baoYFtapwBQqBtUN6EvHFYRkbR-EcEw1b3fQvH9zIAM,204
 janus/language/splitter.py,sha256=hITFp4a9bJ6sP74AWvC2GWa2Poo10MyHXYTj9hviXss,16970
@@ -100,8 +100,8 @@ janus/utils/_tests/test_progress.py,sha256=Rs_u5PiGjP-L-o6C1fhwfE1ig8jYu9Xo9s4p8
 janus/utils/enums.py,sha256=AoilbdiYyMvY2Mp0AM4xlbLSELfut2XMwhIM1S_msP4,27610
 janus/utils/logger.py,sha256=KZeuaMAnlSZCsj4yL0P6N-JzZwpxXygzACWfdZFeuek,2337
 janus/utils/progress.py,sha256=PIpcQec7SrhsfqB25LHj2CDDkfm9umZx90d9LZnAx6k,1469
-janus_llm-3.3.2.dist-info/LICENSE,sha256=_j0st0a-HB6MRbP3_BW3PUqpS16v54luyy-1zVyl8NU,10789
-janus_llm-3.3.2.dist-info/METADATA,sha256=NfyvyOPvC1Qj3jNh264MUYWnsA67hOD6jzL7uiIkDq4,4184
-janus_llm-3.3.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-janus_llm-3.3.2.dist-info/entry_points.txt,sha256=OGhQwzj6pvXp79B0SaBD5apGekCu7Dwe9fZZT_TZ544,39
-janus_llm-3.3.2.dist-info/RECORD,,
+janus_llm-3.4.0.dist-info/LICENSE,sha256=_j0st0a-HB6MRbP3_BW3PUqpS16v54luyy-1zVyl8NU,10789
+janus_llm-3.4.0.dist-info/METADATA,sha256=GmgJ5Oq3MkXOpYw9Vl8YFaLx9dUNcaFZs8GQDV1t8vc,4184
+janus_llm-3.4.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+janus_llm-3.4.0.dist-info/entry_points.txt,sha256=OGhQwzj6pvXp79B0SaBD5apGekCu7Dwe9fZZT_TZ544,39
+janus_llm-3.4.0.dist-info/RECORD,,