janus-llm 4.3.1__py3-none-any.whl → 4.4.5__py3-none-any.whl

janus/converter/diagram.py

@@ -14,6 +14,7 @@ class DiagramGenerator(Documenter):
         self,
         diagram_type="Activity",
         add_documentation=False,
+        extract_variables=False,
         **kwargs,
     ) -> None:
         """Initialize the DiagramGenerator class
@@ -28,24 +29,25 @@ class DiagramGenerator(Documenter):
         self._documenter = Documenter(**kwargs)
 
         super().__init__(**kwargs)
-
-        self.set_prompt("diagram_with_documentation" if add_documentation else "diagram")
+        prompts = []
+        if extract_variables:
+            prompts.append("extract_variables")
+        prompts += ["diagram_with_documentation" if add_documentation else "diagram"]
+        self.set_prompts(prompts)
         self._parser = UMLSyntaxParser(language="plantuml")
 
         self._load_parameters()
 
-    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:
             return RunnableParallel(
                 SOURCE_CODE=self._parser.parse_input,
                 DOCUMENTATION=self._documenter.chain,
                 context=self._retriever,
+                DIAGRAM_TYPE=lambda x: self._diagram_type,
             )
         return RunnableParallel(
             SOURCE_CODE=self._parser.parse_input,
             context=self._retriever,
+            DIAGRAM_TYPE=lambda x: self._diagram_type,
         )

janus/converter/document.py

@@ -5,10 +5,8 @@ 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.parsers.doc_parser import ClozeDocumentationParser, MultiDocumentationParser
+from janus.parsers.parser import JanusParserException
 from janus.utils.enums import LANGUAGES
 from janus.utils.logger import create_logger
 
@@ -21,7 +19,7 @@ class Documenter(Converter):
     ):
         kwargs.update(source_language=source_language)
         super().__init__(**kwargs)
-        self.set_prompt("document")
+        self.set_prompts("document")
 
         if drop_comments:
             comment_node_type = LANGUAGES[source_language].get(
@@ -35,12 +33,14 @@ class Documenter(Converter):
 class MultiDocumenter(Documenter):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.set_prompt("multidocument")
+        self.set_prompts("multidocument")
         self._combiner = JsonCombiner()
         self._parser = MultiDocumentationParser()
 
+        self._load_parameters()
+
 
-class MadLibsDocumenter(Documenter):
+class ClozeDocumenter(Documenter):
     def __init__(
         self,
         comments_per_request: int | None = None,
@@ -48,12 +48,14 @@ class MadLibsDocumenter(Documenter):
     ) -> None:
         kwargs.update(drop_comments=False)
         super().__init__(**kwargs)
-        self.set_prompt("document_madlibs")
+        self.set_prompts("document_cloze")
         self._combiner = JsonCombiner()
-        self._parser = MadlibsDocumentationParser()
+        self._parser = ClozeDocumentationParser()
 
         self.comments_per_request = comments_per_request
 
+        self._load_parameters()
+
     def _add_translation(self, block: TranslatedCodeBlock):
         if block.translated:
             return
@@ -92,7 +94,6 @@ class MadLibsDocumenter(Documenter):
 
         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,
@@ -114,16 +115,26 @@ class MadLibsDocumenter(Documenter):
             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
+            try:
+                super()._add_translation(working_block)
+            except JanusParserException as e:
+                block.text += "\n===============\n" + working_block.text
+                block.tokens = self._llm.get_num_tokens(block.text)
+                raise e
+            finally:
+                # Update metadata to include for all runs
+                block.num_requests += working_block.num_requests
+                block.cost += working_block.cost
+                block.processing_time += working_block.processing_time
+                block.request_input_tokens += working_block.request_input_tokens
+                block.request_output_tokens += working_block.request_output_tokens
 
             # 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))
+            # Set intermediate text, will be overwritten if file
+            # successfully completes
+            block.text = json.dumps(obj)
 
         self._parser.parse_input(block.original)
         block.text = self._parser.parse(json.dumps(obj))

janus/converter/evaluate.py

@@ -1,11 +1,12 @@
 import json
 import re
 from copy import deepcopy
+from pathlib import Path
+from typing import Any
 
 from langchain_core.runnables import Runnable, RunnableLambda, RunnableParallel
 
 from janus.converter.converter import Converter
-from janus.language.block import TranslatedCodeBlock
 from janus.language.combine import JsonCombiner
 from janus.parsers.eval_parsers.incose_parser import IncoseParser
 from janus.parsers.eval_parsers.inline_comment_parser import InlineCommentParser
@@ -30,11 +31,11 @@ class Evaluator(Converter):
 
         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.
+                `OPENAI_API_KEY` environment variable must be set.
             model_arguments: Additional arguments to pass to the LLM constructor.
             max_prompts: The maximum number of prompts to try before giving up.
         """
