janus-llm 2.1.0__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/language/block.py

@@ -216,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/treesitter/treesitter.py

@@ -1,6 +1,7 @@
 import os
 import platform
 from collections import defaultdict
+from ctypes import c_void_p, cdll
 from pathlib import Path
 from typing import Optional
 
@@ -138,7 +139,25 @@ class TreeSitterSplitter(Splitter):
 
         # Load the parser using the generated .so file
         self.parser: tree_sitter.Parser = tree_sitter.Parser()
-        self.parser.set_language(tree_sitter.Language(so_file, self.language))
+        pointer = self._so_to_pointer(so_file)
+        self.parser.set_language(tree_sitter.Language(pointer, self.language))
+
+    def _so_to_pointer(self, so_file: str) -> int:
+        """Convert the .so file to a pointer.
+
+        Taken from `treesitter.Language.__init__` to get past deprecated warning.
+
+        Arguments:
+            so_file: The path to the so file for the language.
+
+        Returns:
+            The pointer to the language.
+        """
+        lib = cdll.LoadLibrary(os.fspath(so_file))
+        language_function = getattr(lib, f"tree_sitter_{self.language}")
+        language_function.restype = c_void_p
+        pointer = language_function()
+        return pointer
 
     def _create_parser(self, so_file: Path | str) -> None:
         """Create the parser for the given language.

janus/llm/model_callbacks.py

@@ -35,6 +35,9 @@ COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
     "ai21.j2-mid-v1": {"input": 0.0125, "output": 0.0125},
     "ai21.j2-ultra-v1": {"input": 0.0188, "output": 0.0188},
     "cohere.command-r-plus-v1:0": {"input": 0.003, "output": 0.015},
+    "mistral.mistral-7b-instruct-v0:2": {"input": 0.00015, "output": 0.0002},
+    "mistral.mixtral-8x7b-instruct-v0:1": {"input": 0.00045, "output": 0.0007},
+    "mistral.mistral-large-2402-v1:0": {"input": 0.004, "output": 0.012},
 }
 
 
@@ -103,53 +106,47 @@ class TokenUsageCallbackHandler(BaseCallbackHandler):
             generation = response.generations[0][0]
         except IndexError:
             generation = None
-        if isinstance(generation, ChatGeneration):
-            try:
-                message = generation.message
-                if isinstance(message, AIMessage):
-                    usage_metadata = message.usage_metadata
-                else:
-                    usage_metadata = None
-            except AttributeError:
-                usage_metadata = None
-        else:
-            usage_metadata = None
-        if usage_metadata:
-            token_usage = {"total_tokens": usage_metadata["total_tokens"]}
-            completion_tokens = usage_metadata["output_tokens"]
-            prompt_tokens = usage_metadata["input_tokens"]
-            if response.llm_output is None:
-                # model name (and therefore cost) is unavailable in
-                # streaming responses
-                model_name = ""
-            else:
-                model_name = response.llm_output.get("model_name", "")
 
+        model_id = ""
+        usage_metadata = None
+        if hasattr(response, "llm_output") and response.llm_output is not None:
+            model_id = response.llm_output.get("model_id", model_id)
+            model_id = response.llm_output.get("model_name", model_id)
+            usage_metadata = response.llm_output.get("usage", usage_metadata)
+            usage_metadata = response.llm_output.get("token_usage", usage_metadata)
+        elif isinstance(generation, ChatGeneration):
+            if hasattr(generation, "response_metadata"):
+                model_id = generation.response_metadata.get("model_id", model_id)
+                model_id = generation.response_metadata.get("model_name", model_id)
+                usage_metadata = generation.response_metadata.get("usage", usage_metadata)
+            elif hasattr(generation, "message"):
+                if isinstance(generation.message, AIMessage):
+                    usage_metadata = generation.message.usage_metadata
+
+        completion_tokens = 0
+        prompt_tokens = 0
+        total_tokens = 0
+        if usage_metadata:
+            prompt_tokens = usage_metadata.get("prompt_tokens", prompt_tokens)
+            prompt_tokens = usage_metadata.get("input_tokens", prompt_tokens)
+            completion_tokens = usage_metadata.get("completion_tokens", completion_tokens)
+            completion_tokens = usage_metadata.get("output_tokens", completion_tokens)
+            total_tokens = usage_metadata.get("total_tokens", total_tokens)
         else:
-            if response.llm_output is None:
-                return None
-
-            if "token_usage" not in response.llm_output:
-                with self._lock:
-                    self.successful_requests += 1
-                return None
-
-            # compute tokens and cost for this request
-            token_usage = response.llm_output["token_usage"]
-            completion_tokens = token_usage.get("completion_tokens", 0)
-            prompt_tokens = token_usage.get("prompt_tokens", 0)
-            model_name = response.llm_output.get("model_name", "")
+            with self._lock:
+                self.successful_requests += 1
+            return None
 
         total_cost = _get_token_cost(
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
-            model_id=model_name,
+            model_id=model_id,
         )
 
         # update shared state behind lock
         with self._lock:
             self.total_cost += total_cost
-            self.total_tokens += token_usage.get("total_tokens", 0)
+            self.total_tokens += total_tokens
             self.prompt_tokens += prompt_tokens
             self.completion_tokens += completion_tokens
             self.successful_requests += 1

janus/llm/models_info.py

@@ -14,6 +14,7 @@ from ..prompts.prompt import (
     CoherePromptEngine,
     Llama2PromptEngine,
     Llama3PromptEngine,
+    MistralPromptEngine,
     PromptEngine,
     TitanPromptEngine,
 )
@@ -85,12 +86,18 @@ titan_models = [
 cohere_models = [
     "bedrock-command-r-plus",
 ]
+mistral_models = [
+    "bedrock-mistral-7b-instruct",
+    "bedrock-mistral-large",
+    "bedrock-mixtral",
+]
 bedrock_models = [
     *claude_models,
     *llama2_models,
     *llama3_models,
     *titan_models,
     *cohere_models,
+    *mistral_models,
 ]
 all_models = [*openai_models, *bedrock_models]
 
@@ -118,6 +125,7 @@ MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
     **{m: Llama3PromptEngine for m in llama3_models},
     **{m: TitanPromptEngine for m in titan_models},
     **{m: CoherePromptEngine for m in cohere_models},
+    **{m: MistralPromptEngine for m in mistral_models},
 }
 
 _open_ai_defaults: dict[str, str] = {
@@ -142,6 +150,9 @@ model_identifiers = {
     "bedrock-jurassic-2-mid": "ai21.j2-mid-v1",
     "bedrock-jurassic-2-ultra": "ai21.j2-ultra-v1",
     "bedrock-command-r-plus": "cohere.command-r-plus-v1:0",
+    "bedrock-mixtral": "mistral.mixtral-8x7b-instruct-v0:1",
+    "bedrock-mistral-7b-instruct": "mistral.mistral-7b-instruct-v0:2",
+    "bedrock-mistral-large": "mistral.mistral-large-2402-v1:0",
 }
 
 MODEL_DEFAULT_ARGUMENTS: dict[str, dict[str, str]] = {
@@ -182,6 +193,9 @@ TOKEN_LIMITS: dict[str, int] = {
     "ai21.j2-mid-v1": 8192,
     "ai21.j2-ultra-v1": 8192,
     "cohere.command-r-plus-v1:0": 128_000,
+    "mistral.mixtral-8x7b-instruct-v0:1": 32_000,
+    "mistral.mistral-7b-instruct-v0:2": 32_000,
+    "mistral.mistral-large-2402-v1:0": 32_000,
 }