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

janus/language/alc/alc.py

@@ -1,10 +1,10 @@
 from langchain.schema.language_model import BaseLanguageModel
 
-from ...utils.logger import create_logger
-from ..block import CodeBlock
-from ..combine import Combiner
-from ..node import NodeType
-from ..treesitter import TreeSitterSplitter
+from janus.language.block import CodeBlock
+from janus.language.combine import Combiner
+from janus.language.node import NodeType
+from janus.language.treesitter import TreeSitterSplitter
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/binary/_tests/test_binary.py

@@ -5,17 +5,17 @@ from unittest.mock import patch
 
 import pytest
 
-from ....llm import load_model
-from ..binary import BinarySplitter
+from janus.language.binary import BinarySplitter
+from janus.llm import load_model
 
 
 class TestBinarySplitter(unittest.TestCase):
     """Tests for the BinarySplitter class."""
 
     def setUp(self):
-        model_name = "gpt-3.5-turbo-0125"
+        model_name = "gpt-4o"
         self.binary_file = Path("janus/language/binary/_tests/hello")
-        self.llm, _, _ = load_model(model_name)
+        self.llm, _, _, _ = load_model(model_name)
         self.splitter = BinarySplitter(model=self.llm)
         os.environ["GHIDRA_INSTALL_PATH"] = "~/programs/ghidra_10.4_PUBLIC"
 

janus/language/binary/binary.py

@@ -7,11 +7,11 @@ from pathlib import Path
 import tree_sitter
 from langchain.schema.language_model import BaseLanguageModel
 
-from ...utils.enums import LANGUAGES
-from ...utils.logger import create_logger
-from ..block import CodeBlock
-from ..combine import Combiner
-from ..treesitter import TreeSitterSplitter
+from janus.language.block import CodeBlock
+from janus.language.combine import Combiner
+from janus.language.treesitter import TreeSitterSplitter
+from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/block.py

@@ -1,8 +1,8 @@
 from functools import total_ordering
 from typing import ForwardRef, Hashable, Optional, Tuple
 
-from ..utils.logger import create_logger
-from .node import NodeType
+from janus.language.node import NodeType
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/combine.py

@@ -1,6 +1,6 @@
-from ..utils.logger import create_logger
-from .block import CodeBlock, TranslatedCodeBlock
-from .file import FileManager
+from janus.language.block import CodeBlock, TranslatedCodeBlock
+from janus.language.file import FileManager
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/file.py

@@ -1,5 +1,5 @@
-from ..utils.enums import LANGUAGES
-from ..utils.logger import create_logger
+from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/mumps/_tests/test_mumps.py

@@ -1,9 +1,9 @@
 import unittest
 from pathlib import Path
 
-from ....llm import load_model
-from ...combine import Combiner
-from ..mumps import MumpsSplitter
+from janus.language.combine import Combiner
+from janus.language.mumps import MumpsSplitter
+from janus.llm import load_model
 
 
 class TestMumpsSplitter(unittest.TestCase):
@@ -11,8 +11,8 @@ class TestMumpsSplitter(unittest.TestCase):
 
     def setUp(self):
         """Set up the tests."""
-        model_name = "gpt-3.5-turbo-0125"
-        llm, _, _ = load_model(model_name)
+        model_name = "gpt-4o"
+        llm, _, _, _ = load_model(model_name)
         self.splitter = MumpsSplitter(model=llm)
         self.combiner = Combiner(language="mumps")
         self.test_file = Path("janus/language/mumps/_tests/mumps.m")

janus/language/mumps/mumps.py

@@ -2,11 +2,11 @@ import re
 
 from langchain.schema.language_model import BaseLanguageModel
 
-from ...utils.logger import create_logger
-from ..block import CodeBlock
-from ..combine import Combiner
-from ..node import NodeType
-from ..splitter import Splitter
+from janus.language.block import CodeBlock
+from janus.language.combine import Combiner
+from janus.language.node import NodeType
+from janus.language.splitter import Splitter
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/mumps/patterns.py

@@ -1,7 +1,7 @@
 import re
 from dataclasses import dataclass
 
-from ...utils.logger import create_logger
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/naive/__init__.py