+        kwargs.update(use_janus_inputs=True)
         super().__init__(**kwargs)
         self._combiner = JsonCombiner()
         self._load_parameters()
@@ -55,15 +56,14 @@ class RequirementEvaluator(Evaluator):
 
         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.
+                `OPENAI_API_KEY` environment variable must be set.
             model_arguments: Additional arguments to pass to the LLM constructor.
             max_prompts: The maximum number of prompts to try before giving up.
         """
         super().__init__(**kwargs)
         self.eval_items_per_request = eval_items_per_request
         self._parser = IncoseParser()
-        self.set_prompt("eval_prompts/incose")
+        self.set_prompts("eval_prompts/incose")
 
     def _input_runnable(self) -> Runnable:
         def _get_code(json_text: str) -> str:
@@ -78,67 +78,55 @@ class RequirementEvaluator(Evaluator):
             context=self._retriever,
         )
 
-    def _add_translation(self, block: TranslatedCodeBlock):
-        if block.translated:
-            return
-
-        if block.original.text is None:
-            block.translated = True
-            return
-
-        if self.eval_items_per_request is None:
-            return super()._add_translation(block)
-
-        input_obj = json.loads(block.original.text)
-        requirements = input_obj.get("requirements", [])
-
-        if not requirements:
-            log.debug(f"[{block.name}] Skipping empty block")
-            block.translated = True
-            block.text = None
-            block.complete = True
-            return
-
-        # For some reason requirements objects are in nested lists?
-        while isinstance(requirements[0], list):
-            requirements = [r for lst in requirements for r in lst]
-
-        if len(requirements) <= self.eval_items_per_request:
-            input_obj["requirements"] = requirements
-            block.original.text = json.dumps(input_obj)
-            return super()._add_translation(block)
-
-        block.processing_time = 0
-        block.cost = 0
-        block.retries = 0
-        obj = {}
-        for i in range(0, len(requirements), self.eval_items_per_request):
-            # Build a new TranslatedBlock using the new working text
-            working_requirements = requirements[i : i + self.eval_items_per_request]
-            working_copy = deepcopy(block.original)
-            working_obj = json.loads(working_copy.text)  # type: ignore
-            working_obj["requirements"] = working_requirements
-            working_copy.text = json.dumps(working_obj)
-            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
-            obj.update(json.loads(working_block.text))
-
-        block.text = json.dumps(obj)
-        block.tokens = self._llm.get_num_tokens(block.text)
-        block.translated = True
-
-        log.debug(
-            f"[{block.name}] Output code:\n{json.dumps(json.loads(block.text), indent=2)}"
-        )
+    def translate_janus_obj(self, obj: Any, name: str, failure_path: Path | None = None):
+        results = []
+        for o in obj["outputs"]:
+            if isinstance(o, dict):
+                results += self.translate_janus_obj(o, name, failure_path)
+            elif isinstance(o, str):
+                temp_obj = deepcopy(obj)
+                requirements = json.loads(o)
+                if not requirements:
+                    log.debug(f"[{name}] Skipping empty output")
+                    continue
+                if (
+                    not self.eval_items_per_request
+                    or len(requirements) < self.eval_items_per_request
+                ):
+                    obj_str = json.dumps(
+                        dict(
+                            requirements=requirements,
+                            code=obj["input"],
+                        )
+                    )
+                    temp_obj["outputs"] = [obj_str]
+                    temp_block = self._janus_object_to_codeblock(temp_obj, name)
+                    translated_block = self.translate_block(temp_block, failure_path)
+                    translated_block.previous_generations[-1] = obj
+                    translated_block.original = self._janus_object_to_codeblock(obj, name)
+                    results.append(translated_block)
+                else:
+                    for i in range(0, len(requirements), self.eval_items_per_request):
+                        working_requirements = requirements[
+                            i : i + self.eval_items_per_request
+                        ]
+                        obj_str = json.dumps(
+                            dict(
+                                requirements=working_requirements,
+                                code=obj["input"],
+                            )
+                        )
+                        temp_obj["outputs"] = [obj_str]
+                        temp_block = self._janus_object_to_codeblock(temp_obj, name)
+                        translated_block = self.translate_block(temp_block, failure_path)
+                        translated_block.previous_generations[-1] = obj
+                        translated_block.original = self._janus_object_to_codeblock(
+                            obj, name
+                        )
+                        results.append(translated_block)
+            else:
+                raise ValueError(f"Error: unable to find janus object: {type(o)}")
+        return results
 
 
 class InlineCommentEvaluator(Evaluator):
@@ -153,89 +141,100 @@ class InlineCommentEvaluator(Evaluator):
 
         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.
+                `OPENAI_API_KEY` environment variable must be set.
             model_arguments: Additional arguments to pass to the LLM constructor.
             max_prompts: The maximum number of prompts to try before giving up.
         """
         super().__init__(**kwargs)
         self._combiner = JsonCombiner()
