janus-llm 4.0.0__tar.gz → 4.1.0__tar.gz

{janus_llm-4.0.0 → janus_llm-4.1.0}/PKG-INFO

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

{janus_llm-4.0.0 → janus_llm-4.1.0}/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__ = "4.0.0"
+__version__ = "4.1.0"
 
 # Ignoring a deprecation warning from langchain_core that I can't seem to hunt down
 warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

{janus_llm-4.0.0 → janus_llm-4.1.0}/janus/cli.py

@@ -39,6 +39,7 @@ from janus.llm.models_info import (
     MODEL_TYPE_CONSTRUCTORS,
     MODEL_TYPES,
     TOKEN_LIMITS,
+    azure_models,
     bedrock_models,
     openai_models,
 )
@@ -952,7 +953,7 @@ def llm_add(
             help="The type of the model",
             click_type=click.Choice(sorted(list(MODEL_TYPE_CONSTRUCTORS.keys()))),
         ),
-    ] = "OpenAI",
+    ] = "Azure",
 ):
     if not MODEL_CONFIG_DIR.exists():
         MODEL_CONFIG_DIR.mkdir(parents=True)
@@ -996,6 +997,7 @@ def llm_add(
             "model_cost": {"input": in_cost, "output": out_cost},
         }
     elif model_type == "OpenAI":