@@ -1,4 +1,4 @@
-from .basic_splitter import FileSplitter
-from .chunk_splitter import ChunkSplitter
-from .simple_ast import get_flexible_ast, get_strict_ast
-from .tag_splitter import TagSplitter
+from janus.language.naive.basic_splitter import FileSplitter
+from janus.language.naive.chunk_splitter import ChunkSplitter
+from janus.language.naive.simple_ast import get_flexible_ast, get_strict_ast
+from janus.language.naive.tag_splitter import TagSplitter

janus/language/naive/basic_splitter.py

@@ -1,7 +1,7 @@
-from ..block import CodeBlock
-from ..naive.chunk_splitter import ChunkSplitter
-from ..naive.registry import register_splitter
-from ..splitter import FileSizeError
+from janus.language.block import CodeBlock
+from janus.language.naive.chunk_splitter import ChunkSplitter
+from janus.language.naive.registry import register_splitter
+from janus.language.splitter import FileSizeError
 
 
 @register_splitter("file")

janus/language/naive/chunk_splitter.py

@@ -1,7 +1,7 @@
-from ..block import CodeBlock
-from ..node import NodeType
-from ..splitter import Splitter
-from .registry import register_splitter
+from janus.language.block import CodeBlock
+from janus.language.naive.registry import register_splitter
+from janus.language.node import NodeType
+from janus.language.splitter import Splitter
 
 
 @register_splitter("chunk")

janus/language/naive/registry.py

@@ -1,6 +1,6 @@
 from typing import Callable, Dict
 
-from ..splitter import Splitter
+from janus.language.splitter import Splitter
 
 CUSTOM_SPLITTERS: Dict[str, Callable[..., Splitter]] = dict()
 

janus/language/naive/simple_ast.py

@@ -1,8 +1,8 @@
-from ...utils.enums import LANGUAGES
-from ..alc.alc import AlcSplitter
-from ..mumps.mumps import MumpsSplitter
-from ..treesitter import TreeSitterSplitter
-from .registry import register_splitter
+from janus.language.alc.alc import AlcSplitter
+from janus.language.mumps.mumps import MumpsSplitter
+from janus.language.naive.registry import register_splitter
+from janus.language.treesitter import TreeSitterSplitter
+from janus.utils.enums import LANGUAGES
 
 
 @register_splitter("ast-flex")

janus/language/naive/tag_splitter.py

@@ -1,7 +1,7 @@
-from ..block import CodeBlock
-from ..node import NodeType
-from ..splitter import Splitter
-from .registry import register_splitter
+from janus.language.block import CodeBlock
+from janus.language.naive.registry import register_splitter
+from janus.language.node import NodeType
+from janus.language.splitter import Splitter
 
 
 @register_splitter("tag")

janus/language/node.py

@@ -1,6 +1,6 @@
 from typing import NewType, Tuple
 
-from ..utils.logger import create_logger
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/splitter.py

@@ -6,10 +6,10 @@ from typing import List
 import tiktoken
 from langchain.schema.language_model import BaseLanguageModel
 
-from ..utils.logger import create_logger
-from .block import CodeBlock
-from .file import FileManager
-from .node import NodeType
+from janus.language.block import CodeBlock
+from janus.language.file import FileManager
+from janus.language.node import NodeType
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/language/treesitter/_tests/test_treesitter.py

@@ -1,9 +1,9 @@
 import unittest
 from pathlib import Path
 
-from ....llm import load_model
-from ...combine import Combiner
-from ..treesitter import TreeSitterSplitter
+from janus.language.combine import Combiner
+from janus.language.treesitter import TreeSitterSplitter
+from janus.llm import load_model
 
 
 class TestTreeSitterSplitter(unittest.TestCase):
@@ -11,9 +11,9 @@ class TestTreeSitterSplitter(unittest.TestCase):
 
     def setUp(self):
         """Set up the tests."""
-        model_name = "gpt-3.5-turbo-0125"
+        model_name = "gpt-4o"
         self.maxDiff = None
-        self.llm, _, _ = load_model(model_name)
+        self.llm, _, _, _ = load_model(model_name)
 
     def _split(self):
         """Split the test file."""

janus/language/treesitter/treesitter.py

