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

janus/converter/evaluate.py

@@ -1,12 +1,11 @@
 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 CodeBlock, 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
@@ -51,7 +50,13 @@ class RequirementEvaluator(Evaluator):
 
     """
 
-    def __init__(self, eval_items_per_request: int | None = None, **kwargs) -> None:
+    def __init__(
+        self,
+        eval_items_per_request: int | None = None,
+        input_types: str | set[str] = set(["requirements"]),
+        output_type: str = "requirements_eval",
+        **kwargs,
+    ) -> None:
         """Initialize the Evaluator class
 
         Arguments:
@@ -60,6 +65,7 @@ class RequirementEvaluator(Evaluator):
             model_arguments: Additional arguments to pass to the LLM constructor.
             max_prompts: The maximum number of prompts to try before giving up.
         """
+        kwargs.update(input_types=input_types, output_type=output_type)
         super().__init__(**kwargs)
         self.eval_items_per_request = eval_items_per_request
         self._parser = IncoseParser()
@@ -78,55 +84,67 @@ class RequirementEvaluator(Evaluator):
             context=self._retriever,
         )
 
-    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"],
-                        )
+    def translate_block(self, input_block: CodeBlock, failure_path: Path | None = None):
+        if len(input_block.previous_generations) == 0:
+            raise ValueError(
+                "Error: Evaluating requirements without previous generations"
+            )
+        if isinstance(input_block.previous_generations[-1], dict):
+            input_str = input_block.previous_generations[-1]["input"]
+        else:
+            input_str = input_block.previous_generations[-1].original.text
+        requirements = json.loads(input_block.text)
+        # The requirements are often a list of lists
+        if isinstance(requirements[0], list):
+            requirements = requirements[0]
+        if not requirements:
+            log.debug(f"[{input_block.name}] Skipping empty output")
+            return []
+        if (
+            not self.eval_items_per_request
+            or len(requirements) < self.eval_items_per_request
+        ):
+            obj_str = json.dumps(
+                dict(
+                    requirements=requirements,
+                    code=input_str,
+                )
+            )
+            temp_block = self._split_text(obj_str, input_block.name)
+            translated_block = super().translate_block(temp_block, failure_path)
+            translated_block.original = input_block
+            translated_block.previous_generations = input_block.previous_generations
+            return translated_block
+        else:
+            translated_blocks = []
+            translated_str: str
+            translate_obj = {}
+            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=input_str,
                     )
-                    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
+                )
+                temp_block = self._split_text(obj_str, input_block.name)
+                translated_block = super().translate_block(temp_block, failure_path)
+                translated_blocks.append(translated_block)
+                translate_obj.update(json.loads(translated_block.text))
+                translated_str = json.dumps(translate_obj)
+
+        translated_block = TranslatedCodeBlock(
+            input_block,
+            self._target_language,
+            self,
+            self._output_type,
+            self._output_label,
+        )
+        translated_block.text = translated_str
+        translated_block.children = translated_blocks
+        translated_block.tokens = self._llm.get_num_tokens(translated_str)
+        translated_block.translated = True
+        return translated_block
 
 
 class InlineCommentEvaluator(Evaluator):
@@ -136,7 +154,13 @@ class InlineCommentEvaluator(Evaluator):
     with an associated prompt.
     """
 
-    def __init__(self, eval_items_per_request: int | None = None, **kwargs) -> None:
+    def __init__(
+        self,
+        eval_items_per_request: int | None = None,
+        input_types: str | set[str] = set(["cloze_comments"]),
+        output_type: str = "cloze_comments_eval",
+        **kwargs,
+    ) -> None:
         """Initialize the Evaluator class
 
         Arguments:
@@ -145,6 +169,7 @@ class InlineCommentEvaluator(Evaluator):
             model_arguments: Additional arguments to pass to the LLM constructor.
             max_prompts: The maximum number of prompts to try before giving up.
         """
+        kwargs.update(input_types=input_types, output_type=output_type)
         super().__init__(**kwargs)
         self._combiner = JsonCombiner()
         self._parser = InlineCommentParser()
@@ -176,65 +201,75 @@ class InlineCommentEvaluator(Evaluator):
         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):
+    def translate_block(self, input_block: CodeBlock, failure_path: Path | None = None):
         comment_pattern = r"<(?:INLINE|BLOCK)_COMMENT \w{8}>.*$"