+        print("DEPRECATED: Use 'Azure' instead. CTRL+C to exit.")
         model_id = typer.prompt(
             "Enter the model ID (list model IDs with `janus llm ls -a`)",
             default="gpt-4o",
@@ -1017,6 +1019,28 @@ def llm_add(
             "token_limit": max_tokens,
             "model_cost": model_cost,
         }
+    elif model_type == "Azure":
+        model_id = typer.prompt(
+            "Enter the model ID (list model IDs with `janus llm ls -a`)",
+            default="gpt-4o",
+            type=click.Choice(azure_models),
+            show_choices=False,
+        )
+        params = dict(
+            # Azure uses the "azure_deployment" key for what we're calling "long_model_id"
+            azure_deployment=MODEL_ID_TO_LONG_ID[model_id],
+            temperature=0.7,
+            n=1,
+        )
+        max_tokens = TOKEN_LIMITS[MODEL_ID_TO_LONG_ID[model_id]]
+        model_cost = COST_PER_1K_TOKENS[MODEL_ID_TO_LONG_ID[model_id]]
+        cfg = {
+            "model_type": model_type,
+            "model_id": model_id,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": model_cost,
+        }
     elif model_type == "BedrockChat" or model_type == "Bedrock":
         model_id = typer.prompt(
             "Enter the model ID (list model IDs with `janus llm ls -a`)",

{janus_llm-4.0.0 → janus_llm-4.1.0}/janus/converter/_tests/test_translate.py

@@ -90,14 +90,14 @@ class TestDiagramGenerator(unittest.TestCase):
     def setUp(self):
         """Set up the tests."""
         self.diagram_generator = DiagramGenerator(
-            model="gpt-4o",
+            model="gpt-4o-mini",
             source_language="fortran",
             diagram_type="Activity",
         )
 
     def test_init(self):
         """Test __init__ method."""
-        self.assertEqual(self.diagram_generator._model_name, "gpt-4o")
+        self.assertEqual(self.diagram_generator._model_name, "gpt-4o-mini")
         self.assertEqual(self.diagram_generator._source_language, "fortran")
         self.assertEqual(self.diagram_generator._diagram_type, "Activity")
 

{janus_llm-4.0.0 → janus_llm-4.1.0}/janus/llm/model_callbacks.py

@@ -23,6 +23,11 @@ openai_model_reroutes = {
     "gpt-3.5-turbo-16k-0613": "gpt-3.5-turbo-0125",
 }
 
+azure_model_reroutes = {
+    "gpt-4o": "gpt-4o-2024-08-06",
+    "gpt-4o-mini": "gpt-4o-mini",
+    "gpt-3.5-turbo-16k": "gpt35-turbo-16k",
+}
 
 # Updated 2024-06-21
 COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
@@ -31,6 +36,10 @@ COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
     "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
     "gpt-4-0613": {"input": 0.03, "output": 0.06},
     "gpt-4o-2024-05-13": {"input": 0.005, "output": 0.015},
+    "gpt-4o-2024-08-06": {"input": 0.00275, "output": 0.011},
+    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
+    "gpt35-turbo-16k": {"input": 0.003, "output": 0.004},
+    "gpt-35-turbo-16k": {"input": 0.003, "output": 0.004},
     "anthropic.claude-v2": {"input": 0.008, "output": 0.024},
     "anthropic.claude-instant-v1": {"input": 0.0008, "output": 0.0024},
     "anthropic.claude-3-haiku-20240307-v1:0": {"input": 0.00025, "output": 0.00125},

{janus_llm-4.0.0 → janus_llm-4.1.0}/janus/llm/models_info.py

@@ -1,15 +1,14 @@
 import json
 import os
-import time
 from pathlib import Path
-from typing import Protocol, TypeVar
+from typing import Callable, Protocol, TypeVar
 
 from dotenv import load_dotenv
 from langchain_community.llms import HuggingFaceTextGenInference
 from langchain_core.runnables import Runnable
-from langchain_openai import ChatOpenAI
+from langchain_openai import AzureChatOpenAI
 
-from janus.llm.model_callbacks import COST_PER_1K_TOKENS, openai_model_reroutes
+from janus.llm.model_callbacks import COST_PER_1K_TOKENS, azure_model_reroutes
 from janus.prompts.prompt import (
     ChatGptPromptEngine,
     ClaudePromptEngine,
@@ -46,7 +45,7 @@ except ImportError:
 
 ModelType = TypeVar(
     "ModelType",
-    ChatOpenAI,
+    AzureChatOpenAI,
     HuggingFaceTextGenInference,
     Bedrock,
     BedrockChat,
@@ -72,7 +71,6 @@ class JanusModel(Runnable, JanusModelProtocol):
 
 load_dotenv()
 
-
 openai_models = [
     "gpt-4o",
     "gpt-4o-mini",
@@ -82,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",
@@ -120,18 +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, ModelType] = {
-    "OpenAI": ChatOpenAI,
+    # "OpenAI": ChatOpenAI,
     "HuggingFace": HuggingFaceTextGenInference,
+    "Azure": AzureChatOpenAI,
     "Bedrock": Bedrock,
     "BedrockChat": BedrockChat,
     "HuggingFaceLocal": HuggingFacePipeline,
 }
 
-MODEL_PROMPT_ENGINES: dict[str, type[PromptEngine]] = {
-    **{m: ChatGptPromptEngine for m in openai_models},
+
+MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
+    # **{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},
@@ -141,7 +147,8 @@ MODEL_PROMPT_ENGINES: dict[str, type[PromptEngine]] = {
 }
 
 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",
@@ -171,8 +178,9 @@ DEFAULT_MODELS = list(MODEL_DEFAULT_ARGUMENTS.keys())
 
 MODEL_CONFIG_DIR = Path.home().expanduser() / ".janus" / "llm"
 
-MODEL_TYPES: dict[str, str] = {
-    **{m: "OpenAI" for m in openai_models},
+MODEL_TYPES: dict[str, PromptEngine] = {
+    # **{m: "OpenAI" for m in openai_models},
+    **{m: "Azure" for m in azure_models},
     **{m: "BedrockChat" for m in bedrock_models},
 }
 
@@ -182,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,
@@ -270,11 +281,21 @@ def load_model(model_id) -> JanusModel:
             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...")
+        # 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)
+        # 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]

{janus_llm-4.0.0 → janus_llm-4.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "janus-llm"
-version = "4.0.0"
+version = "4.1.0"
 description = "A transcoding library using LLMs."
 authors = ["Michael Doyle <mdoyle@mitre.org>", "Chris Glasz <cglasz@mitre.org>",
            "Chris Tohline <ctohline@mitre.org>", "William Macke <wmacke@mitre.org>",