@@ -9,10 +9,10 @@ import tree_sitter
 from git import Repo
 from langchain.schema.language_model import BaseLanguageModel
 
-from ...utils.enums import LANGUAGES
-from ...utils.logger import create_logger
-from ..block import CodeBlock, NodeType
-from ..splitter import Splitter
+from janus.language.block import CodeBlock, NodeType
+from janus.language.splitter import Splitter
+from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/llm/__init__.py

@@ -1 +1 @@
-from .models_info import load_model
+from janus.llm.models_info import load_model

janus/llm/model_callbacks.py

@@ -8,7 +8,7 @@ from langchain_core.messages import AIMessage
 from langchain_core.outputs import ChatGeneration, LLMResult
 from langchain_core.tracers.context import register_configure_hook
 
-from ..utils.logger import create_logger
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/llm/models_info.py

@@ -8,7 +8,8 @@ from langchain_community.llms import HuggingFaceTextGenInference
 from langchain_core.language_models import BaseLanguageModel
 from langchain_openai import ChatOpenAI
 
-from ..prompts.prompt import (
+from janus.llm.model_callbacks import COST_PER_1K_TOKENS
+from janus.prompts.prompt import (
     ChatGptPromptEngine,
     ClaudePromptEngine,
     CoherePromptEngine,
@@ -18,8 +19,7 @@ from ..prompts.prompt import (
     PromptEngine,
     TitanPromptEngine,
 )
-from ..utils.logger import create_logger
-from .model_callbacks import COST_PER_1K_TOKENS
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 
@@ -55,11 +55,12 @@ openai_model_reroutes = {
 }
 
 openai_models = [
-    "gpt-4-0613",
-    "gpt-4-1106-preview",
-    "gpt-4-0125-preview",
-    "gpt-4o-2024-05-13",
-    "gpt-3.5-turbo-0125",
+    "gpt-4o",
+    "gpt-4",
+    "gpt-4-turbo",
+    "gpt-4-turbo-preview",
+    "gpt-3.5-turbo",
+    "gpt-3.5-turbo-16k",
 ]
 claude_models = [
     "bedrock-claude-v2",
@@ -133,8 +134,8 @@ _open_ai_defaults: dict[str, str] = {
     "openai_organization": os.getenv("OPENAI_ORG_ID"),
 }
 
-model_identifiers = {
-    **{m: m for m in openai_models},
+MODEL_ID_TO_LONG_ID = {
+    **{m: mr for m, mr in openai_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",
@@ -157,7 +158,7 @@ model_identifiers = {
 
 MODEL_DEFAULT_ARGUMENTS: dict[str, dict[str, str]] = {
     k: (dict(model_name=k) if k in openai_models else dict(model_id=v))
-    for k, v in model_identifiers.items()
+    for k, v in MODEL_ID_TO_LONG_ID.items()
 }
 
 DEFAULT_MODELS = list(MODEL_DEFAULT_ARGUMENTS.keys())
@@ -199,22 +200,38 @@ TOKEN_LIMITS: dict[str, int] = {
 }
 
 
-def load_model(model_name: str) -> tuple[BaseLanguageModel, int, dict[str, float]]:
+def get_available_model_names() -> list[str]:
+    avaialable_models = []
+    for file in MODEL_CONFIG_DIR.iterdir():
+        if file.is_file():
+            avaialable_models.append(MODEL_CONFIG_DIR.stem)
+    return avaialable_models
+
+
+def load_model(user_model_name: str) -> tuple[BaseLanguageModel, int, dict[str, float]]:
     if not MODEL_CONFIG_DIR.exists():
         MODEL_CONFIG_DIR.mkdir(parents=True)
-    model_config_file = MODEL_CONFIG_DIR / f"{model_name}.json"
+    model_config_file = MODEL_CONFIG_DIR / f"{user_model_name}.json"
     if not model_config_file.exists():
-        if model_name not in DEFAULT_MODELS:
-            if model_name in openai_model_reroutes:
-                model_name = openai_model_reroutes[model_name]
-            else:
-                raise ValueError(f"Error: could not find model {model_name}")
+        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_name],
-            "model_args": MODEL_DEFAULT_ARGUMENTS[model_name],
-            "token_limit": TOKEN_LIMITS.get(model_identifiers[model_name], 4096),
+            "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_identifiers[model_name], {"input": 0, "output": 0}
+                MODEL_ID_TO_LONG_ID[model_id], {"input": 0, "output": 0}
             ),
         }
         with open(model_config_file, "w") as f:
@@ -227,4 +244,9 @@ def load_model(model_name: str) -> tuple[BaseLanguageModel, int, dict[str, float
     if model_config["model_type"] == "OpenAI":
         model_args.update(_open_ai_defaults)
     model = model_constructor(**model_args)
-    return model, model_config["token_limit"], model_config["model_cost"]
+    return (
+        model,
+        model_config["model_id"],
+        model_config["token_limit"],
+        model_config["model_cost"],
+    )

janus/metrics/_tests/test_bleu.py

@@ -2,7 +2,7 @@ import unittest
 
 from sacrebleu import sentence_bleu
 
-from ..bleu import bleu
+from janus.metrics.bleu import bleu
 
 
 class TestBLEU(unittest.TestCase):

janus/metrics/_tests/test_chrf.py

@@ -2,7 +2,7 @@ import unittest
 
 from sacrebleu import sentence_chrf
 
-from ..chrf import chrf
+from janus.metrics.chrf import chrf
 
 
 class TestChrF(unittest.TestCase):

janus/metrics/_tests/test_file_pairing.py

@@ -3,7 +3,7 @@
 import unittest
 from pathlib import Path
 
-from ..file_pairing import (
+from janus.metrics.file_pairing import (
     FILE_PAIRING_METHODS,
     pair_by_file,
     pair_by_line,

janus/metrics/_tests/test_llm.py

@@ -3,8 +3,8 @@ from unittest.mock import patch
 
 import pytest
 
-from ...llm.models_info import load_model
-from ..llm_metrics import llm_evaluate_option, llm_evaluate_ref_option
+from janus.llm.models_info import load_model
+from janus.metrics.llm_metrics import llm_evaluate_option, llm_evaluate_ref_option
 
 
 class TestLLMMetrics(unittest.TestCase):
@@ -53,7 +53,7 @@ class TestLLMMetrics(unittest.TestCase):
             self.bad_code,
             metric="quality",
             language="python",
-            llm=load_model("gpt-3.5-turbo-0125")[0],
+            llm=load_model("gpt-4o")[0],
         )
         self.assertLess(bad_code_quality, 5)
 
@@ -63,7 +63,7 @@ class TestLLMMetrics(unittest.TestCase):
             self.impressive_code,
             metric="quality",
             language="python",
-            llm=load_model("gpt-3.5-turbo-0125")[0],
+            llm=load_model("gpt-4o")[0],
         )
         self.assertGreater(impressive_code_quality, 5)
 
@@ -81,7 +81,7 @@ class TestLLMMetrics(unittest.TestCase):
             self.impressive_code_reference,
             metric="faithfulness",
             language="python",
-            llm=load_model("gpt-3.5-turbo-0125")[0],
+            llm=load_model("gpt-4o")[0],
         )
         self.assertGreater(faithfulness, 8)
 

janus/metrics/_tests/test_reading.py

@@ -1,6 +1,6 @@
 import unittest
 
-from ..reading import _repeat_text, flesch, gunning_fog
+from janus.metrics.reading import _repeat_text, flesch, gunning_fog
 
 
 class TestReading(unittest.TestCase):

janus/metrics/_tests/test_rouge_score.py

@@ -1,6 +1,6 @@
 import unittest
 
-from ..rouge_score import rouge
+from janus.metrics.rouge_score import rouge
 
 
 class TestRouge(unittest.TestCase):

janus/metrics/_tests/test_similarity_score.py

@@ -1,6 +1,6 @@
 import unittest
 
-from ..similarity import similarity_score
+from janus.metrics.similarity import similarity_score
 
 
 class TestSimilarityScore(unittest.TestCase):

janus/metrics/_tests/test_treesitter_metrics.py

@@ -3,8 +3,8 @@ from pathlib import Path
 
 from typer.testing import CliRunner
 
-from ...cli import app
-from ..complexity_metrics import (
+from janus.cli import app
+from janus.metrics.complexity_metrics import (
     TreeSitterMetric,
     cyclomatic_complexity,
     difficulty,

janus/metrics/bleu.py

@@ -4,7 +4,7 @@ import click
 import typer
 from sacrebleu import sentence_bleu
 
-from .metric import metric
+from janus.metrics.metric import metric
 
 # from sacrebleu import sentence_chrf
 

janus/metrics/chrf.py

@@ -1,7 +1,7 @@
 import typer
 from sacrebleu import sentence_chrf
 
-from .metric import metric
+from janus.metrics.metric import metric
 
 
 @metric(help="chrF score using sacrebleu")

janus/metrics/complexity_metrics.py

@@ -1,10 +1,10 @@
 import math
 from typing import List, Optional
 
-from ..language.block import CodeBlock
-from ..language.treesitter.treesitter import TreeSitterSplitter
-from ..utils.enums import LANGUAGES
-from .metric import metric
+from janus.language.block import CodeBlock
+from janus.language.treesitter import TreeSitterSplitter
+from janus.metrics.metric import metric
+from janus.utils.enums import LANGUAGES
 
 
 class NodeException(Exception):

janus/metrics/file_pairing.py

@@ -1,10 +1,10 @@
 from typing import Any, Callable
 
-from ..language.binary import BinarySplitter
-from ..language.mumps import MumpsSplitter
-from ..language.node import NodeType
-from ..language.treesitter import TreeSitterSplitter
-from ..utils.enums import CUSTOM_SPLITTERS
+from janus.language.binary import BinarySplitter
+from janus.language.mumps import MumpsSplitter
+from janus.language.node import NodeType
+from janus.language.treesitter import TreeSitterSplitter
+from janus.utils.enums import CUSTOM_SPLITTERS
 
 FILE_PAIRING_METHODS: dict[str, Callable[[str, str], list[tuple[str, str]]]] = {}
 

janus/metrics/llm_metrics.py

@@ -9,7 +9,7 @@ from langchain_core.prompts import PromptTemplate
 from langchain_core.pydantic_v1 import BaseModel, Field
 from typing_extensions import Annotated
 
-from .metric import metric
+from janus.metrics.metric import metric
 
 
 class LLMMetricOutput(BaseModel):

janus/metrics/metric.py

@@ -7,13 +7,13 @@ import click
 import typer
 from typing_extensions import Annotated
 
-from ..llm import load_model
-from ..utils.enums import LANGUAGES
-from ..utils.logger import create_logger
-from ..utils.progress import track
-from .cli import evaluate
-from .file_pairing import FILE_PAIRING_METHODS
-from .splitting import SPLITTING_METHODS
+from janus.llm import load_model
+from janus.metrics.cli import evaluate
+from janus.metrics.file_pairing import FILE_PAIRING_METHODS
+from janus.metrics.splitting import SPLITTING_METHODS
+from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
+from janus.utils.progress import track
 
 log = create_logger(__name__)
 
@@ -112,7 +112,7 @@ def metric(
                         "-L",
                         help="The custom name of the model set with 'janus llm add'.",
                     ),
-                ] = "gpt-3.5-turbo-0125",
+                ] = "gpt-4o",
                 progress: Annotated[
                     bool,
                     typer.Option(
@@ -135,7 +135,7 @@ def metric(
                 **kwargs,
             ):
                 out = []
-                llm, token_limit, model_cost = load_model(llm_name)
+                llm, _, token_limit, model_cost = load_model(llm_name)
                 if json_file_name is not None:
                     with open(json_file_name, "r") as f:
                         json_obj = json.load(f)
@@ -274,7 +274,7 @@ def metric(
                         "-L",
                         help="The custom name of the model set with 'janus llm add'.",
                     ),
-                ] = "gpt-3.5-turbo-0125",
+                ] = "gpt-4o",
                 progress: Annotated[
                     bool,
                     typer.Option(
@@ -296,7 +296,7 @@ def metric(
                 *args,
                 **kwargs,
             ):
-                llm, token_limit, model_cost = load_model(llm_name)
+                llm, _, token_limit, model_cost = load_model(llm_name)
                 if json_file_name is not None:
                     with open(json_file_name, "r") as f:
                         json_obj = json.load(f)