janus-llm 3.5.2__py3-none-any.whl → 4.0.0__py3-none-any.whl

janus/converter/diagram.py

@@ -1,14 +1,6 @@
-import math
+from langchain_core.runnables import Runnable, RunnableParallel
 
-from langchain.output_parsers import RetryWithErrorOutputParser
-from langchain_core.exceptions import OutputParserException
-from langchain_core.runnables import RunnableLambda, RunnableParallel
-
-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.parsers.refiner_parser import RefinerParser
 from janus.parsers.uml import UMLSyntaxParser
 from janus.utils.logger import create_logger
 
@@ -16,10 +8,7 @@ log = create_logger(__name__)
 
 
 class DiagramGenerator(Documenter):
-    """DiagramGenerator
-
-    A class that translates code from one programming language to a set of diagrams.
-    """
+    """A Converter that translates code into a set of PLANTUML diagrams."""
 
     def __init__(
         self,
@@ -30,110 +19,33 @@ class DiagramGenerator(Documenter):
         """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
+            add_documentation: Whether to add a documentation step prior to
+                diagram generation.
         """
-        super().__init__(**kwargs)
         self._diagram_type = diagram_type
         self._add_documentation = add_documentation
-        self._documenter = None
-        self._diagram_parser = UMLSyntaxParser(language="plantuml")
-        if add_documentation:
-            self._diagram_prompt_template_name = "diagram_with_documentation"
-        else:
-            self._diagram_prompt_template_name = "diagram"
-        self._load_diagram_prompt_engine()
+        self._documenter = Documenter(**kwargs)
 
-    def _run_chain(self, block: TranslatedCodeBlock) -> str:
-        input = self._parser.parse_input(block.original)
-        n1 = round(self.max_prompts ** (1 / 3))
+        super().__init__(**kwargs)
+
+        self.set_prompt("diagram_with_documentation" if add_documentation else "diagram")
+        self._parser = UMLSyntaxParser(language="plantuml")
 
-        # Retries with the input, output, and error
-        n2 = round((self.max_prompts // n1) ** (1 / 2))
+        self._load_parameters()
 
-        # Retries with just the input
-        n3 = math.ceil(self.max_prompts / (n1 * n2))
+    def _load_prompt(self):
+        super()._load_prompt()
+        self._prompt = self._prompt.partial(DIAGRAM_TYPE=self._diagram_type)
 
+    def _input_runnable(self) -> Runnable:
         if self._add_documentation:
-            documentation_text = super()._run_chain(block)
-            refine_output = RefinerParser(
-                parser=self._diagram_parser,
-                initial_prompt=self._diagram_prompt.format(
-                    **{
-                        "SOURCE_CODE": input,
-                        "DOCUMENTATION": documentation_text,
-                        "DIAGRAM_TYPE": self._diagram_type,
-                    }
-                ),
-                refiner=self._refiner,
-                max_retries=n1,
-                llm=self._llm,
-            )
-        else:
-            refine_output = RefinerParser(
-                parser=self._diagram_parser,
-                initial_prompt=self._diagram_prompt.format(
-                    **{
-                        "SOURCE_CODE": input,
-                        "DIAGRAM_TYPE": self._diagram_type,
-                    }
-                ),
-                refiner=self._refiner,
-                max_retries=n1,
-                llm=self._llm,
+            return RunnableParallel(
+                SOURCE_CODE=self._parser.parse_input,
+                DOCUMENTATION=self._documenter.chain,
+                context=self._retriever,
             )
-        retry = RetryWithErrorOutputParser.from_llm(
-            llm=self._llm,
-            parser=refine_output,
-            max_retries=n2,
-        )
-        completion_chain = self._prompt | self._llm
-        chain = RunnableParallel(
-            completion=completion_chain, prompt_value=self._diagram_prompt
-        ) | RunnableLambda(lambda x: retry.parse_with_prompt(**x))
-        for _ in range(n3):
-            try:
-                if self._add_documentation:
-                    return chain.invoke(
-                        {
-                            "SOURCE_CODE": input,
-                            "DOCUMENTATION": documentation_text,
-                            "DIAGRAM_TYPE": self._diagram_type,
-                        }
-                    )
-                else:
-                    return chain.invoke(
-                        {
-                            "SOURCE_CODE": input,
-                            "DIAGRAM_TYPE": self._diagram_type,
-                        }
-                    )
-            except OutputParserException:
-                pass
-
-        raise OutputParserException(f"Failed to parse after {n1*n2*n3} retries")
-
-    @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_id](
-            source_language=self._source_language,
-            target_language="text",
-            target_version=None,
-            prompt_template=self._diagram_prompt_template_name,
+        return RunnableParallel(
+            SOURCE_CODE=self._parser.parse_input,
+            context=self._retriever,
         )
-        self._diagram_prompt = self._diagram_prompt_engine.prompt

janus/converter/translate.py

@@ -90,7 +90,7 @@ class Translator(Converter):
                     f"({self._source_language} != {self._target_language})"
                 )
 
-        prompt_engine = MODEL_PROMPT_ENGINES[self._model_id](
+        prompt_engine = MODEL_PROMPT_ENGINES[self._llm.short_model_id](
             source_language=self._source_language,
             target_language=self._target_language,
             target_version=self._target_version,

janus/language/alc/_tests/test_alc.py

@@ -12,7 +12,7 @@ class TestAlcSplitter(unittest.TestCase):
     def setUp(self):
         """Set up the tests."""
         model_name = "gpt-4o"
-        llm, _, _, _ = load_model(model_name)
+        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")

janus/language/alc/alc.py

@@ -1,12 +1,11 @@
 import re
 from typing import Optional
 
-from langchain.schema.language_model import BaseLanguageModel
-
 from janus.language.block import CodeBlock
 from janus.language.combine import Combiner
 from janus.language.node import NodeType
 from janus.language.treesitter import TreeSitterSplitter
+from janus.llm.models_info import JanusModel
 from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
@@ -27,7 +26,7 @@ class AlcSplitter(TreeSitterSplitter):
 
     def __init__(
         self,
-        model: None | BaseLanguageModel = None,
+        model: JanusModel | None = None,
         max_tokens: int = 4096,
         protected_node_types: tuple[str, ...] = (),
         prune_node_types: tuple[str, ...] = (),
@@ -63,7 +62,7 @@ class AlcSplitter(TreeSitterSplitter):
             #  instruction and containing all the subsequent nodes up until the
             #  next csect or dsect instruction
             sects: list[list[CodeBlock]] = [[]]
-            for c in block.children:
+            for c in sorted(block.children):
                 if c.node_type == "csect_instruction":
                     c.context_tags["alc_section"] = "CSECT"
                     sects.append([c])
@@ -101,7 +100,7 @@ class AlcListingSplitter(AlcSplitter):
 
     def __init__(
         self,
-        model: None | BaseLanguageModel = None,
+        model: JanusModel | None = None,
         max_tokens: int = 4096,
         protected_node_types: tuple[str, ...] = (),
         prune_node_types: tuple[str, ...] = (),
@@ -129,12 +128,18 @@ class AlcListingSplitter(AlcSplitter):
             prune_unprotected=prune_unprotected,
         )
 
-    def _get_ast(self, code: str) -> CodeBlock:
+    def split_string(self, code: str, name: str) -> CodeBlock:
+        # Override split_string to use processed code and track active usings
         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
+        processed_code = self.preproccess_assembly(code)
+        root = super().split_string(processed_code, name)
+        if active_usings is not None:
+            stack = [root]
+            while stack:
+                block = stack.pop()
+                block.context_tags["active_usings"] = active_usings
+                stack.extend(block.children)
+        return root
 
     def preproccess_assembly(self, code: str) -> str:
         """Remove non-essential lines from an assembly snippet"""
@@ -142,7 +147,7 @@ class AlcListingSplitter(AlcSplitter):
         lines = code.splitlines()
         lines = self.strip_header_and_left(lines)
         lines = self.strip_addresses(lines)
-        return "".join(str(line) for line in lines)
+        return "\n".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"""

janus/language/binary/_tests/test_binary.py

@@ -15,7 +15,7 @@ class TestBinarySplitter(unittest.TestCase):
     def setUp(self):
         model_name = "gpt-4o"
         self.binary_file = Path("janus/language/binary/_tests/hello")
-        self.llm, _, _, _ = load_model(model_name)
+        self.llm = load_model(model_name)
         self.splitter = BinarySplitter(model=self.llm)
         os.environ["GHIDRA_INSTALL_PATH"] = "~/programs/ghidra_10.4_PUBLIC"
 

janus/language/binary/binary.py

@@ -5,11 +5,11 @@ import tempfile
 from pathlib import Path
 
 import tree_sitter
-from langchain.schema.language_model import BaseLanguageModel
 
 from janus.language.block import CodeBlock
 from janus.language.combine import Combiner
 from janus.language.treesitter import TreeSitterSplitter
+from janus.llm.models_info import JanusModel
 from janus.utils.enums import LANGUAGES
 from janus.utils.logger import create_logger
 
@@ -31,7 +31,7 @@ class BinarySplitter(TreeSitterSplitter):
 
     def __init__(
         self,
-        model: None | BaseLanguageModel = None,
+        model: JanusModel | None = None,
         max_tokens: int = 4096,
         protected_node_types: tuple[str] = (),
         prune_node_types: tuple[str] = (),

janus/language/mumps/_tests/test_mumps.py

@@ -12,7 +12,7 @@ class TestMumpsSplitter(unittest.TestCase):
     def setUp(self):
         """Set up the tests."""
         model_name = "gpt-4o"
-        llm, _, _, _ = load_model(model_name)
+        llm = load_model(model_name)
         self.splitter = MumpsSplitter(model=llm)
         self.combiner = Combiner(language="mumps")
         self.test_file = Path("janus/language/mumps/_tests/mumps.m")

janus/language/mumps/mumps.py

@@ -1,11 +1,10 @@
 import re
 
-from langchain.schema.language_model import BaseLanguageModel
-
 from janus.language.block import CodeBlock
 from janus.language.combine import Combiner
 from janus.language.node import NodeType
 from janus.language.splitter import Splitter
+from janus.llm.models_info import JanusModel
 from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
@@ -44,7 +43,7 @@ class MumpsSplitter(Splitter):
 
     def __init__(
         self,
-        model: None | BaseLanguageModel = None,
+        model: JanusModel | None = None,
         max_tokens: int = 4096,
         protected_node_types: tuple[str] = ("routine_definition",),
         prune_node_types: tuple[str] = (),

janus/language/naive/simple_ast.py

@@ -19,6 +19,7 @@ def get_flexible_ast(language: str, **kwargs) -> Splitter:
     Returns:
         A flexible AST splitter for the given language.
     """
+    kwargs.update(protected_node_types=())
     if language == "ibmhlasm":
         return AlcSplitter(**kwargs)
     elif language == "mumps":
@@ -28,7 +29,7 @@ def get_flexible_ast(language: str, **kwargs) -> Splitter:
 
 
 @register_splitter("ast-strict")
-def get_strict_ast(language: str, **kwargs) -> Splitter:
+def get_strict_ast(language: str, prune_unprotected=True, **kwargs) -> Splitter:
     """Get a strict AST splitter for the given language.
 
     The strict splitter will only return nodes that are of a functional type.
@@ -41,7 +42,7 @@ def get_strict_ast(language: str, **kwargs) -> Splitter:
     """
     kwargs.update(
         protected_node_types=LANGUAGES[language]["functional_node_types"],
-        prune_unprotected=True,
+        prune_unprotected=prune_unprotected,
     )
     if language == "ibmhlasm":
         return AlcSplitter(**kwargs)

janus/language/splitter.py

@@ -4,11 +4,11 @@ from pathlib import Path
 from typing import List
 
 import tiktoken
-from langchain.schema.language_model import BaseLanguageModel
 
 from janus.language.block import CodeBlock
 from janus.language.file import FileManager
 from janus.language.node import NodeType
+from janus.llm.models_info import JanusModel
 from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
@@ -44,7 +44,7 @@ class Splitter(FileManager):
     def __init__(
         self,
         language: str,
-        model: None | BaseLanguageModel = None,
+        model: JanusModel | None = None,
         max_tokens: int = 4096,
         skip_merge: bool = False,
         protected_node_types: tuple[str, ...] = (),
@@ -387,7 +387,10 @@ class Splitter(FileManager):
             return
 
         if self._is_protected(node):
-            raise TokenLimitError(r"Irreducible node too large for context!")
+            log.error(
+                "Protected node too large for context!"
+                f" ({node.tokens} > {self.max_tokens})"
+            )
 
         if node.children:
             for child in node.children:
@@ -423,7 +426,7 @@ class Splitter(FileManager):
             name = f"{node.name}-L#{node_line}"
             tokens = self._count_tokens(line)
             if tokens > self.max_tokens:
-                raise TokenLimitError(
+                log.error(
                     "Irreducible node too large for context!"
                     f" ({tokens} > {self.max_tokens})"
                 )

janus/language/treesitter/_tests/test_treesitter.py

@@ -13,7 +13,7 @@ class TestTreeSitterSplitter(unittest.TestCase):
         """Set up the tests."""
         model_name = "gpt-4o"
         self.maxDiff = None
-        self.llm, _, _, _ = load_model(model_name)
+        self.llm = load_model(model_name)
 
     def _split(self):
         """Split the test file."""

janus/language/treesitter/treesitter.py

@@ -7,10 +7,10 @@ from typing import Optional
 
 import tree_sitter
 from git import Repo
-from langchain.schema.language_model import BaseLanguageModel
 
 from janus.language.block import CodeBlock, NodeType
 from janus.language.splitter import Splitter
+from janus.llm.models_info import JanusModel
 from janus.utils.enums import LANGUAGES
 from janus.utils.logger import create_logger
 
@@ -25,7 +25,7 @@ class TreeSitterSplitter(Splitter):
     def __init__(
         self,
         language: str,
-        model: None | BaseLanguageModel = None,
+        model: JanusModel | None = None,
         max_tokens: int = 4096,
         protected_node_types: tuple[str, ...] = (),
         prune_node_types: tuple[str, ...] = (),

janus/llm/model_callbacks.py

@@ -12,6 +12,17 @@ from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 
+openai_model_reroutes = {
+    "gpt-4o": "gpt-4o-2024-05-13",
+    "gpt-4o-mini": "gpt-4o-mini",
+    "gpt-4": "gpt-4-0613",
+    "gpt-4-turbo": "gpt-4-turbo-2024-04-09",
+    "gpt-4-turbo-preview": "gpt-4-0125-preview",
+    "gpt-3.5-turbo": "gpt-3.5-turbo-0125",
+    "gpt-3.5-turbo-16k": "gpt-3.5-turbo-0125",
+    "gpt-3.5-turbo-16k-0613": "gpt-3.5-turbo-0125",
+}
+
 
 # Updated 2024-06-21
 COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
@@ -45,6 +56,8 @@ def _get_token_cost(
     prompt_tokens: int, completion_tokens: int, model_id: str | None
 ) -> float:
     """Get the cost of tokens according to model ID"""
+    if model_id in openai_model_reroutes:
+        model_id = openai_model_reroutes[model_id]
     if model_id not in COST_PER_1K_TOKENS:
         raise ValueError(
             f"Unknown model: {model_id}. Please provide a valid model name."