janus-llm 3.5.3__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, ...] = (),
@@ -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/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, ...] = (),

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."

janus/llm/models_info.py

@@ -2,14 +2,14 @@ import json
 import os
 import time
 from pathlib import Path
-from typing import Any, Callable
+from typing import Protocol, TypeVar
 
 from dotenv import load_dotenv
 from langchain_community.llms import HuggingFaceTextGenInference
-from langchain_core.language_models import BaseLanguageModel
+from langchain_core.runnables import Runnable
 from langchain_openai import ChatOpenAI
 
-from janus.llm.model_callbacks import COST_PER_1K_TOKENS
+from janus.llm.model_callbacks import COST_PER_1K_TOKENS, openai_model_reroutes
 from janus.prompts.prompt import (
     ChatGptPromptEngine,
     ClaudePromptEngine,
@@ -44,17 +44,34 @@ except ImportError:
     )
 
 
+ModelType = TypeVar(
+    "ModelType",
+    ChatOpenAI,
+    HuggingFaceTextGenInference,
+    Bedrock,
+    BedrockChat,
+    HuggingFacePipeline,
+)
+
+
+class JanusModelProtocol(Protocol):
+    model_id: str
+    model_type_name: str
+    token_limit: int
+    input_token_cost: float
+    output_token_cost: float
+    prompt_engine: type[PromptEngine]
+
+    def get_num_tokens(self, text: str) -> int:
+        ...
+
+
+class JanusModel(Runnable, JanusModelProtocol):
+    ...
+
+
 load_dotenv()
 
