janus-llm 3.2.1__py3-none-any.whl → 3.3.0__py3-none-any.whl

janus/__init__.py

@@ -5,7 +5,7 @@ from langchain_core._api.deprecation import LangChainDeprecationWarning
 from janus.converter.translate import Translator
 from janus.metrics import *  # noqa: F403
 
-__version__ = "3.2.1"
+__version__ = "3.3.0"
 
 # Ignoring a deprecation warning from langchain_core that I can't seem to hunt down
 warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

janus/converter/_tests/test_translate.py

@@ -45,16 +45,17 @@ class TestTranslator(unittest.TestCase):
     def setUp(self):
         """Set up the tests."""
         self.translator = Translator(
-            model="gpt-4o",
+            model="gpt-4o-mini",
             source_language="fortran",
             target_language="python",
             target_version="3.10",
+            splitter_type="ast-flex",
         )
         self.test_file = Path("janus/language/treesitter/_tests/languages/fortran.f90")
         self.TEST_FILE_EMBEDDING_COUNT = 14
 
         self.req_translator = RequirementsDocumenter(
-            model="gpt-4o",
+            model="gpt-4o-mini",
             source_language="fortran",
             prompt_template="requirements",
         )

janus/converter/converter.py

@@ -6,7 +6,6 @@ from pathlib import Path
 from typing import Any
 
 from langchain.output_parsers import RetryWithErrorOutputParser
-from langchain.output_parsers.fix import OutputFixingParser
 from langchain_core.exceptions import OutputParserException
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.output_parsers import BaseOutputParser
@@ -29,6 +28,8 @@ from janus.llm import load_model
 from janus.llm.model_callbacks import get_model_callback
 from janus.llm.models_info import MODEL_PROMPT_ENGINES
 from janus.parsers.code_parser import GenericParser
+from janus.parsers.refiner_parser import RefinerParser
+from janus.refiners.refiner import BasicRefiner, Refiner
 from janus.utils.enums import LANGUAGES
 from janus.utils.logger import create_logger
 
@@ -75,6 +76,7 @@ class Converter:
         protected_node_types: tuple[str, ...] = (),
         prune_node_types: tuple[str, ...] = (),
         splitter_type: str = "file",
+        refiner_type: str = "basic",
     ) -> None:
         """Initialize a Converter instance.
 
@@ -84,6 +86,17 @@ class Converter:
                 values are `"code"`, `"text"`, `"eval"`, and `None` (default). If `None`,
                 the `Converter` assumes you won't be parsing an output (i.e., adding to an
                 embedding DB).
+            max_prompts: The maximum number of prompts to try before giving up.
+            max_tokens: The maximum number of tokens to use in the LLM. If `None`, the
+                converter will use half the model's token limit.
+            prompt_template: The name of the prompt template to use.
+            db_path: The path to the database to use for vectorization.
+            db_config: The configuration for the database.
+            protected_node_types: A set of node types that aren't to be merged.
+            prune_node_types: A set of node types which should be pruned.
+            splitter_type: The type of splitter to use. Valid values are `"file"`,
+                `"tag"`, `"chunk"`, `"ast-strict"`, and `"ast-flex"`.
+            refiner_type: The type of refiner to use. Valid values are `"basic"`.
         """
         self._changed_attrs: set = set()
 
@@ -116,7 +129,11 @@ class Converter:
         self._parser: BaseOutputParser = GenericParser()
         self._combiner: Combiner = Combiner()
 
+        self._refiner_type: str
+        self._refiner: Refiner
+
         self.set_splitter(splitter_type=splitter_type)
+        self.set_refiner(refiner_type=refiner_type)
         self.set_model(model_name=model, **model_arguments)
         self.set_prompt(prompt_template=prompt_template)
         self.set_source_language(source_language)
@@ -142,6 +159,7 @@ class Converter:
         self._load_prompt()
         self._load_splitter()
         self._load_vectorizer()
+        self._load_refiner()
         self._changed_attrs.clear()
 
     def set_model(self, model_name: str, **custom_arguments: dict[str, Any]):
