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

janus/refiners/refiner.py

@@ -3,6 +3,7 @@ from typing import Any
 
 from langchain.output_parsers import RetryWithErrorOutputParser
 from langchain_core.exceptions import OutputParserException
+from langchain_core.messages import AIMessage
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompt_values import PromptValue
 from langchain_core.runnables import RunnableSerializable
@@ -108,7 +109,7 @@ class ReflectionRefiner(JanusRefiner):
     def parse_completion(
         self, completion: str, prompt_value: PromptValue, **kwargs
     ) -> Any:
-        log.info(f"Reflection Prompt: {self.reflection_prompt_name}")
+        log.debug(f"Reflection Prompt: {self.reflection_prompt_name}")
         for retry_number in range(self.max_retries):
             reflection = self.reflection_chain.invoke(
                 dict(
@@ -119,8 +120,117 @@ class ReflectionRefiner(JanusRefiner):
             if re.search(r"\bLGTM\b", reflection) is not None:
                 return self.parser.parse(completion)
             if not retry_number:
-                log.info(f"Completion:\n{completion}")
-            log.info(f"Reflection:\n{reflection}")
+                log.debug(f"Completion:\n{completion}")
+            log.debug(f"Reflection:\n{reflection}")
+            completion = self.revision_chain.invoke(
+                dict(
+                    prompt=prompt_value.to_string(),
+                    completion=completion,
+                    reflection=reflection,
+                )
+            )
+            log.debug(f"Revision:\n{completion}")
+
+        return self.parser.parse(completion)
+
+
+class RequirementsReflectionRefiner(JanusRefiner):
+    """
+    This requirements-specific refiner is intended to address a common issue with
+    requirements reflection, where over the course of several reflection loops,
+    requirements lists grow increasingly verbose, eventually becoming too wordy
+    to be useful. To reduce this, this refiner interlaces an additional reflection
+    -> revision loop which de-duplicates requirements.
+    """
+
+    max_retries: int
+    reflection_chain: RunnableSerializable
+    revision_chain: RunnableSerializable
+    reflect_duplication_chain: RunnableSerializable
+    revise_duplication_chain: RunnableSerializable
+    reflection_prompt_name: str
+
+    def __init__(
+        self,
+        llm: JanusModel,
+        parser: JanusParser,
+        max_retries: int,
+        prompt_template_name: str = "refinement/reflection/incose",
+    ):
+        # Main reflection loop
+        reflection_prompt = MODEL_PROMPT_ENGINES[llm.short_model_id](
+            source_language="text",
+            prompt_template="refinement/reflection/incose",
+        ).prompt
+        revision_prompt = MODEL_PROMPT_ENGINES[llm.short_model_id](
+            source_language="text",
+            prompt_template="refinement/revision/incose",
+        ).prompt
+        # De-duplication loop
+        reflect_duplication_prompt = MODEL_PROMPT_ENGINES[llm.short_model_id](
+            source_language="text",
+            prompt_template="refinement/reflection/incose_deduplicate",
+        ).prompt
+        revise_duplication_prompt = MODEL_PROMPT_ENGINES[llm.short_model_id](
+            source_language="text",
+            prompt_template="refinement/revision/incose_deduplicate",
+        ).prompt
+
+        reflection_chain = reflection_prompt | llm | StrOutputParser()
+        revision_chain = revision_prompt | llm | StrOutputParser()
+        reflect_duplication_chain = reflect_duplication_prompt | llm | StrOutputParser()
+        revise_duplication_chain = revise_duplication_prompt | llm | StrOutputParser()
+        super().__init__(
+            reflection_prompt_name=prompt_template_name,
+            reflection_chain=reflection_chain,
+            revision_chain=revision_chain,
+            reflect_duplication_chain=reflect_duplication_chain,
+            revise_duplication_chain=revise_duplication_chain,
+            parser=parser,
+            max_retries=max_retries,
+        )
+
+    def parse_completion(
+        self, completion: str, prompt_value: PromptValue, **kwargs
+    ) -> Any:
+        log.debug(f"Reflection Prompt: {self.reflection_prompt_name}")
+        if isinstance(completion, AIMessage):
+            completion = completion.content
+        for retry_number in range(self.max_retries):
+            # First, check if the generated requirements are redundant or too specific
+            duplication_reflection = self.reflect_duplication_chain.invoke(
+                dict(
+                    prompt=prompt_value.to_string(),
+                    completion=completion,
+                )
+            )
+            if re.search(r"\bLGTM\b", duplication_reflection) is not None:
+                log.debug(
+                    "No de-duplication suggested in reflection, "
+                    "passing to next reflection step"
+                )
+            else:
+                completion = self.revise_duplication_chain.invoke(
+                    dict(
+                        prompt=prompt_value.to_string(),
+                        completion=completion,
+                        reflection=duplication_reflection,
+                    )
+                )
+
+            # Once we're happy with the results or trimmed them down,
+            # continue with the typical reflection process,
+            # except with specific INCOSE-focused prompts
+            reflection = self.reflection_chain.invoke(
+                dict(
+                    prompt=prompt_value.to_string(),
+                    completion=completion,
+                )
+            )
+            if re.search(r"\bLGTM\b", reflection) is not None:
+                return self.parser.parse(completion)
+            if not retry_number:
+                log.debug(f"Completion:\n{completion}")
             completion = self.revision_chain.invoke(
                 dict(
                     prompt=prompt_value.to_string(),
@@ -128,7 +238,6 @@ class ReflectionRefiner(JanusRefiner):
                     reflection=reflection,
                 )
             )
-            log.info(f"Revision:\n{completion}")
 
         return self.parser.parse(completion)