-        self._load_parameters()
         self._parser = InlineCommentParser()
-        self.set_prompt("eval_prompts/inline_comments")
+        self.set_prompts("eval_prompts/inline_comments")
         self.eval_items_per_request = eval_items_per_request
+        self._load_parameters()
 
-    def _add_translation(self, block: TranslatedCodeBlock):
-        if block.translated:
-            return
-
-        if block.original.text is None:
-            block.translated = True
-            return
-
-        if self.eval_items_per_request is None:
-            return super()._add_translation(block)
-
+    def _process_comments(self, input_str: str, generated_comments: dict[str, str]):
+        comment_patterns = [
+            (r"<BLOCK_COMMENT (\w{8})>", "<BLOCK_COMMENT {}>", "<BLOCK_COMMENT {}>"),
+            (r"<INLINE_COMMENT (\w{8})>", "<INLINE_COMMENT {}>", "<INLINE_COMMENT {}>"),
+            (r"<MODULE (\w{8})>", "<MODULE {}>", "<BLOCK_COMMENT {}>"),
+        ]
+        missing_comments = 0
+        for pattern, find_template, repl_template in comment_patterns:
+            matches = re.findall(pattern, input_str)
+
+            for comment_id in matches:
+                find_tag = find_template.format(comment_id)
+                repl_tag = repl_template.format(comment_id)
+
+                if comment_id not in generated_comments:
+                    missing_comments += 1
+                comment = generated_comments.get(comment_id, "[comment missing]")
+                comment = comment.replace("\n", "\\n")
+
+                # Replace the tag in the code with the comment appended.
+                input_str = input_str.replace(find_tag, f"{repl_tag} {comment}")
+        processed_str = re.sub(r"\s*<JANUS_PARTITION>\s*\n", "\n", input_str)
+        return processed_str.strip("\n"), missing_comments
+
+    def translate_janus_obj(self, obj: Any, name: str, failure_path: Path | None = None):
         comment_pattern = r"<(?:INLINE|BLOCK)_COMMENT \w{8}>.*$"