-        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
+        if len(input_block.previous_generations) == 0:
+            raise ValueError(
+                "Error: cannot evaluate block, no previous generations found"
+            )
+        if isinstance(input_block.previous_generations[-1], dict):
+            input_str = input_block.previous_generations[-1]["input"]
+        else:
+            input_str = input_block.previous_generations[-1].original.text
+        generated_comments = json.loads(input_block.text)
+        processed_input, missing_comments = self._process_comments(
+            input_str, generated_comments
+        )
+        if missing_comments:
+            log.info(f"[{input_block.name}] Warning: missing {missing_comments} comments")
+        comments = list(re.finditer(comment_pattern, processed_input, flags=re.MULTILINE))
+        if not comments:
+            log.info(f"[{input_block.name}] Skipping commentless block")
+            return []
+        if (
+            self.eval_items_per_request is None
+            or len(comments) < self.eval_items_per_request
+        ):
+            temp_block = self._split_text(processed_input, input_block.name)
+            translated_block = super().translate_block(temp_block, failure_path)
+            translated_block.original = input_block
+            translated_block.previous_generations = input_block.previous_generations
+            return translated_block
+        else:
+            comment_group_indices = list(
+                range(0, len(comments), self.eval_items_per_request)
+            )
+            log.debug(
+                f"[{input_block.name}]"
+                f" Block contains more than {self.eval_items_per_request}"
+                f" comments, splitting {len(comments)} comments into"
+                f" {len(comment_group_indices)} groups"
+            )
+            translated_blocks = []
+            translated_str: str
+            translate_obj = {}
+            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_block = self._split_text(prefix + keeper + suffix, input_block.name)
+                translated_block = super().translate_block(temp_block, failure_path)
+                translated_blocks.append(translated_block)
+                translate_obj.update(json.loads(translated_block.text))
+                translated_str = json.dumps(translate_obj)
+            translated_block = TranslatedCodeBlock(
+                input_block,
+                self._target_language,
+                self,
+                self._output_type,
+                self._output_label,
+            )
+            translated_block.children = translated_blocks
+            translated_block.text = translated_str
+            translated_block.tokens = self._llm.get_num_tokens(translated_str)
+            translated_block.translated = True
+            return translated_block

janus/converter/partition.py

@@ -6,7 +6,10 @@ log = create_logger(__name__)
 
 
 class Partitioner(Converter):
-    def __init__(self, partition_token_limit: int, **kwargs):
+    def __init__(
+        self, partition_token_limit: int, output_type: str = "partition", **kwargs
+    ):
+        kwargs.update(output_type=output_type)
         super().__init__(**kwargs)
         self.set_prompts("partition")
         self._load_model()

janus/converter/passthrough.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+from janus.converter.converter import Converter
+from janus.language.block import CodeBlock, TranslatedCodeBlock
+
+
+class ConverterPassthrough(Converter):
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def translate_block(
+        self, input_block: CodeBlock, failure_path: Path | None = None
+    ) -> TranslatedCodeBlock:
+        self._output_label = input_block.block_label
+        self._output_type = input_block.block_type
+        res = super().translate_block(input_block, failure_path)
+        if isinstance(input_block.previous_generations[-1], dict):
+            res.original = self._split_text(
+                input_block.previous_generations[-1]["input"], res.name
+            )
+        else:
+            res.original = input_block.previous_generations[-1].original
+        res.previous_generations = input_block.previous_generations[:-1]
+        return res
+
+    def _add_translation(self, block: TranslatedCodeBlock) -> None:
+        block.text = block.original.text
+        block.tokens = block.original.tokens
+        block.translated = True