@@ -179,6 +197,16 @@ class Converter:
         """
         self._splitter_type = splitter_type
 
+    def set_refiner(self, refiner_type: str) -> None:
+        """Validate and set the refiner name
+
+        The affected objects will not be updated until translate is called
+
+        Arguments:
+            refiner_type: the name of the refiner to use
+        """
+        self._refiner_type = refiner_type
+
     def set_source_language(self, source_language: str) -> None:
         """Validate and set the source language.
 
@@ -249,10 +277,24 @@ class Converter:
         )
 
         if self._splitter_type == "tag":
-            kwargs["tag"] = "<ITMOD_ALC_SPLIT>"
+            kwargs["tag"] = "<ITMOD_ALC_SPLIT>"  # Hardcoded for now
 
         self._splitter = CUSTOM_SPLITTERS[self._splitter_type](**kwargs)
 
+    @run_if_changed("_refiner_type", "_model_name")
+    def _load_refiner(self) -> None:
+        """Load the refiner 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._refiner_type == "basic":
+            self._refiner = BasicRefiner(
+                "basic_refinement", self._model_name, self._source_language
+            )
+        else:
+            raise ValueError(f"Error: unknown refiner type {self._refiner_type}")
+
     @run_if_changed("_model_name", "_custom_model_arguments")
     def _load_model(self) -> None:
         """Load the model according to this instance's attributes.
@@ -561,22 +603,22 @@ class Converter:
         # Retries with just the input
         n3 = math.ceil(self.max_prompts / (n1 * n2))
 
-        fix_format = OutputFixingParser.from_llm(
-            llm=self._llm,
+        refine_output = RefinerParser(
             parser=self._parser,
+            initial_prompt=self._prompt.format(**{"SOURCE_CODE": block.original.text}),
+            refiner=self._refiner,
             max_retries=n1,
+            llm=self._llm,
         )
         retry = RetryWithErrorOutputParser.from_llm(
             llm=self._llm,
-            parser=fix_format,
+            parser=refine_output,
             max_retries=n2,
         )
-
         completion_chain = self._prompt | self._llm
         chain = RunnableParallel(
             completion=completion_chain, prompt_value=self._prompt
         ) | RunnableLambda(lambda x: retry.parse_with_prompt(**x))
-
         for _ in range(n3):
             try:
                 return chain.invoke({"SOURCE_CODE": block.original.text})

janus/converter/diagram.py

@@ -1,10 +1,14 @@
-import json
-from copy import deepcopy
+import math
+
+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
 
@@ -47,65 +51,74 @@ class DiagramGenerator(Documenter):
             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}")
-
+    def _run_chain(self, block: TranslatedCodeBlock) -> str:
         self._parser.set_reference(block.original)
+        n1 = round(self.max_prompts ** (1 / 3))
 
