janus-llm 1.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

janus/parsers/reqs_parser.py

@@ -0,0 +1,72 @@
+import json
+import re
+
+from langchain.output_parsers.json import parse_json_markdown
+from langchain.schema.output_parser import BaseOutputParser
+from langchain_core.exceptions import OutputParserException
+from langchain_core.messages import AIMessage
+
+from ..language.block import CodeBlock
+from ..utils.logger import create_logger
+from .code_parser import JanusParser
+
+log = create_logger(__name__)
+
+
+class RequirementsParser(BaseOutputParser[str], JanusParser):
+    block_name: str = ""
+
+    def __init__(self):
+        super().__init__(expected_keys=[])
+
+    def set_reference(self, block: CodeBlock):
+        self.block_name = block.name
+
+    def parse(self, text: str) -> str:
+        if isinstance(text, AIMessage):
+            text = text.content
+        text = text.lstrip("```json")
+        text = text.rstrip("```")
+        try:
+            obj = parse_json_markdown(text)
+        except json.JSONDecodeError as e:
+            log.debug(f"Invalid JSON object. Output:\n{text}")
+            raise OutputParserException(f"Got invalid JSON object. Error: {e}")
+
+        if not isinstance(obj, dict):
+            raise OutputParserException(
+                f"Got invalid return object. Expected a dictionary, but got {type(obj)}"
+            )
+        return json.dumps(obj)
+
+    def parse_combined_output(self, text: str):
+        """Parse the output text from the LLM when multiple inputs are combined.
+
+        Arguments:
+            text: The output text from the LLM.
+
+        Returns:
+            A parsed version of the text.
+        """
+        json_strings = re.findall(r"\{.*?\}", text)
+        output_list = list()
+        for i, json_string in enumerate(json_strings, 1):
+            json_dict = json.loads(json_string)
+            output_list.append(json_dict["requirements"])
+        return output_list
+
+    def get_format_instructions(self) -> str:
+        """Get the format instructions for the parser.
+
+        Returns:
+            The format instructions for the LLM.
+        """
+        return (
+            "Output must contain an ieee style requirements specification "
+            "all in a json-formatted string, including the following field: "
+            '"requirements".'
+        )
+
+    @property
+    def _type(self) -> str:
+        return self.__class__.name

janus/prompts/prompt.py

@@ -1,15 +1,14 @@
 import json
+from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import List
 
+from langchain import PromptTemplate
 from langchain.prompts import ChatPromptTemplate
 from langchain.prompts.chat import (
     HumanMessagePromptTemplate,
     SystemMessagePromptTemplate,
 )
-from langchain.schema.messages import BaseMessage
 
-from ..language.block import CodeBlock
 from ..utils.enums import LANGUAGES
 from ..utils.logger import create_logger
 
@@ -18,12 +17,13 @@ log = create_logger(__name__)
 
 # Prompt names (self.template_map keys) that should output text,
 # regardless of the `output-lang` argument.
-TEXT_OUTPUT = ["document", "requirements"]
+TEXT_OUTPUT = []
+
 # Prompt names (self.template_map keys) that should output the
 # same language as the input, regardless of the `output-lang` argument.
 SAME_OUTPUT = ["document_inline"]
 
-JSON_OUTPUT = ["evaluate"]
+JSON_OUTPUT = ["evaluate", "document", "document_madlibs", "requirements"]
 
 # Directory containing Janus prompt template directories and files
 JANUS_PROMPT_TEMPLATES_DIR = Path(__file__).parent / "templates"
@@ -34,7 +34,7 @@ HUMAN_PROMPT_TEMPLATE_FILENAME = "human.txt"
 PROMPT_VARIABLES_FILENAME = "variables.json"
 
 