janus/converter/pool.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+
+from janus.converter.converter import Converter
+from janus.language.block import BlockCollection, CodeBlock, TranslatedCodeBlock
+
+
+class ConverterPool(Converter):
+    def __init__(self, *args, **kwargs):
+        if len(args) == 0:
+            raise ValueError("Error: Converter chain must be passed at least 1 converter")
+        for converter in args:
+            if not isinstance(converter, Converter):
+                raise ValueError(f"Error: unrecognized type: {type(converter)}")
+        self._converters = args
+        if "source_language" in kwargs:
+            for c in self._converters:
+                c.set_source_language(kwargs["source_language"])
+        if "model" in kwargs:
+            for c in self._converters:
+                c.set_model(kwargs["model"])
+        super().__init__(**kwargs)
+
+    def translate_blocks(
+        self, input_blocks: CodeBlock | BlockCollection, failure_path: Path | None = None
+    ):
+        output_blocks = []
+        for c in self._converters:
+            collection = c.translate_blocks(input_blocks)
+            for b in collection.blocks:
+                c._combiner.combine(b)
+            output_blocks += collection.blocks
+        return BlockCollection(output_blocks, input_blocks.previous_generations)
+
+    def _get_output_obj(
+        self,
+        block: TranslatedCodeBlock | BlockCollection | dict,
+        combine_children: bool = True,
+        include_previous_outputs: bool = True,
+    ) -> dict[str, int | float | str | dict[str, str] | dict[str, float]]:
+        outputs = []
+        for b in block.blocks:
+            for c in self._converters:
+                if c == b.converter:
+                    outputs.append(c._get_output_obj(b, c._combine_output, False))
+                    break
+
+        def _get_input(block):
+            if isinstance(block, BlockCollection):
+                return self._combine_inputs([_get_input(b) for b in block.blocks])
+            return block.original.text or ""
+
+        out = dict(
+            input=_get_input(block),
+            metadata=dict(
+                cost=block.total_cost,
+                processing_time=block.total_processing_time,
+                num_requests=block.total_num_requests,
+                input_tokens=block.total_request_input_tokens,
+                output_tokens=block.total_request_output_tokens,
+                converter_name=self.__class__.__name__,
+                type=block.block_type,
+                label=block.block_label,
+            ),
+            outputs=outputs,
+        )
+        if include_previous_outputs and len(block.previous_generations) > 0:
+            intermediate_outputs = [
+                self._get_output_obj(g, combine_children, False)
+                for g in block.previous_generations
+                if isinstance(g, dict)
+            ]
+            if len(intermediate_outputs) > 0:
+                out["intermediate_outputs"] = intermediate_outputs
+        return out

janus/converter/requirements.py

@@ -12,7 +12,10 @@ class RequirementsDocumenter(Documenter):
     A class that translates code from one programming language to its requirements.
     """
 
-    def __init__(self, combine_output: bool = False, **kwargs):
+    def __init__(
+        self, combine_output: bool = False, output_type: str = "requirements", **kwargs
+    ):
+        kwargs.update(output_type=output_type)
         super().__init__(combine_output=combine_output, **kwargs)
         self.set_prompts("requirements")
         self._combiner = ChunkCombiner()

janus/language/_tests/test_combine.py

@@ -36,6 +36,7 @@ class TestCombiner(unittest.TestCase):
         self.translated_block = TranslatedCodeBlock(
             self.block,
             language="python",
+            converter=None,
         )
 
     def test_combine(self):

janus/language/block.py

@@ -1,9 +1,12 @@
 from functools import total_ordering
-from typing import ForwardRef, Hashable, Optional, Tuple
+from typing import TYPE_CHECKING, ForwardRef, Hashable, Optional, Tuple
 
 from janus.language.node import NodeType
 from janus.utils.logger import create_logger
 
+if TYPE_CHECKING:
+    from janus.converter.converter import Converter
+
 log = create_logger(__name__)
 
 
@@ -47,6 +50,8 @@ class CodeBlock:
         affixes: Tuple[str, str] = ("", ""),
         context_tags: dict[str, str] = {},
         previous_generations: list["TranslatedCodeBlock"] = [],
+        block_type: str | None = None,
+        block_label: str | None = None,
     ) -> None:
         self.id: Hashable = id
         self.name: Optional[str] = name
@@ -67,6 +72,8 @@ class CodeBlock:
         self.omit_prefix = True
         self.omit_suffix = False
         self.previous_generations = previous_generations
+        self.block_type = block_type
+        self.block_label = block_label
 
         if self.children:
             self.children[0].omit_prefix = False
@@ -186,12 +193,23 @@ class TranslatedCodeBlock(CodeBlock):
         translated: Whether this block has been successfully translated
     """
 
-    def __init__(self, original: CodeBlock, language: str) -> None:
+    def __init__(
+        self,
+        original: CodeBlock,
+        language: str,
+        converter: ForwardRef("Converter"),
+        block_type: str | None = None,
+        block_label: str | None = None,
+    ) -> None:
         """Create an "empty" `TranslatedCodeBlock` from the given original
 
         Arguments:
             original: The original code block
             language: The language to translate to
+            converter: the converter used to translate
+            block_type: type of the block
+            block_label: label for block
+            (for mapping outputs to inputs through ConverterChain)
 
         Returns:
             A `TranslatedCodeBlock` with the same attributes as the original, except
@@ -209,12 +227,17 @@ class TranslatedCodeBlock(CodeBlock):
             end_byte=None,
             tokens=0,
             children=[
-                TranslatedCodeBlock(child, language) for child in original.children
+                TranslatedCodeBlock(child, language, block_type, block_label)
+                for child in original.children
             ],
             affixes=original.affixes,
             previous_generations=original.previous_generations,
+            block_type=block_type,
+            block_label=block_label,
         )
+
         self.original = original
+        self.converter = converter
 
         self.complete = original.complete
         self.translated = False
@@ -279,6 +302,11 @@ class TranslatedCodeBlock(CodeBlock):
         children_sum = sum(c.total_num_requests for c in self.children)
         return children_sum + self.num_requests
 
+    @property
+    def total_processing_time(self) -> float:
+        children_sum = sum(c.total_processing_time for c in self.children)
+        return children_sum + self.processing_time
+
     @property
     def translation_completed(self) -> bool:
         """Whether or not the code block was successfully translated
