janus-llm 3.5.3__py3-none-any.whl → 4.1.0__py3-none-any.whl

janus/llm/models_info.py

@@ -1,15 +1,14 @@
 import json
 import os
-import time
 from pathlib import Path
-from typing import Any, Callable
+from typing import Callable, Protocol, TypeVar
 
 from dotenv import load_dotenv
 from langchain_community.llms import HuggingFaceTextGenInference
-from langchain_core.language_models import BaseLanguageModel
-from langchain_openai import ChatOpenAI
+from langchain_core.runnables import Runnable
+from langchain_openai import AzureChatOpenAI
 
-from janus.llm.model_callbacks import COST_PER_1K_TOKENS
+from janus.llm.model_callbacks import COST_PER_1K_TOKENS, azure_model_reroutes
 from janus.prompts.prompt import (
     ChatGptPromptEngine,
     ClaudePromptEngine,
@@ -44,17 +43,33 @@ except ImportError:
     )
 
 
-load_dotenv()
+ModelType = TypeVar(
+    "ModelType",
+    AzureChatOpenAI,
+    HuggingFaceTextGenInference,
+    Bedrock,
+    BedrockChat,
+    HuggingFacePipeline,
+)
 
-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",
-    "gpt-3.5-turbo": "gpt-3.5-turbo-0125",
-    "gpt-3.5-turbo-16k": "gpt-3.5-turbo-0125",
-}
+
+class JanusModelProtocol(Protocol):
+    model_id: str
+    model_type_name: str
+    token_limit: int
+    input_token_cost: float
+    output_token_cost: float
+    prompt_engine: type[PromptEngine]
+
+    def get_num_tokens(self, text: str) -> int:
+        ...
+
+
+class JanusModel(Runnable, JanusModelProtocol):
+    ...
+
+
+load_dotenv()
 
 openai_models = [
     "gpt-4o",
@@ -65,6 +80,11 @@ openai_models = [
     "gpt-3.5-turbo",
     "gpt-3.5-turbo-16k",
 ]
+azure_models = [
+    "gpt-4o",
+    "gpt-4o-mini",
+    "gpt-3.5-turbo-16k",
+]
 claude_models = [
     "bedrock-claude-v2",
     "bedrock-claude-instant-v1",
@@ -103,27 +123,21 @@ bedrock_models = [
     *cohere_models,
     *mistral_models,
 ]
-all_models = [*openai_models, *bedrock_models]
+all_models = [*azure_models, *bedrock_models]
 
-MODEL_TYPE_CONSTRUCTORS: dict[str, Callable[[Any], BaseLanguageModel]] = {
-    "OpenAI": ChatOpenAI,
+MODEL_TYPE_CONSTRUCTORS: dict[str, ModelType] = {
+    # "OpenAI": ChatOpenAI,
     "HuggingFace": HuggingFaceTextGenInference,
+    "Azure": AzureChatOpenAI,
+    "Bedrock": Bedrock,
+    "BedrockChat": BedrockChat,
+    "HuggingFaceLocal": HuggingFacePipeline,
 }
 
-try:
-    MODEL_TYPE_CONSTRUCTORS.update(
-        {
-            "HuggingFaceLocal": HuggingFacePipeline.from_model_id,
-            "Bedrock": Bedrock,
-            "BedrockChat": BedrockChat,
-        }
-    )
-except NameError:
-    pass
-
 
 MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
-    **{m: ChatGptPromptEngine for m in openai_models},
+    # **{m: ChatGptPromptEngine for m in openai_models},
+    **{m: ChatGptPromptEngine for m in azure_models},
     **{m: ClaudePromptEngine for m in claude_models},
     **{m: Llama2PromptEngine for m in llama2_models},
     **{m: Llama3PromptEngine for m in llama3_models},
@@ -132,13 +146,9 @@ MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
     **{m: MistralPromptEngine for m in mistral_models},
 }
 
-_open_ai_defaults: dict[str, str] = {
-    "openai_api_key": os.getenv("OPENAI_API_KEY"),
-    "openai_organization": os.getenv("OPENAI_ORG_ID"),
-}
-
 MODEL_ID_TO_LONG_ID = {
-    **{m: mr for m, mr in openai_model_reroutes.items()},
+    # **{m: mr for m, mr in openai_model_reroutes.items()},
+    **{m: mr for m, mr in azure_model_reroutes.items()},
     "bedrock-claude-v2": "anthropic.claude-v2",
     "bedrock-claude-instant-v1": "anthropic.claude-instant-v1",
     "bedrock-claude-haiku": "anthropic.claude-3-haiku-20240307-v1:0",
@@ -169,7 +179,8 @@ DEFAULT_MODELS = list(MODEL_DEFAULT_ARGUMENTS.keys())
 MODEL_CONFIG_DIR = Path.home().expanduser() / ".janus" / "llm"
 
 MODEL_TYPES: dict[str, PromptEngine] = {
-    **{m: "OpenAI" for m in openai_models},
+    # **{m: "OpenAI" for m in openai_models},
+    **{m: "Azure" for m in azure_models},
     **{m: "BedrockChat" for m in bedrock_models},
 }
 
@@ -179,7 +190,10 @@ TOKEN_LIMITS: dict[str, int] = {
     "gpt-4-1106-preview": 128_000,
     "gpt-4-0125-preview": 128_000,
     "gpt-4o-2024-05-13": 128_000,
+    "gpt-4o-2024-08-06": 128_000,
+    "gpt-4o-mini": 128_000,
     "gpt-3.5-turbo-0125": 16_384,
+    "gpt35-turbo-16k": 16_384,
     "text-embedding-ada-002": 8191,
     "gpt4all": 16_384,
     "anthropic.claude-v2": 100_000,
@@ -211,53 +225,100 @@ def get_available_model_names() -> list[str]:
     return avaialable_models
 
 
-def load_model(
-    user_model_name: str,
-) -> tuple[BaseLanguageModel, str, int, dict[str, float]]:
+def load_model(model_id) -> JanusModel:
     if not MODEL_CONFIG_DIR.exists():
         MODEL_CONFIG_DIR.mkdir(parents=True)
-    model_config_file = MODEL_CONFIG_DIR / f"{user_model_name}.json"
-    if not model_config_file.exists():
-        log.warning(
-            f"Model {user_model_name} not found in user-defined models, searching "
-            f"default models for {user_model_name}."
-        )
-        model_id = user_model_name
-        if user_model_name not in DEFAULT_MODELS:
-            message = (
-                f"Model {user_model_name} not found in default models. Make sure to run "
-                "`janus llm add` first."
-            )
-            log.error(message)
-            raise ValueError(message)
-        model_config = {
-            "model_type": MODEL_TYPES[model_id],
-            "model_id": model_id,
-            "model_args": MODEL_DEFAULT_ARGUMENTS[model_id],
-            "token_limit": TOKEN_LIMITS.get(MODEL_ID_TO_LONG_ID[model_id], 4096),
-            "model_cost": COST_PER_1K_TOKENS.get(
-                MODEL_ID_TO_LONG_ID[model_id], {"input": 0, "output": 0}
-            ),
-        }
-        with open(model_config_file, "w") as f:
-            json.dump(model_config, f)
-    else:
+    model_config_file = MODEL_CONFIG_DIR / f"{model_id}.json"
+
+    if model_config_file.exists():
+        log.info(f"Loading {model_id} from {model_config_file}.")
         with open(model_config_file, "r") as f:
             model_config = json.load(f)
-    model_constructor = MODEL_TYPE_CONSTRUCTORS[model_config["model_type"]]
-    model_args = model_config["model_args"]
-    if model_config["model_type"] == "OpenAI":
-        model_args.update(_open_ai_defaults)
-        log.warning("Do NOT use this model in sensitive environments!")
-        log.warning("If you would like to cancel, please press Ctrl+C.")
-        log.warning("Waiting 10 seconds...")
+        model_type_name = model_config["model_type"]
+        model_id = model_config["model_id"]
+        model_args = model_config["model_args"]
+        token_limit = model_config["token_limit"]
+        input_token_cost = model_config["model_cost"]["input"]
+        output_token_cost = model_config["model_cost"]["output"]
+
+    elif model_id in DEFAULT_MODELS:
+        model_id = model_id
+        model_long_id = MODEL_ID_TO_LONG_ID[model_id]
+        model_type_name = MODEL_TYPES[model_id]
+        model_args = MODEL_DEFAULT_ARGUMENTS[model_id]
+
+        token_limit = 0
+        input_token_cost = 0.0
+        output_token_cost = 0.0
+        if model_long_id in TOKEN_LIMITS:
+            token_limit = TOKEN_LIMITS[model_long_id]
+        if model_long_id in COST_PER_1K_TOKENS:
+            token_limits = COST_PER_1K_TOKENS[model_long_id]
+            input_token_cost = token_limits["input"]
+            output_token_cost = token_limits["output"]
+
+    else:
+        model_list = "\n\t".join(DEFAULT_MODELS)
+        message = (
+            f"Model {model_id} not found in user-defined model directory "
+            f"({MODEL_CONFIG_DIR}), and is not a default model. Valid default "
+            f"models:\n\t{model_list}\n"
+            f"To use a custom model, first run `janus llm add`."
+        )
+        log.error(message)
+        raise ValueError(message)
+
+    if model_type_name == "HuggingFaceLocal":
+        model = HuggingFacePipeline.from_model_id(
+            model_id=model_id,
+            task="text-generation",
+            model_kwargs=model_args,
+        )
+        model_args.update(pipeline=model.pipeline)
+
+    elif model_type_name == "OpenAI":
+        model_args.update(
+            openai_api_key=str(os.getenv("OPENAI_API_KEY")),
+            openai_organization=str(os.getenv("OPENAI_ORG_ID")),
+        )
+        # log.warning("Do NOT use this model in sensitive environments!")
+        # log.warning("If you would like to cancel, please press Ctrl+C.")
+        # log.warning("Waiting 10 seconds...")
         # Give enough time for the user to read the warnings and cancel
-        time.sleep(10)
-
-    model = model_constructor(**model_args)
-    return (
-        model,
-        model_config["model_id"],
-        model_config["token_limit"],
-        model_config["model_cost"],
+        # time.sleep(10)
+        raise DeprecationWarning("OpenAI models are no longer supported.")
+
+    elif model_type_name == "Azure":
+        model_args.update(
+            {
+                "api_key": os.getenv("AZURE_OPENAI_API_KEY"),
+                "azure_endpoint": os.getenv("AZURE_OPENAI_ENDPOINT"),
+                "api_version": os.getenv("OPENAI_API_VERSION", "2024-02-01"),
+            }
+        )
+
+    model_type = MODEL_TYPE_CONSTRUCTORS[model_type_name]
+    prompt_engine = MODEL_PROMPT_ENGINES[model_id]
+
+    class JanusModel(model_type):
+        model_id: str
+        short_model_id: str
+        model_type_name: str
+        token_limit: int
+        input_token_cost: float
+        output_token_cost: float
+        prompt_engine: type[PromptEngine]
+
+    model_args.update(
+        model_id=MODEL_ID_TO_LONG_ID[model_id],
+        short_model_id=model_id,
+    )
+
+    return JanusModel(
+        model_type_name=model_type_name,
+        token_limit=token_limit,
+        input_token_cost=input_token_cost,
+        output_token_cost=output_token_cost,
+        prompt_engine=prompt_engine,
+        **model_args,
     )

janus/metrics/metric.py

@@ -8,6 +8,7 @@ import typer
 from typing_extensions import Annotated
 
 from janus.llm import load_model
+from janus.llm.model_callbacks import COST_PER_1K_TOKENS
 from janus.metrics.cli import evaluate
 from janus.metrics.file_pairing import FILE_PAIRING_METHODS
 from janus.metrics.splitting import SPLITTING_METHODS
@@ -135,7 +136,7 @@ def metric(
                 **kwargs,
             ):
                 out = []
-                llm, _, token_limit, model_cost = load_model(llm_name)
+                llm = load_model(llm_name)
                 if json_file_name is not None:
                     with open(json_file_name, "r") as f:
                         json_obj = json.load(f)
@@ -171,8 +172,8 @@ def metric(
                         out_file=out_file,
                         lang=language,
                         llm=llm,
-                        token_limit=token_limit,
-                        model_cost=model_cost,
+                        token_limit=llm.token_limit,
+                        model_cost=COST_PER_1K_TOKENS[llm.model_id],
                     )
                 else:
                     raise ValueError(
@@ -187,8 +188,8 @@ def metric(
                             progress,
                             language,
                             llm,
-                            token_limit,
-                            model_cost,
+                            llm.token_limit,
+                            COST_PER_1K_TOKENS[llm.model_id],
                             *args,
                             **kwargs,
                         )
@@ -199,8 +200,8 @@ def metric(
                         progress,
                         language,
                         llm,
-                        token_limit,
-                        model_cost,
+                        llm.token_limit,
+                        COST_PER_1K_TOKENS[llm.model_id],
                         *args,
                         **kwargs,
                     )
@@ -296,7 +297,7 @@ def metric(
                 *args,
                 **kwargs,
             ):
-                llm, _, token_limit, model_cost = load_model(llm_name)
+                llm = load_model(llm_name)
                 if json_file_name is not None:
                     with open(json_file_name, "r") as f:
                         json_obj = json.load(f)
@@ -328,8 +329,8 @@ def metric(
                         out_file=out_file,
                         lang=language,
                         llm=llm,
-                        token_limit=token_limit,
-                        model_cost=model_cost,
+                        token_limit=llm.token_limit,
+                        model_cost=COST_PER_1K_TOKENS[llm.model_id],
                     )
                 else:
                     raise ValueError(
@@ -344,8 +345,8 @@ def metric(
                             progress,
                             language,
                             llm,
-                            token_limit,
-                            model_cost,
+                            llm.token_limit,
+                            COST_PER_1K_TOKENS[llm.model_id],
                             *args,
                             **kwargs,
                         )
@@ -356,8 +357,8 @@ def metric(
                         progress,
                         language,
                         llm,
-                        token_limit,
-                        model_cost,
+                        llm.token_limit,
+                        COST_PER_1K_TOKENS[llm.model_id],
                         *args,
                         **kwargs,
                     )

janus/parsers/uml.py

@@ -1,7 +1,7 @@
 import re
 import subprocess  # nosec
 from pathlib import Path
-from typing import List, Tuple
+from tempfile import NamedTemporaryFile
 
 from langchain_core.exceptions import OutputParserException
 from langchain_core.messages import BaseMessage
@@ -13,39 +13,76 @@ log = create_logger(__name__)
 
 
 class UMLSyntaxParser(CodeParser):
-    def _get_uml_output(self, file: Path) -> Tuple[str, str]:
-        # NOTE: running subprocess with shell=False, added nosec to label that we know
-        # risk exists
+    def _check_plantuml(self, text: str) -> None:
+        # Leading newlines can break the parser, remove them
+        text = text.replace("\\n", "\n").strip()
+
+        # Write the text to a temporary file (automatically deleted)
+        file = NamedTemporaryFile()
+        fname = file.name
+        with open(fname, "w") as fin:
+            fin.write(text)
+
         try:
             plantuml_path = Path.home().expanduser() / ".janus/lib/plantuml.jar"
+            # NOTE: running subprocess with shell=False, added nosec to
+            # label that we know risk exists
             res = subprocess.run(
-                ["java", "-jar", plantuml_path, file],
+                ["java", "-jar", plantuml_path, fname],
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
             )  # nosec
             stdout = res.stdout.decode("utf-8")
             stderr = res.stderr.decode("utf-8")
         except FileNotFoundError:
-            log.warning("Plant UML executable not found, skipping syntax check")
-            stdout = ""
-            stderr = ""
-        return stdout, stderr
+            err_txt = (
+                "Plant UML executable not found. Either choose a different parser"
+                " or install with `bash scripts/install_plantuml.sh`. Java and"
+                " graphviz are dependencies for the tool, they must also be installed."
+            )
+            log.error(err_txt)
+            raise Exception(err_txt)
+
+        # Check for bad outputs, raise OutputParserExceptions if so
+        if "Error" in stderr or "Error" in stdout:
+            err_txt = "Recieved UML parsing error(s)."
+
+            line_nos = self._get_error_lines(stderr) + self._get_error_lines(stdout)
+            lines = text.split("\n")
+            for i in line_nos:
+                i0 = max(0, i - 3)
+                i1 = min(len(lines) - 1, i + 2)
+                err_lines = [
+                    f"> {lines[j]}" if j == i - 1 else f"  {lines[j]}"
+                    for j in range(i0, i1)
+                ]
+                if i0:
+                    err_lines.insert(0, "  ...")
+                if i1 < (len(lines) - 1):
+                    err_lines.append("  ...")
 
-    def _get_errs(self, s: str) -> List[str]:
-        return [x.group() for x in re.finditer(r"Error (.*)\n", s)]
+                err_txt += f"\nError located at line {i} must be fixed:\n"
+                err_txt += "\n".join(err_lines)
+            log.warning(err_txt)
+            raise OutputParserException(err_txt)
+
+        if "Warning" in stdout or "Warning" in stderr:
+            err_txt = "Recieved UML parsing warning (often due to missing PLANTUML)."
+            if stderr:
+                err_txt += f"\nSTDERR:\n```\n{stderr.strip()}\n```\n"
+            if stdout:
+                err_txt += f"\nSTDOUT:\n```\n{stdout.strip()}\n```\n"
+
+            log.warning(err_txt)
+            raise OutputParserException(err_txt)
+
+    def _get_error_lines(self, s: str) -> list[int]:
+        return [int(x.group(1)) for x in re.finditer(r"Error line (\d+) in file:", s)]
+
+    def _get_warns(self, s: str) -> list[str]:
+        return [x.group() for x in re.finditer(r"Warning: (.*)\n", s)]
 
     def parse(self, text: str | BaseMessage) -> str:
         text = super().parse(text)
-        janus_path = Path.home().expanduser() / Path(".janus")
-        if not janus_path.exists():
-            janus_path.mkdir()
-        temp_file_path = janus_path / "tmp.txt"
-        with open(temp_file_path, "w") as f:
-            f.write(text)
-        uml_std_out, uml_std_err = self._get_uml_output(temp_file_path)
-        uml_errs = self._get_errs(uml_std_out) + self._get_errs(uml_std_err)
-        if len(uml_errs) > 0:
-            raise OutputParserException(
-                "Error: Received UML Errors:\n" + "\n".join(uml_errs)
-            )
+        self._check_plantuml(text)
         return text

janus/refiners/refiner.py

@@ -1,73 +1,115 @@
-from langchain_core.prompts import ChatPromptTemplate
+from typing import Any
 
-from janus.llm.models_info import MODEL_PROMPT_ENGINES
+from langchain.output_parsers import RetryWithErrorOutputParser
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompt_values import PromptValue
+from langchain_core.runnables import RunnableSerializable
 
+from janus.llm.models_info import MODEL_PROMPT_ENGINES, JanusModel
+from janus.parsers.parser import JanusParser
+from janus.utils.logger import create_logger
 
-class Refiner:
-    def refine(
-        self,
-        original_prompt: str,
-        previous_prompt: str,
-        previous_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
-        """
+log = create_logger(__name__)
+
+
+class JanusRefiner(JanusParser):
+    parser: JanusParser
+
+    def parse_runnable(self, input: dict[str, Any]) -> Any:
+        return self.parse_completion(**input)
+
+    def parse_completion(self, completion: str, **kwargs) -> Any:
+        return self.parser.parse(completion)
+
+    def parse(self, text: str) -> str:
         raise NotImplementedError
 
 
-class BasicRefiner(Refiner):
+class FixParserExceptions(JanusRefiner, RetryWithErrorOutputParser):
+    def __init__(self, llm: JanusModel, parser: JanusParser, max_retries: int):
+        retry_prompt = MODEL_PROMPT_ENGINES[llm.model_id](
+            source_language="text",
+            prompt_template="refinement/fix_exceptions",
+        ).prompt
+        chain = retry_prompt | llm | StrOutputParser()
+        RetryWithErrorOutputParser.__init__(
+            self, parser=parser, retry_chain=chain, max_retries=max_retries
+        )
+
+    def parse_completion(
+        self, completion: str, prompt_value: PromptValue, **kwargs
+    ) -> Any:
+        return self.parse_with_prompt(completion, prompt_value=prompt_value)
+
+
+class ReflectionRefiner(JanusRefiner):
+    max_retries: int
+    reflection_chain: RunnableSerializable
+    revision_chain: RunnableSerializable
+
     def __init__(
         self,
-        prompt_name: str,
-        model_id: 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_id: ID of the llm to use. Found in models_info.py
-            source_language: source_langauge to use
-        """
-        self._prompt_name = prompt_name
-        self._model_id = model_id
-        self._source_language = source_language
-
-    def refine(
-        self,
-        original_prompt: str,
-        previous_prompt: str,
-        previous_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_id](
-            prompt_template=self._prompt_name,
-            source_language=self._source_language,
+        llm: JanusModel,
+        parser: JanusParser,
+        max_retries: int,
+        prompt_template_name: str = "refinement/reflection",
+    ):
+        reflection_prompt = MODEL_PROMPT_ENGINES[llm.model_id](
+            source_language="text",
+            prompt_template=prompt_template_name,
+        ).prompt
+        revision_prompt = MODEL_PROMPT_ENGINES[llm.model_id](
+            source_language="text",
+            prompt_template="refinement/revision",
+        ).prompt
+
+        reflection_chain = reflection_prompt | llm | StrOutputParser()
+        revision_chain = revision_prompt | llm | StrOutputParser()
+        super().__init__(
+            reflection_chain=reflection_chain,
+            revision_chain=revision_chain,
+            parser=parser,
+            max_retries=max_retries,
+        )
+
+    def parse_completion(
+        self, completion: str, prompt_value: PromptValue, **kwargs
+    ) -> Any:
+        for retry_number in range(self.max_retries):
+            reflection = self.reflection_chain.invoke(
+                dict(
+                    prompt=prompt_value.to_string(),
+                    completion=completion,
+                )
+            )
+            if reflection.strip() == "LGTM":
+                return self.parser.parse(completion)
+            if not retry_number:
+                log.info(f"Completion:\n{completion}")
+            log.info(f"Reflection:\n{reflection}")
+            completion = self.revision_chain.invoke(
+                dict(
+                    prompt=prompt_value.to_string(),
+                    completion=completion,
+                    reflection=reflection,
+                )
+            )
+            log.info(f"Revision:\n{completion}")
+
+        return self.parser.parse(completion)
+
+
+class HallucinationRefiner(ReflectionRefiner):
+    def __init__(self, **kwargs):
+        super().__init__(
+            prompt_template_name="refinement/hallucination",
+            **kwargs,
         )
-        prompt_arguments = {
-            "ORIGINAL_PROMPT": original_prompt,
-            "OUTPUT": previous_output,
-            "ERRORS": errors,
-        }
-        return prompt_engine.prompt, prompt_arguments
+
+
+REFINERS = dict(
+    none=JanusRefiner,
+    parser=FixParserExceptions,
+    reflection=ReflectionRefiner,
+    hallucination=HallucinationRefiner,
+)