-        comments = list(
-            re.finditer(comment_pattern, block.original.text, flags=re.MULTILINE)
-        )
-
-        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.eval_items_per_request:
-            return super()._add_translation(block)
-
-        comment_group_indices = list(range(0, len(comments), self.eval_items_per_request))
-        log.debug(
-            f"[{block.name}] Block contains more than {self.eval_items_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.eval_items_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.eval_items_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, flags=re.MULTILINE)
-            suffix = re.sub(comment_pattern, "", suffix, flags=re.MULTILINE)
-
-            # 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
-            obj.update(json.loads(working_block.text))
-
-        block.text = json.dumps(obj)
-        block.tokens = self._llm.get_num_tokens(block.text)
-        block.translated = True
-
-        log.debug(
-            f"[{block.name}] Output code:\n{json.dumps(json.loads(block.text), indent=2)}"
-        )
+        results = []
+        input_str = obj["input"]
+        for o in obj["outputs"]:
+            if isinstance(o, dict):
+                results += self.translate_janus_obj(o, name, failure_path)
+            elif isinstance(o, str):
+                temp_obj = deepcopy(obj)
+                generated_comments = json.loads(o)
+                processed_input, missing_comments = self._process_comments(
+                    input_str, generated_comments
+                )
+                if missing_comments:
+                    log.info(f"[{name}] Warning: missing {missing_comments} comments")
+                comments = list(
+                    re.finditer(comment_pattern, processed_input, flags=re.MULTILINE)
+                )
+                if not comments:
+                    log.info(f"[{name}] Skipping commentless block")
+                    continue
+                if (
+                    self.eval_items_per_request is None
+                    or len(comments) < self.eval_items_per_request
+                ):
+                    temp_obj["outputs"] = [processed_input]
+                    temp_block = self._janus_object_to_codeblock(temp_obj, name)
+                    translated_block = self.translate_block(temp_block, failure_path)
+                    translated_block.previous_generations[-1] = obj
+                    translated_block.original = self._janus_object_to_codeblock(obj, name)
+                    results.append(translated_block)
+                    continue
+                comment_group_indices = list(
+                    range(0, len(comments), self.eval_items_per_request)
+                )
+                log.debug(
+                    f"[{name}] Block contains more than {self.eval_items_per_request}"
+                    f" comments, splitting {len(comments)} comments into"
+                    f" {len(comment_group_indices)} groups"
+                )
+                for comment_ind in comment_group_indices:
+                    working_comments = comments[
+                        comment_ind : comment_ind + self.eval_items_per_request
+                    ]
+                    start_idx = working_comments[0].start()
+                    end_idx = working_comments[-1].end()
+                    prefix = processed_input[:start_idx]
+                    keeper = processed_input[start_idx:end_idx]
+                    suffix = processed_input[end_idx:]
+
+                    # Strip all comment placeholders outside of the section of interest
+                    prefix = re.sub(comment_pattern, "", prefix, flags=re.MULTILINE)
+                    suffix = re.sub(comment_pattern, "", suffix, flags=re.MULTILINE)
+                    temp_obj["outputs"] = [prefix + keeper + suffix]
+                    temp_block = self._janus_object_to_codeblock(temp_obj, name)
+                    translated_block = self.translate_block(temp_block, failure_path)
+                    translated_block.previous_generations[-1] = obj
+                    translated_block.original = self._janus_object_to_codeblock(obj, name)
+                    results.append(translated_block)
+            else:
+                raise ValueError(f"Error: unrecognized janus object type: {type(o)}")
+        return results

janus/converter/partition.py

@@ -1,7 +1,4 @@
-from pathlib import Path
-
 from janus.converter.converter import Converter
-from janus.language.block import TranslatedCodeBlock
 from janus.parsers.partition_parser import PartitionParser
 from janus.utils.logger import create_logger
 
@@ -11,17 +8,12 @@ log = create_logger(__name__)
 class Partitioner(Converter):
     def __init__(self, partition_token_limit: int, **kwargs):
         super().__init__(**kwargs)
-        self.set_prompt("partition")
+        self.set_prompts("partition")
         self._load_model()
         self._parser = PartitionParser(
             token_limit=partition_token_limit,
             model=self._llm,
         )
         self._target_language = self._source_language
-        self._target_suffix = self._source_suffix
+        self._target_suffix = self._source_suffixes[0]
         self._load_parameters()
-
-    def _save_to_file(self, block: TranslatedCodeBlock, out_path: Path) -> None:
-        output_str = self._parser.parse_combined_output(block.complete_text)
-        out_path.parent.mkdir(parents=True, exist_ok=True)
-        out_path.write_text(output_str, encoding="utf-8")

janus/converter/requirements.py

@@ -1,8 +1,4 @@
-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
@@ -16,41 +12,9 @@ class RequirementsDocumenter(Documenter):
     A class that translates code from one programming language to its requirements.
     """
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.set_prompt("requirements")
+    def __init__(self, combine_output: bool = False, **kwargs):
+        super().__init__(combine_output=combine_output, **kwargs)
+        self.set_prompts("requirements")
         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.
-
-        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
-        blocks = [block for block in block.children] if len(block.children) else [block]
-        for block in blocks:
-            code = block.original.text
-            requirements = self._parser.parse_combined_output(block.complete_text)
-            metadata = dict(
-                retries=block.total_retries,
-                cost=block.total_cost,
-                processing_time=block.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")
+        self._load_parameters()

janus/converter/translate.py

@@ -1,8 +1,5 @@
 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__)
@@ -21,8 +18,7 @@ class Translator(Converter):
 
         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.
+                `OPENAI_API_KEY` environment variable must be set.
             model_arguments: Additional arguments to pass to the LLM constructor.
             source_language: The source programming language.
             target_language: The target programming language.
@@ -30,13 +26,11 @@ class Translator(Converter):
             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.
+            prompt_templates: name of prompt template directories
+                (see janus/prompts/templates) or paths to directories.
         """
         super().__init__(**kwargs)
 
-        self._target_version: str | None
-
         self.set_target_language(
             target_language=target_language,
             target_version=target_version,
@@ -48,56 +42,6 @@ class Translator(Converter):
         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._llm.short_model_id](
-            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.