@@ -317,6 +345,8 @@ class TranslatedCodeBlock(CodeBlock):
             children=[child.to_codeblock() for child in self.children],
             affixes=self.affixes,
             previous_generations=self.previous_generations + [self],
+            block_type=self.block_type,
+            block_label=self.block_label,
         )
 
     def __iadd__(self, other):
@@ -326,3 +356,54 @@ class TranslatedCodeBlock(CodeBlock):
         self.request_input_tokens += other.request_input_tokens
         self.request_output_tokens += other.request_output_tokens
         return self
+
+
+class BlockCollection:
+    def __init__(
+        self,
+        blocks: list[CodeBlock],
+        previous_generations: list[ForwardRef("BlockCollection")] = [],
+    ):
+        self.blocks = blocks
+        self.previous_generations = previous_generations
+
+    def to_codeblock(self) -> ForwardRef("BlockCollection"):
+        return BlockCollection(
+            [b.to_codeblock() for b in self.blocks], self.previous_generations + [self]
+        )
+
+    @property
+    def total_cost(self):
+        return sum(b.total_cost for b in self.blocks)
+
+    @property
+    def total_processing_time(self):
+        return sum(b.total_processing_time for b in self.blocks)
+
+    @property
+    def total_request_input_tokens(self):
+        return sum(b.total_request_input_tokens for b in self.blocks)
+
+    @property
+    def total_request_output_tokens(self):
+        return sum(b.total_request_output_tokens for b in self.blocks)
+
+    @property
+    def total_num_requests(self):
+        return sum(b.total_num_requests for b in self.blocks)
+
+    @property
+    def block_type(self):
+        return None
+
+    @property
+    def block_label(self):
+        return None
+
+    @property
+    def translation_completed(self):
+        return all(b.translation_completed for b in self.blocks)
+
+    @property
+    def complete(self):
+        return all(b.complete for b in self.blocks)

janus/llm/model_callbacks.py

@@ -44,12 +44,18 @@ COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
     "anthropic.claude-instant-v1": {"input": 0.0008, "output": 0.0024},
     "anthropic.claude-3-haiku-20240307-v1:0": {"input": 0.00025, "output": 0.00125},
     "anthropic.claude-3-sonnet-20240229-v1:0": {"input": 0.003, "output": 0.015},
+    "anthropic.claude-3-5-sonnet-20240620-v1:0": {"input": 0.003, "output": 0.015},
+    "anthropic.claude-3-5-sonnet-20241022-v2:0": {"input": 0.003, "output": 0.015},
     "meta.llama2-13b-chat-v1": {"input": 0.00075, "output": 0.001},
     "meta.llama2-70b-chat-v1": {"input": 0.00195, "output": 0.00256},
     "meta.llama2-13b-v1": {"input": 0.0, "output": 0.0},
     "meta.llama2-70b-v1": {"input": 0.00265, "output": 0.0035},
     "meta.llama3-8b-instruct-v1:0": {"input": 0.0003, "output": 0.0006},
     "meta.llama3-70b-instruct-v1:0": {"input": 0.00265, "output": 0.0035},
+    "meta.llama3-3-70b-instruct-v1:0": {"input": 0.00072, "output": 0.00072},
+    "amazon.nova-lite-v1:0": {"input": 0.00006, "output": 0.00024},
+    "amazon.nova-micro-v1:0": {"input": 0.000035, "output": 0.00014},
+    "amazon.nova-pro-v1:0": {"input": 0.0008, "output": 0.0032},
     "amazon.titan-text-lite-v1": {"input": 0.00015, "output": 0.0002},
     "amazon.titan-text-express-v1": {"input": 0.0002, "output": 0.0006},
     "ai21.j2-mid-v1": {"input": 0.0125, "output": 0.0125},