-        query_and_parse = self.diagram_prompt | self._llm | self._diagram_parser
+        # Retries with the input, output, and error
+        n2 = round((self.max_prompts // n1) ** (1 / 2))
+
+        # Retries with just the input
+        n3 = math.ceil(self.max_prompts / (n1 * n2))
 
         if self._add_documentation:
-            block.text = query_and_parse.invoke(
-                {
-                    "SOURCE_CODE": block.original.text,
-                    "DIAGRAM_TYPE": self._diagram_type,
-                    "DOCUMENTATION": documentation,
-                }
+            documentation_text = super()._run_chain(block)
+            refine_output = RefinerParser(
+                parser=self._diagram_parser,
+                initial_prompt=self._diagram_prompt.format(
+                    **{
+                        "SOURCE_CODE": block.original.text,
+                        "DOCUMENTATION": documentation_text,
+                        "DIAGRAM_TYPE": self._diagram_type,
+                    }
+                ),
+                refiner=self._refiner,
+                max_retries=n1,
+                llm=self._llm,
             )
         else:
-            block.text = query_and_parse.invoke(
-                {
-                    "SOURCE_CODE": block.original.text,
-                    "DIAGRAM_TYPE": self._diagram_type,
-                }
+            refine_output = RefinerParser(
+                parser=self._diagram_parser,
+                initial_prompt=self._diagram_prompt.format(
+                    **{
+                        "SOURCE_CODE": block.original.text,
+                        "DIAGRAM_TYPE": self._diagram_type,
+                    }
+                ),
+                refiner=self._refiner,
+                max_retries=n1,
+                llm=self._llm,
             )
-        block.tokens = self._llm.get_num_tokens(block.text)
-        block.translated = True
-
-        log.debug(f"[{block.name}] Output code:\n{block.text}")
+        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": block.original.text,
+                            "DOCUMENTATION": documentation_text,
+                            "DIAGRAM_TYPE": self._diagram_type,
+                        }
+                    )
+                else:
+                    return chain.invoke(
+                        {
+                            "SOURCE_CODE": block.original.text,
+                            "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",
@@ -123,4 +136,4 @@ class DiagramGenerator(Documenter):
             target_version=None,
             prompt_template=self._diagram_prompt_template_name,
         )
-        self.diagram_prompt = self._diagram_prompt_engine.prompt
+        self._diagram_prompt = self._diagram_prompt_engine.prompt

janus/llm/models_info.py

@@ -47,6 +47,7 @@ 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",
@@ -56,6 +57,7 @@ openai_model_reroutes = {
 
 openai_models = [
     "gpt-4o",
+    "gpt-4o-mini",
     "gpt-4",
     "gpt-4-turbo",
     "gpt-4-turbo-preview",

janus/parsers/refiner_parser.py

@@ -0,0 +1,49 @@
+from langchain_core.exceptions import OutputParserException
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.output_parsers import BaseOutputParser
+
+from janus.refiners.refiner import Refiner
+
+
+class RefinerParser(BaseOutputParser):
+    """Parser for performing refinement with a refiner
+
+    Properties:
+        llm: the language model to use
+        parser: the parser to use for parsing llm output
+        initial_prompt: initial prompt used to generate output
+        refiner: refiner that gives new subsequent prompts
+        max_retires: maximum number of times to attempt refining
+    """
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    llm: BaseLanguageModel
+    parser: BaseOutputParser
+    initial_prompt: str
+    refiner: Refiner
+    max_retries: int
+
+    def parse(self, text: str) -> str:
+        """Parses the text using the refiner
+
+        Arguments:
+            text: text to parse
+
+        Returns:
+            Parsed text
+        """
+        last_prompt = self.initial_prompt
+        for _ in range(self.max_retries):
+            try:
+                return self.parser.parse(text)
+            except OutputParserException as oe:
+                err = str(oe)
+                new_prompt, prompt_arguments = self.refiner.refine(last_prompt, text, err)
+                new_chain = new_prompt | self.llm
+                text = new_chain.invoke(prompt_arguments)
+                last_prompt = new_prompt.format(**prompt_arguments)
+        raise OutputParserException(
+            f"Error: unable to correct output after {self.max_retries} attempts"
+        )

janus/refiners/refiner.py

@@ -0,0 +1,63 @@
+from langchain_core.prompts import ChatPromptTemplate
+
+from janus.llm.models_info import MODEL_PROMPT_ENGINES
+
+
+class Refiner:
+    def refine(
+        self, original_prompt: str, original_output: str, errors: str, **kwargs
+    ) -> tuple[ChatPromptTemplate, dict[str, str]]:
+        """Creates a new prompt based on feedback from original results
+
+        Arguments:
+            original_prompt: original prompt used to produce output
+            original_output: origial output of llm
+            errors: list of errors detected by parser
+
+        Returns:
+            Tuple of new prompt and prompt arguments
+        """
+        raise NotImplementedError
+
+
+class BasicRefiner(Refiner):
+    def __init__(
+        self,
+        prompt_name: str,
+        model_name: str,
+        source_language: str,
+    ) -> None:
+        """Basic refiner, asks llm to fix output of previous prompt given errors
+
+        Arguments:
+            prompt_name: refinement prompt name to use
+            model_name: name of llm to use
+            source_language: source_langauge to use
+        """
+        self._prompt_name = prompt_name
+        self._model_name = model_name
+        self._source_language = source_language
+
+    def refine(
+        self, original_prompt: str, original_output: str, errors: str, **kwargs
+    ) -> tuple[ChatPromptTemplate, dict[str, str]]:
+        """Creates a new prompt based on feedback from original results
+
+        Arguments:
+            original_prompt: original prompt used to produce output
+            original_output: origial output of llm
+            errors: list of errors detected by parser
+
+        Returns:
+            Tuple of new prompt and prompt arguments
+        """
+        prompt_engine = MODEL_PROMPT_ENGINES[self._model_name](
+            prompt_template=self._prompt_name,
+            source_language=self._source_language,
+        )
+        prompt_arguments = {
+            "ORIGINAL_PROMPT": original_prompt,
+            "OUTPUT": original_output,
+            "ERRORS": errors,
+        }
+        return prompt_engine.prompt, prompt_arguments

{janus_llm-3.2.1.dist-info → janus_llm-3.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: janus-llm
-Version: 3.2.1
+Version: 3.3.0
 Summary: A transcoding library using LLMs.
 Home-page: https://github.com/janus-llm/janus-llm
 License: Apache 2.0

{janus_llm-3.2.1.dist-info → janus_llm-3.3.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-janus/__init__.py,sha256=pQaY-5WZCc5Axe5Vk7HvXbyvNOvjZni9lIAlfsd9Evs,361
+janus/__init__.py,sha256=U4HbuBrititICOf7Slcy-oI52SlTuid-hoIdIBbsUy8,361
 janus/__main__.py,sha256=lEkpNtLVPtFo8ySDZeXJ_NXDHb0GVdZFPWB4gD4RPS8,64
 janus/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/_tests/conftest.py,sha256=V7uW-oq3YbFiRPvrq15YoVVrA1n_83pjgiyTZ-IUGW8,963
@@ -6,9 +6,9 @@ janus/_tests/test_cli.py,sha256=oYJsUGWfpBJWEGRG5NGxdJedU5DU_m6fwJ7xEbJVYl0,4244
 janus/cli.py,sha256=S6IaQyUG55xSB166xARn6TQc_cOCIQ_ZMdIi7vj6BMk,31626
 janus/converter/__init__.py,sha256=U2EOMcCykiC0ZqhorNefOP_04hOF18qhYoPKrVp1Vrk,345
 janus/converter/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-janus/converter/_tests/test_translate.py,sha256=QylggIuLNOCovqz067L59rmky-bJlxcedzlJEwkR_T8,5583
-janus/converter/converter.py,sha256=u0XySdsrBjDHEYdPlvU8lsRLolciDil5aYtf0QvmEeo,23573
-janus/converter/diagram.py,sha256=JsJNDf-P8bPejpDxbCVEHvw-0kewiMrXh5qLhsL5JOA,4730
+janus/converter/_tests/test_translate.py,sha256=yzcFEGc_z8QmBBBmC9dZnfL9tT8CD1rkpc8Hz44Jp4c,5631
+janus/converter/converter.py,sha256=m0eGrybrPWw2B0nXrd1Fx6o0BO1JvY4PfC0rFUtpLxY,25633
+janus/converter/diagram.py,sha256=5mo1H3Y1uIBPYdIsWz9kxluN5DNyuUMZrtcJmGF2Uw0,5335
 janus/converter/document.py,sha256=hsW512veNjFWbdl5WriuUdNmMEqZy8ktRvqn9rRmA6E,4566
 janus/converter/evaluate.py,sha256=APWQUY3gjAXqkJkPzvj0UA4wPK3Cv9QSJLM-YK9t-ng,476
 janus/converter/requirements.py,sha256=6YvrJRVH9BuPCOPxnXmaJQFYmoLYYvCu3zTntDLHeNg,1832
@@ -57,7 +57,7 @@ janus/language/treesitter/_tests/test_treesitter.py,sha256=4S_UdH6AfJ0j6hyInZ2CF
 janus/language/treesitter/treesitter.py,sha256=FU86H8j2cfRLtwyNeEnf9A6gzZEvDwmnFCIrQymUJso,7541
 janus/llm/__init__.py,sha256=TKLYvnsWKWfxMucy-lCLQ-4bkN9ENotJZDywDEQmrKg,45
 janus/llm/model_callbacks.py,sha256=K7P5NY-rf7IYRAFHnZ3kzhrQWE6g_najx8uxlaSiz3E,7110
-janus/llm/models_info.py,sha256=1SX7NSxY20BBzLz7b6_XpKMRsCFsjTCW_WIGZapJ9QU,8305
+janus/llm/models_info.py,sha256=FHyv9PKXAUxYCgO06bw4qzoniFsqYN21uD450BdpvyA,8358
 janus/metrics/__init__.py,sha256=AsxtZJUzZiXJPr2ehPPltuYP-ddechjg6X85WZUO7mA,241
 janus/metrics/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/metrics/_tests/reference.py,sha256=hiaJPP9CXkvFBV_wL-gOe_BzELTw0nvB6uCxhxtIiE8,13
@@ -87,10 +87,12 @@ janus/parsers/_tests/test_code_parser.py,sha256=RVgMmLvg8_57g0uJphfX-jZZsyBqOSuG
 janus/parsers/code_parser.py,sha256=SZBsYThG4iszKlu4fHoWrs-6cbJiUFjWv4cLSr5bzDM,1790
 janus/parsers/doc_parser.py,sha256=bJiOE5M7npUZur_1MWJ14C2HZl7-yXExqRXiC5ZBJvI,5679
 janus/parsers/eval_parser.py,sha256=L1Lu2aNimcqUshe0FQee_9Zqj1rrqyZPXCgEAS05VJ4,2740
+janus/parsers/refiner_parser.py,sha256=72tOEhpHwCZqHDb2T4aS5bPsiXN3pQXUk_oOPupa3Ts,1621
 janus/parsers/reqs_parser.py,sha256=6YzpF63rjuDPqpKWfYvtjpsluWQ-UboWlsKoGrGQogA,2380
 janus/parsers/uml.py,sha256=ZRyGY8YxvYibacTd-WZEAAaW3XjmvJhPJE3o29f71t8,1825
 janus/prompts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/prompts/prompt.py,sha256=3796YXIzzIec9b0iUzd8VZlq-AdQbzq8qUGXLy4KH-0,10586
+janus/refiners/refiner.py,sha256=O4i5JaPEWH_ijmHunTKP4YzX_ZwZIyOIckn4Hmf1ZOI,2084
 janus/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/utils/_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 janus/utils/_tests/test_logger.py,sha256=jkkvrCTKwsFCsZtmyuvc-WJ0rC7LJi2Z91sIe4IiKzA,2209
@@ -98,8 +100,8 @@ janus/utils/_tests/test_progress.py,sha256=Rs_u5PiGjP-L-o6C1fhwfE1ig8jYu9Xo9s4p8
 janus/utils/enums.py,sha256=AoilbdiYyMvY2Mp0AM4xlbLSELfut2XMwhIM1S_msP4,27610
 janus/utils/logger.py,sha256=KZeuaMAnlSZCsj4yL0P6N-JzZwpxXygzACWfdZFeuek,2337
 janus/utils/progress.py,sha256=PIpcQec7SrhsfqB25LHj2CDDkfm9umZx90d9LZnAx6k,1469
-janus_llm-3.2.1.dist-info/LICENSE,sha256=_j0st0a-HB6MRbP3_BW3PUqpS16v54luyy-1zVyl8NU,10789
-janus_llm-3.2.1.dist-info/METADATA,sha256=w_DcqGuRCD4p2o63oEXEOZYdoQuxCMx-czxzR6rrh-U,4184
-janus_llm-3.2.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-janus_llm-3.2.1.dist-info/entry_points.txt,sha256=OGhQwzj6pvXp79B0SaBD5apGekCu7Dwe9fZZT_TZ544,39
-janus_llm-3.2.1.dist-info/RECORD,,
+janus_llm-3.3.0.dist-info/LICENSE,sha256=_j0st0a-HB6MRbP3_BW3PUqpS16v54luyy-1zVyl8NU,10789
+janus_llm-3.3.0.dist-info/METADATA,sha256=RpMSwpF_2g3KcVoS6tHHQjQlYHC-n3h989-qJAk6Epk,4184
+janus_llm-3.3.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+janus_llm-3.3.0.dist-info/entry_points.txt,sha256=OGhQwzj6pvXp79B0SaBD5apGekCu7Dwe9fZZT_TZ544,39
+janus_llm-3.3.0.dist-info/RECORD,,