-class PromptEngine:
+class PromptEngine(ABC):
     """A class defining prompting schemes for the LLM."""
 
     def __init__(
@@ -59,22 +59,14 @@ class PromptEngine:
         template_path = self.get_prompt_template_path(prompt_template)
         self._template_path = template_path
         self._template_name = prompt_template
-        system_prompt_path = SystemMessagePromptTemplate.from_template(
-            (template_path / SYSTEM_PROMPT_TEMPLATE_FILENAME).read_text()
-        )
-        human_prompt_path = HumanMessagePromptTemplate.from_template(
-            (template_path / HUMAN_PROMPT_TEMPLATE_FILENAME).read_text()
-        )
-        self.prompt = ChatPromptTemplate.from_messages(
-            [system_prompt_path, human_prompt_path]
-        )
+        self.prompt = self.load_prompt_template(template_path)
 
         # Define variables to be passed in to the prompt formatter
         source_language = source_language.lower()
         target_language = target_language.lower()
         self.variables = dict(
-            SOURCE_LANGUAGE=source_language.lower(),
-            TARGET_LANGUAGE=target_language.lower(),
+            SOURCE_LANGUAGE=source_language,
+            TARGET_LANGUAGE=target_language,
             TARGET_LANGUAGE_VERSION=str(target_version),
             FILE_SUFFIX=LANGUAGES[source_language]["suffix"],
             SOURCE_CODE_EXAMPLE=LANGUAGES[source_language]["example"],
@@ -83,20 +75,11 @@ class PromptEngine:
         variables_path = template_path / PROMPT_VARIABLES_FILENAME
         if variables_path.exists():
             self.variables.update(json.loads(variables_path.read_text()))
+        self.prompt = self.prompt.partial(**self.variables)
 
-    def create(self, code: CodeBlock) -> List[BaseMessage]:
-        """Convert a code block to a Chat GPT prompt.
-
-        Arguments:
-            code: The code block to convert.
-
-        Returns:
-            The converted prompt as a list of messages.
-        """
-        return self.prompt.format_prompt(
-            SOURCE_CODE=code.text,
-            **self.variables,
-        ).to_messages()
+    @abstractmethod
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        pass
 
     @staticmethod
     def get_prompt_template_path(template_name: str) -> Path:
@@ -146,3 +129,93 @@ class PromptEngine:
                 f"Specified prompt template directory {template_path} is "
                 f"missing a {HUMAN_PROMPT_TEMPLATE_FILENAME}"
             )
+
+
+class ChatGptPromptEngine(PromptEngine):
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        system_prompt_path = template_path / SYSTEM_PROMPT_TEMPLATE_FILENAME
+        system_prompt = system_prompt_path.read_text()
+        system_message = SystemMessagePromptTemplate.from_template(system_prompt)
+
+        human_prompt_path = template_path / HUMAN_PROMPT_TEMPLATE_FILENAME
+        human_prompt = human_prompt_path.read_text()
+        human_message = HumanMessagePromptTemplate.from_template(human_prompt)
+        return ChatPromptTemplate.from_messages([system_message, human_message])
+
+
+class ClaudePromptEngine(PromptEngine):
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        prompt_path = template_path / HUMAN_PROMPT_TEMPLATE_FILENAME
+        prompt = prompt_path.read_text()
+        return PromptTemplate.from_template(f"Human: {prompt}\n\nAssistant: ")
+
+
+class TitanPromptEngine(PromptEngine):
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        prompt_path = template_path / HUMAN_PROMPT_TEMPLATE_FILENAME
+        prompt = prompt_path.read_text()
+        return PromptTemplate.from_template(f"User: {prompt}\n\nAssistant: ")
+
+
+class Llama2PromptEngine(PromptEngine):
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        system_prompt_path = template_path / SYSTEM_PROMPT_TEMPLATE_FILENAME
+        system_prompt = system_prompt_path.read_text()
+
+        human_prompt_path = template_path / HUMAN_PROMPT_TEMPLATE_FILENAME
+        human_prompt = human_prompt_path.read_text()
+
+        return PromptTemplate.from_template(
+            f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{human_prompt} [/INST]"
+        )
+
+
+class Llama3PromptEngine(PromptEngine):
+    # see https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-3
+    #            /#special-tokens-used-with-meta-llama-3
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        system_prompt_path = template_path / SYSTEM_PROMPT_TEMPLATE_FILENAME
+        system_prompt = system_prompt_path.read_text()
+
+        human_prompt_path = template_path / HUMAN_PROMPT_TEMPLATE_FILENAME
+        human_prompt = human_prompt_path.read_text()
+
+        return PromptTemplate.from_template(
+            f"<|begin_of_text|>"
+            f"<|start_header_id|>"
+            f"system"
+            f"<|end_header_id|>"
+            f"\n\n{system_prompt}"
+            f"<|eot_id|>"
+            f"<|start_header_id|>"
+            f"user"
+            f"<|end_header_id|>"
+            f"\n\n{human_prompt}"
+            f"<|eot_id|>"
+            f"<|start_header_id|>"
+            f"assistant"
+            f"<|end_header_id|>"
+            f"\n\n"
+        )
+
+
+class CoherePromptEngine(PromptEngine):
+    # see https://docs.cohere.com/docs/prompting-command-r
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        system_prompt_path = template_path / SYSTEM_PROMPT_TEMPLATE_FILENAME
+        system_prompt = system_prompt_path.read_text()
+
+        human_prompt_path = template_path / HUMAN_PROMPT_TEMPLATE_FILENAME
+        human_prompt = human_prompt_path.read_text()
+
+        return PromptTemplate.from_template(
+            f"<BOS_TOKEN>"
+            f"<|START_OF_TURN_TOKEN|>"
+            f"<|SYSTEM_TOKEN|>"
+            f"{system_prompt}"
+            f"<|END_OF_TURN_TOKEN|>"
+            f"<|START_OF_TURN_TOKEN|>"
+            f"<|USER_TOKEN|>"
+            f"{human_prompt}"
+            f"<|END_OF_TURN_TOKEN|>"
+        )