-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",
-}
 
 openai_models = [
     "gpt-4o",
@@ -105,24 +122,15 @@ bedrock_models = [
 ]
 all_models = [*openai_models, *bedrock_models]
 
-MODEL_TYPE_CONSTRUCTORS: dict[str, Callable[[Any], BaseLanguageModel]] = {
+MODEL_TYPE_CONSTRUCTORS: dict[str, ModelType] = {
     "OpenAI": ChatOpenAI,
     "HuggingFace": HuggingFaceTextGenInference,
+    "Bedrock": Bedrock,
+    "BedrockChat": BedrockChat,
+    "HuggingFaceLocal": HuggingFacePipeline,
 }
 
-try:
-    MODEL_TYPE_CONSTRUCTORS.update(
-        {
-            "HuggingFaceLocal": HuggingFacePipeline.from_model_id,
-            "Bedrock": Bedrock,
-            "BedrockChat": BedrockChat,
-        }
-    )
-except NameError:
-    pass
-
-
-MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
+MODEL_PROMPT_ENGINES: dict[str, type[PromptEngine]] = {
     **{m: ChatGptPromptEngine for m in openai_models},
     **{m: ClaudePromptEngine for m in claude_models},
     **{m: Llama2PromptEngine for m in llama2_models},
@@ -132,11 +140,6 @@ MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
     **{m: MistralPromptEngine for m in mistral_models},
 }
 
-_open_ai_defaults: dict[str, str] = {
-    "openai_api_key": os.getenv("OPENAI_API_KEY"),
-    "openai_organization": os.getenv("OPENAI_ORG_ID"),
-}
-
 MODEL_ID_TO_LONG_ID = {
     **{m: mr for m, mr in openai_model_reroutes.items()},
     "bedrock-claude-v2": "anthropic.claude-v2",
@@ -168,7 +171,7 @@ DEFAULT_MODELS = list(MODEL_DEFAULT_ARGUMENTS.keys())
 
 MODEL_CONFIG_DIR = Path.home().expanduser() / ".janus" / "llm"
 
-MODEL_TYPES: dict[str, PromptEngine] = {
+MODEL_TYPES: dict[str, str] = {
     **{m: "OpenAI" for m in openai_models},
     **{m: "BedrockChat" for m in bedrock_models},
 }
@@ -211,53 +214,90 @@ def get_available_model_names() -> list[str]:
     return avaialable_models
 
 
-def load_model(
-    user_model_name: str,
-) -> tuple[BaseLanguageModel, str, int, dict[str, float]]:
+def load_model(model_id) -> JanusModel:
     if not MODEL_CONFIG_DIR.exists():
         MODEL_CONFIG_DIR.mkdir(parents=True)
-    model_config_file = MODEL_CONFIG_DIR / f"{user_model_name}.json"
-    if not model_config_file.exists():
-        log.warning(
-            f"Model {user_model_name} not found in user-defined models, searching "
-            f"default models for {user_model_name}."
-        )
-        model_id = user_model_name
-        if user_model_name not in DEFAULT_MODELS:
-            message = (
-                f"Model {user_model_name} not found in default models. Make sure to run "
-                "`janus llm add` first."
-            )
-            log.error(message)
-            raise ValueError(message)
-        model_config = {
-            "model_type": MODEL_TYPES[model_id],
-            "model_id": model_id,
-            "model_args": MODEL_DEFAULT_ARGUMENTS[model_id],
-            "token_limit": TOKEN_LIMITS.get(MODEL_ID_TO_LONG_ID[model_id], 4096),
-            "model_cost": COST_PER_1K_TOKENS.get(
-                MODEL_ID_TO_LONG_ID[model_id], {"input": 0, "output": 0}
-            ),
-        }
-        with open(model_config_file, "w") as f:
-            json.dump(model_config, f)
-    else:
+    model_config_file = MODEL_CONFIG_DIR / f"{model_id}.json"
+
+    if model_config_file.exists():
+        log.info(f"Loading {model_id} from {model_config_file}.")
         with open(model_config_file, "r") as f:
             model_config = json.load(f)
-    model_constructor = MODEL_TYPE_CONSTRUCTORS[model_config["model_type"]]
-    model_args = model_config["model_args"]
-    if model_config["model_type"] == "OpenAI":
-        model_args.update(_open_ai_defaults)
+        model_type_name = model_config["model_type"]
+        model_id = model_config["model_id"]
+        model_args = model_config["model_args"]
+        token_limit = model_config["token_limit"]
+        input_token_cost = model_config["model_cost"]["input"]
+        output_token_cost = model_config["model_cost"]["output"]
+
+    elif model_id in DEFAULT_MODELS:
+        model_id = model_id
+        model_long_id = MODEL_ID_TO_LONG_ID[model_id]
+        model_type_name = MODEL_TYPES[model_id]
+        model_args = MODEL_DEFAULT_ARGUMENTS[model_id]
+
+        token_limit = 0
+        input_token_cost = 0.0
+        output_token_cost = 0.0
+        if model_long_id in TOKEN_LIMITS:
+            token_limit = TOKEN_LIMITS[model_long_id]
+        if model_long_id in COST_PER_1K_TOKENS:
+            token_limits = COST_PER_1K_TOKENS[model_long_id]
+            input_token_cost = token_limits["input"]
+            output_token_cost = token_limits["output"]
+
+    else:
+        model_list = "\n\t".join(DEFAULT_MODELS)
+        message = (
+            f"Model {model_id} not found in user-defined model directory "
+            f"({MODEL_CONFIG_DIR}), and is not a default model. Valid default "
+            f"models:\n\t{model_list}\n"
+            f"To use a custom model, first run `janus llm add`."
+        )
+        log.error(message)
+        raise ValueError(message)
+
+    if model_type_name == "HuggingFaceLocal":
+        model = HuggingFacePipeline.from_model_id(
+            model_id=model_id,
+            task="text-generation",
+            model_kwargs=model_args,
+        )
+        model_args.update(pipeline=model.pipeline)
+
+    elif model_type_name == "OpenAI":
+        model_args.update(
+            openai_api_key=str(os.getenv("OPENAI_API_KEY")),
+            openai_organization=str(os.getenv("OPENAI_ORG_ID")),
+        )
         log.warning("Do NOT use this model in sensitive environments!")
         log.warning("If you would like to cancel, please press Ctrl+C.")
         log.warning("Waiting 10 seconds...")
         # Give enough time for the user to read the warnings and cancel
         time.sleep(10)
 
-    model = model_constructor(**model_args)
-    return (
-        model,
-        model_config["model_id"],
-        model_config["token_limit"],
-        model_config["model_cost"],
+    model_type = MODEL_TYPE_CONSTRUCTORS[model_type_name]
+    prompt_engine = MODEL_PROMPT_ENGINES[model_id]
+
+    class JanusModel(model_type):
+        model_id: str
+        short_model_id: str
+        model_type_name: str
+        token_limit: int
+        input_token_cost: float
+        output_token_cost: float
+        prompt_engine: type[PromptEngine]
+
+    model_args.update(
+        model_id=MODEL_ID_TO_LONG_ID[model_id],
+        short_model_id=model_id,
+    )
+
+    return JanusModel(
+        model_type_name=model_type_name,
+        token_limit=token_limit,
+        input_token_cost=input_token_cost,
+        output_token_cost=output_token_cost,
+        prompt_engine=prompt_engine,
+        **model_args,
     )