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

janus/__init__.py

@@ -2,10 +2,10 @@ import warnings
 
 from langchain_core._api.deprecation import LangChainDeprecationWarning
 
-from .converter.translate import Translator
-from .metrics import *  # noqa: F403
+from janus.converter.translate import Translator
+from janus.metrics import *  # noqa: F403
 
-__version__ = "3.1.1"
+__version__ = "3.2.1"
 
 # Ignoring a deprecation warning from langchain_core that I can't seem to hunt down
 warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

janus/_tests/test_cli.py

@@ -4,9 +4,9 @@ from unittest.mock import ANY, patch
 
 from typer.testing import CliRunner
 
-from ..cli import app, translate
-from ..embedding.embedding_models_info import EMBEDDING_MODEL_CONFIG_DIR
-from ..llm.models_info import MODEL_CONFIG_DIR
+from janus.cli import app, translate
+from janus.embedding.embedding_models_info import EMBEDDING_MODEL_CONFIG_DIR
+from janus.llm.models_info import MODEL_CONFIG_DIR
 
 
 class TestCli(unittest.TestCase):

janus/cli.py

@@ -32,8 +32,12 @@ from janus.language.treesitter import TreeSitterSplitter
 from janus.llm.model_callbacks import COST_PER_1K_TOKENS
 from janus.llm.models_info import (
     MODEL_CONFIG_DIR,
+    MODEL_ID_TO_LONG_ID,
     MODEL_TYPE_CONSTRUCTORS,
+    MODEL_TYPES,
     TOKEN_LIMITS,
+    bedrock_models,
+    openai_models,
 )
 from janus.metrics.cli import evaluate
 from janus.utils.enums import LANGUAGES
@@ -104,7 +108,7 @@ embedding = typer.Typer(
 
 def version_callback(value: bool) -> None:
     if value:
-        from . import __version__ as version
+        from janus import __version__ as version
 
         print(f"Janus CLI [blue]v{version}[/blue]")
         raise typer.Exit()
@@ -179,7 +183,7 @@ def translate(
             "-L",
             help="The custom name of the model set with 'janus llm add'.",
         ),
-    ] = "gpt-3.5-turbo-0125",
+    ] = "gpt-4o",
     max_prompts: Annotated[
         int,
         typer.Option(
@@ -301,7 +305,7 @@ def document(
             "-L",
             help="The custom name of the model set with 'janus llm add'.",
         ),
-    ] = "gpt-3.5-turbo-0125",
+    ] = "gpt-4o",
     max_prompts: Annotated[
         int,
         typer.Option(
@@ -437,7 +441,7 @@ def diagram(
             "-L",
             help="The custom name of the model set with 'janus llm add'.",
         ),
-    ] = "gpt-3.5-turbo-0125",
+    ] = "gpt-4o",
     max_prompts: Annotated[
         int,
         typer.Option(
@@ -800,16 +804,44 @@ def llm_add(
             "model_cost": {"input": in_cost, "output": out_cost},
         }
     elif model_type == "OpenAI":
-        model_name = typer.prompt("Enter the model name", default="gpt-3.5-turbo-0125")
+        model_id = typer.prompt(
+            "Enter the model ID (list model IDs with `janus llm ls -a`)",
+            default="gpt-4o",
+            type=click.Choice(openai_models),
+            show_choices=False,
+        )
         params = dict(
-            model_name=model_name,
+            # OpenAI uses the "model_name" key for what we're calling "long_model_id"
+            model_name=MODEL_ID_TO_LONG_ID[model_id],
             temperature=0.7,
             n=1,
         )
-        max_tokens = TOKEN_LIMITS[model_name]
-        model_cost = COST_PER_1K_TOKENS[model_name]
+        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`)",
+            default="bedrock-claude-sonnet",
+            type=click.Choice(bedrock_models),
+            show_choices=False,
+        )
+        params = dict(
+            # Bedrock uses the "model_id" key for what we're calling "long_model_id"
+            model_id=MODEL_ID_TO_LONG_ID[model_id],
+            model_kwargs={"temperature": 0.7},
+        )
+        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,
@@ -821,6 +853,31 @@ def llm_add(
     print(f"Model config written to {model_cfg}")
 
 
+@llm.command("ls", help="List all of the user-configured models")
+def llm_ls(
+    all: Annotated[
+        bool,
+        typer.Option(
+            "--all",
+            "-a",
+            is_flag=True,
+            help="List all models, including the default model IDs.",
+            click_type=click.Choice(sorted(list(MODEL_TYPE_CONSTRUCTORS.keys()))),
+        ),
+    ] = False,
+):
+    print("\n[green]User-configured models[/green]:")
+    for model_cfg in MODEL_CONFIG_DIR.glob("*.json"):
+        with open(model_cfg, "r") as f:
+            cfg = json.load(f)
+        print(f"\t[blue]{model_cfg.stem}[/blue]: [purple]{cfg['model_type']}[/purple]")
+
+    if all:
+        print("\n[green]Available model IDs[/green]:")
+        for model_id, model_type in MODEL_TYPES.items():
+            print(f"\t[blue]{model_id}[/blue]: [purple]{model_type}[/purple]")
+
+
 @embedding.command("add", help="Add an embedding model config to janus")
 def embedding_add(
     model_name: Annotated[

janus/converter/__init__.py

@@ -1,6 +1,6 @@
-from .converter import Converter
-from .diagram import DiagramGenerator
-from .document import Documenter, MadLibsDocumenter, MultiDocumenter
-from .evaluate import Evaluator
-from .requirements import RequirementsDocumenter
-from .translate import Translator
+from janus.converter.converter import Converter
+from janus.converter.diagram import DiagramGenerator
+from janus.converter.document import Documenter, MadLibsDocumenter, MultiDocumenter
+from janus.converter.evaluate import Evaluator
+from janus.converter.requirements import RequirementsDocumenter
+from janus.converter.translate import Translator

janus/converter/_tests/test_translate.py

@@ -7,37 +7,11 @@ from langchain.schema import Document
 from langchain.schema.embeddings import Embeddings
 from langchain.schema.vectorstore import VST, VectorStore
 
+from janus.converter.diagram import DiagramGenerator
+from janus.converter.requirements import RequirementsDocumenter
+from janus.converter.translate import Translator
 from janus.language.block import CodeBlock, TranslatedCodeBlock
 
-from ..diagram import DiagramGenerator
-from ..requirements import RequirementsDocumenter
-from ..translate import Translator
-
-# from langchain.vectorstores import Chroma
-
-
-# from ..utils.enums import EmbeddingType
-
-
-def print_query_results(query, n_results):
-    # print(f"\n{query}")
-    # count = 1
-    # for t in n_results:
-    #     short_code = (
-    #         (t[0].page_content[0:50] + "..")
-    #         if (len(t[0].page_content) > 50)
-    #         else t[0].page_content
-    #     )
-    #     return_index = short_code.find("\n")
-    #     if -1 != return_index:
-    #         short_code = short_code[0:return_index] + ".."
-    #     print(
-    #         f"{count}. @ {t[0].metadata['start_line']}-{t[0].metadata['end_line']}"
-    #         f" -- {t[1]} -- {short_code}"
-    #     )
-    #     count += 1
-    pass
-
 
 class MockCollection(VectorStore):
     """Vector store for testing"""
@@ -65,21 +39,13 @@ class MockCollection(VectorStore):
         raise NotImplementedError("from_texts() not implemented!")
 
 
-# class MockEmbeddingsFactory(EmbeddingsFactory):
-#     """Embeddings for testing - uses MockCollection"""
-#
-#     def get_embeddings(self) -> Embeddings:
-#         return MockCollection()
-#
-
-
 class TestTranslator(unittest.TestCase):
     """Tests for the Translator class."""
 
     def setUp(self):
         """Set up the tests."""
         self.translator = Translator(
-            model="gpt-3.5-turbo-0125",
+            model="gpt-4o",
             source_language="fortran",
             target_language="python",
             target_version="3.10",
@@ -88,7 +54,7 @@ class TestTranslator(unittest.TestCase):
         self.TEST_FILE_EMBEDDING_COUNT = 14
 
         self.req_translator = RequirementsDocumenter(
-            model="gpt-3.5-turbo-0125",
+            model="gpt-4o",
             source_language="fortran",
             prompt_template="requirements",
         )
@@ -105,200 +71,6 @@ class TestTranslator(unittest.TestCase):
         # unit tests anyway
         self.assertTrue(python_file.exists())
 
-        # def test_embeddings(self):
-        #     """Testing access to embeddings"""
-        #     vector_store = self.translator.embeddings(EmbeddingType.SOURCE)
-        #     self.assertIsInstance(vector_store, Chroma, "Unexpected vector store type!")
-        #     self.assertEqual(
-        #         0, vector_store._collection.count(), "Non-empty initial vector store?"
-        #     )
-        #
-        #     self.translator.set_model("llama")
-        #     self.translator._load_parameters()
-        #     vector_store = self.translator.embeddings(EmbeddingType.SOURCE)
-        #     self.assertIsInstance(vector_store, Chroma)
-        #     self.assertEqual(
-        #         0, vector_store._collection.count(), "Non-empty initial vector store?"
-        #     )
-
-        # def test_embed_split_source(self):
-        #     """Characterize _embed method"""
-        #     mock_embeddings = MockEmbeddingsFactory()
-        #     self.translator.set_embeddings(mock_embeddings)
-        #     self.translator._load_parameters()
-        #     input_block = self.translator.splitter.split(self.test_file)
-        #     self.assertIsNone(
-        #         input_block.text, "Root node of input text shouldn't contain text"
-        #     )
-        #     self.assertIsNone(input_block.embedding_id, "Precondition failed")
-        #
-        #     result = self.translator._embed(
-        #         input_block, EmbeddingType.SOURCE, self.test_file.name
-        #     )
-        #
-        #     self.assertFalse(result, "Nothing to embed, so should have no result")
-        #     self.assertIsNone(
-        # input_block.embedding_id, "Embeddings should not have changed")
-
-        # def test_embed_has_values_for_each_non_empty_node(self):
-        #     """Characterize our sample fortran file"""
-        #     mock_embeddings = MockEmbeddingsFactory()
-        #     self.translator.set_embeddings(mock_embeddings)
-        #     self.translator._load_parameters()
-        #     input_block = self.translator.splitter.split(self.test_file)
-        #     self.translator._embed_nodes_recursively(
-        #         input_block, EmbeddingType.SOURCE, self.test_file.name
-        #     )
-        #     has_text_count = 0
-        #     has_embeddings_count = 0
-        #     nodes = [input_block]
-        #     while nodes:
-        #         node = nodes.pop(0)
-        #         if node.text:
-        #             has_text_count += 1
-        #         if node.embedding_id:
-        #             has_embeddings_count += 1
-        #         nodes.extend(node.children)
-        #     self.assertEqual(
-        #         self.TEST_FILE_EMBEDDING_COUNT,
-        #         has_text_count,
-        #         "Parsing of test_file has changed!",
-        #     )
-        #     self.assertEqual(
-        #         self.TEST_FILE_EMBEDDING_COUNT,
-        #         has_embeddings_count,
-        #         "Not all non-empty nodes have embeddings!",
-        #     )
-
-        # def test_embed_nodes_recursively(self):
-        #     mock_embeddings = MockEmbeddingsFactory()
-        #     self.translator.set_embeddings(mock_embeddings)
-        #     self.translator._load_parameters()
-        #     input_block = self.translator.splitter.split(self.test_file)
-        #     self.translator._embed_nodes_recursively(
-        #         input_block, EmbeddingType.SOURCE, self.test_file.name
-        #     )
-        #     nodes = [input_block]
-        #     while nodes:
-        #         node = nodes.pop(0)
-        #         self.assertEqual(node.text is not None, node.embedding_id is not None)
-        #         nodes.extend(node.children)
-
-        # @pytest.mark.slow
-        # def test_translate_file_adds_source_embeddings(self):
-        #     mock_embeddings = MockEmbeddingsFactory()
-        #     self.translator.set_embeddings(mock_embeddings)
-        #     self.translator._load_parameters()
-        #     vector_store = self.translator.embeddings(EmbeddingType.SOURCE)
-        #     self.assertEqual(0, vector_store._add_texts_calls, "precondition")
-        #
-        #     self.translator.translate_file(self.test_file)
-        #
-        #     self.assertEqual(
-        #         self.TEST_FILE_EMBEDDING_COUNT,
-        #         vector_store._add_texts_calls,
-        #         "Did not find expected source embeddings",
-        #     )
-
-        # @pytest.mark.slow
-        # def test_embeddings_usage(self):
-        #     """Noodling on use of embeddings
-        #     To see results have to uncomment print_query_results() above
-        #     """
-        #     input_block = self.translator.splitter.split(self.test_file)
-        #     self.translator._embed_nodes_recursively(
-        #         input_block, EmbeddingType.SOURCE, self.test_file.name
-        #     )
-        #     vector_store = self.translator.embeddings(EmbeddingType.SOURCE)
-        #
-        #     # this symbol has the lowest relevance scores of any in this test, but
-        #     # still not very low; multiple embedded nodes contain it
-        #     QUERY_STRING = "IWX_BAND_START"
-        #     query = self.translator._embeddings._embeddings.embed_query(QUERY_STRING)
-        #     n_results = vector_store.similarity_search_by_vector_with_relevance_scores(
-        #         embedding=query,
-        #         k=10,
-        #         where_document={"$contains": QUERY_STRING},
-        #     )
-        #     self.assertTrue(len(n_results) > 1, "Why was valid symbol not found?")
-        #     print_query_results(QUERY_STRING, n_results)
-
-        # in the XYZZY test, the least dissimilar results were the start and finish lines
-        # 0, and 415, which produced a similarity score of 0.47:
-
-        # QUERY_STRING = "XYZZY"
-        # query = self.translator._embeddings.embed_query(QUERY_STRING)
-        # n_results = vector_store.similarity_search_by_vector_with_relevance_scores(
-        #     embedding=query,
-        #     k=10,
-        #     # filter={"end_line": 15},
-        #     # filter={"$and": [{"end_line": 15}, {"tokens": {"$gte": 21}}]},
-        #     # where_document={"$contains": QUERY_STRING},
-        # )
-        # print_query_results(QUERY_STRING, n_results)
-        # # self.assertTrue(len(n_results) == 0, "Invalid symbol was found?")
-
-        # # only returns a single result because only 1 embedded node contains
-        # #   CSV_ICASEARR:
-        # QUERY_STRING = "What is the use of CSV_ICASEARR?"
-        # query = self.translator._embeddings._embeddings.embed_query(QUERY_STRING)
-        # n_results = vector_store.similarity_search_by_vector_with_relevance_scores(
-        #     embedding=query,
-        #     k=10,
-        #     # where_document={"$contains": QUERY_STRING},
-        #     where_document={"$contains": "CSV_ICASEARR"},
-        # )
-        # print_query_results(QUERY_STRING, n_results)
-        # self.assertTrue(len(n_results) == 1, "Was splitting changed?")
-        #
-        # # trimmed out some characters from line 43, and still not very similar scoring
-        # QUERY_STRING = "IYL_EDGEBUFFER EDGEBUFFER IGN_MASK CELLSIZE"
-        # query = self.translator._embeddings._embeddings.embed_query(QUERY_STRING)
-        # n_results = vector_store.similarity_search_by_vector_with_relevance_scores(
-        #     embedding=query,
-        #     k=10,
-        #     # where_document={"$contains": QUERY_STRING},
-        # )
-        # print_query_results(QUERY_STRING, n_results)
-        #
-        # # random string (as bad as XYZZY), but searching for a specific line
-        # QUERY_STRING = "ghost in the invisible moon"
-        # query = self.translator._embeddings._embeddings.embed_query(QUERY_STRING)
-        # n_results = vector_store.similarity_search_by_vector_with_relevance_scores(
-        #     embedding=query,
-        #     k=10,
-        #     filter={"$and": [{"end_line": 90}, {"tokens": {"$gte": 21}}]},
-        # )
-        # print_query_results(QUERY_STRING, n_results)
-        # self.assertTrue(len(n_results) == 1, "Was splitting changed?")
-
-    # @pytest.mark.slow
-    # def test_document_embeddings_added_by_translate(self):
-    #     vector_store = self.req_translator.embeddings(EmbeddingType.REQUIREMENT)
-    #     self.assertEqual(0, vector_store._add_texts_calls, "Precondition failed")
-    #     self.req_translator.translate(self.test_file.parent, self.test_file.parent,
-    #                                   True)
-    #     self.assertTrue(vector_store._add_texts_calls > 0, "Why no documentation?")
-
-    # @pytest.mark.slow
-    # def test_embed_requirements(self):
-    #     vector_store = self.req_translator.embeddings(EmbeddingType.REQUIREMENT)
-    #     translated = self.req_translator.translate_file(self.test_file)
-    #     self.assertEqual(
-    #         0,
-    #         vector_store._add_texts_calls,
-    #         "Unexpected requirements added in translate_file",
-    #     )
-    #     result = self.req_translator._embed(
-    #         translated, EmbeddingType.REQUIREMENT, self.test_file.name
-    #     )
-    #     self.assertFalse(result, "No text in root node, so should generate no docs")
-    #     self.assertIsNotNone(translated.children[0].text, "Data changed?")
-    #     result = self.req_translator._embed(
-    #         translated.children[0], EmbeddingType.REQUIREMENT, self.test_file.name
-    #     )
-    #     self.assertTrue(result, "No docs generated for first child node?")
-
     def test_invalid_selections(self) -> None:
         """Tests that settings values for the translator will raise exceptions"""
         self.assertRaises(
@@ -317,14 +89,14 @@ class TestDiagramGenerator(unittest.TestCase):
     def setUp(self):
         """Set up the tests."""
         self.diagram_generator = DiagramGenerator(
-            model="gpt-3.5-turbo-0125",
+            model="gpt-4o",
             source_language="fortran",
             diagram_type="Activity",
         )
 
     def test_init(self):
         """Test __init__ method."""
-        self.assertEqual(self.diagram_generator._model_name, "gpt-3.5-turbo-0125")
+        self.assertEqual(self.diagram_generator._model_name, "gpt-4o")
         self.assertEqual(self.diagram_generator._source_language, "fortran")
         self.assertEqual(self.diagram_generator._diagram_type, "Activity")
 
@@ -370,8 +142,8 @@ def test_language_combinations(
     """Tests that translator target language settings are consistent
     with prompt template expectations.
     """
-    translator = Translator(model="gpt-3.5-turbo-0125")
-    translator.set_model("gpt-3.5-turbo-0125")
+    translator = Translator(model="gpt-4o")
+    translator.set_model("gpt-4o")
     translator.set_source_language(source_language)
     translator.set_target_language(expected_target_language, expected_target_version)
     translator.set_prompt(prompt_template)
@@ -379,5 +151,5 @@ def test_language_combinations(
     assert translator._target_language == expected_target_language  # nosec
     assert translator._target_version == expected_target_version  # nosec
     assert translator._splitter.language == source_language  # nosec
-    assert translator._splitter.model.model_name == "gpt-3.5-turbo-0125"  # nosec
+    assert translator._splitter.model.model_name == "gpt-4o"  # nosec
     assert translator._prompt_template_name == prompt_template  # nosec

janus/converter/converter.py

@@ -64,7 +64,7 @@ class Converter:
 
     def __init__(
         self,
-        model: str = "gpt-3.5-turbo-0125",
+        model: str = "gpt-4o",
         model_arguments: dict[str, Any] = {},
         source_language: str = "fortran",
         max_prompts: int = 10,
@@ -92,6 +92,7 @@ class Converter:
         self.override_token_limit: bool = max_tokens is not None
 
         self._model_name: str
+        self._model_id: str
         self._custom_model_arguments: dict[str, Any]
 
         self._source_language: str
@@ -265,7 +266,9 @@ class Converter:
         # model_arguments.update(self._custom_model_arguments)
 
         # Load the model
-        self._llm, token_limit, self.model_cost = load_model(self._model_name)
+        self._llm, self._model_id, token_limit, self.model_cost = load_model(
+            self._model_name
+        )
         # Set the max_tokens to less than half the model's limit to allow for enough
         # tokens at output
         # Only modify max_tokens if it is not specified by user
@@ -283,7 +286,7 @@ class Converter:
         If the relevant fields have not been changed since the last time this
         method was called, nothing happens.
         """
-        prompt_engine = MODEL_PROMPT_ENGINES[self._model_name](
+        prompt_engine = MODEL_PROMPT_ENGINES[self._model_id](
             source_language=self._source_language,
             prompt_template=self._prompt_template_name,
         )

janus/converter/translate.py

@@ -90,7 +90,7 @@ class Translator(Converter):
                     f"({self._source_language} != {self._target_language})"
                 )
 
-        prompt_engine = MODEL_PROMPT_ENGINES[self._model_name](
+        prompt_engine = MODEL_PROMPT_ENGINES[self._model_id](
             source_language=self._source_language,
             target_language=self._target_language,
             target_version=self._target_version,

janus/embedding/_tests/test_collections.py

@@ -4,8 +4,8 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from ...utils.enums import EmbeddingType
-from ..collections import Collections
+from janus.embedding.collections import Collections
+from janus.utils.enums import EmbeddingType
 
 
 class TestCollections(unittest.TestCase):

janus/embedding/_tests/test_database.py

@@ -2,7 +2,7 @@ import unittest
 from pathlib import Path
 from unittest.mock import patch
 
-from ..database import ChromaEmbeddingDatabase, uri_to_path
+from janus.embedding.database import ChromaEmbeddingDatabase, uri_to_path
 
 
 class TestDatabase(unittest.TestCase):

janus/embedding/_tests/test_vectorize.py

@@ -5,9 +5,9 @@ from unittest.mock import MagicMock
 
 from chromadb.api.client import Client
 
-from ...language.treesitter import TreeSitterSplitter
-from ...utils.enums import EmbeddingType
-from ..vectorize import Vectorizer, VectorizerFactory
+from janus.embedding.vectorize import Vectorizer, VectorizerFactory
+from janus.language.treesitter import TreeSitterSplitter
+from janus.utils.enums import EmbeddingType
 
 
 class MockDBVectorizer(VectorizerFactory):

janus/embedding/collections.py

@@ -5,8 +5,8 @@ from typing import Dict, Optional, Sequence
 from chromadb import Client, Collection
 from langchain_community.vectorstores import Chroma
 
-from ..utils.enums import EmbeddingType
-from .embedding_models_info import load_embedding_model
+from janus.embedding.embedding_models_info import load_embedding_model
+from janus.utils.enums import EmbeddingType
 
 # See https://docs.trychroma.com/telemetry#in-chromas-backend-using-environment-variables
 os.environ["ANONYMIZED_TELEMETRY"] = "False"

janus/embedding/database.py

@@ -5,7 +5,7 @@ from urllib.request import url2pathname
 
 import chromadb
 
-from ..utils.logger import create_logger
+from janus.utils.logger import create_logger
 
 log = create_logger(__name__)
 

janus/embedding/embedding_models_info.py

@@ -8,7 +8,7 @@ from langchain_community.embeddings.huggingface import HuggingFaceInferenceAPIEm
 from langchain_core.embeddings import Embeddings
 from langchain_openai import OpenAIEmbeddings
 
-from ..utils.logger import create_logger
+from janus.utils.logger import create_logger
 
 load_dotenv()
 

janus/embedding/vectorize.py

@@ -6,10 +6,10 @@ from typing import Any, Dict, Optional, Sequence
 from chromadb import Client, Collection
 from langchain_community.vectorstores import Chroma
 
-from ..language.block import CodeBlock, TranslatedCodeBlock
-from ..utils.enums import EmbeddingType
-from .collections import Collections
-from .database import ChromaEmbeddingDatabase
+from janus.embedding.collections import Collections
+from janus.embedding.database import ChromaEmbeddingDatabase
+from janus.language.block import CodeBlock, TranslatedCodeBlock
+from janus.utils.enums import EmbeddingType
 
 
 class Vectorizer(object):
@@ -59,7 +59,7 @@ class Vectorizer(object):
         self,
         code_block: CodeBlock,
         collection_name: EmbeddingType | str,
-        filename: str  # perhaps this should be a relative path from the source, but for
+        filename: str,  # perhaps this should be a relative path from the source, but for
         # now we're all in 1 directory
     ) -> None:
         """Calculate `code_block` embedding, returning success & storing in `embedding_id`

janus/language/_tests/test_combine.py

@@ -1,6 +1,6 @@
 import unittest
 
-from ..combine import CodeBlock, Combiner, TranslatedCodeBlock
+from janus.language.combine import CodeBlock, Combiner, TranslatedCodeBlock
 
 
 class TestCombiner(unittest.TestCase):

janus/language/_tests/test_splitter.py

@@ -1,6 +1,6 @@
 import unittest
 
-from ..splitter import Splitter
+from janus.language.splitter import Splitter
 
 
 class TestSplitter(unittest.TestCase):

janus/language/alc/_tests/test_alc.py

@@ -1,9 +1,9 @@
 import unittest
 from pathlib import Path
 
-from ....llm import load_model
-from ...combine import Combiner
-from ..alc import AlcSplitter
+from janus.language.alc import AlcSplitter
+from janus.language.combine import Combiner
+from janus.llm import load_model
 
 
 class TestAlcSplitter(unittest.TestCase):
@@ -11,8 +11,8 @@ class TestAlcSplitter(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 = AlcSplitter(model=llm)
         self.combiner = Combiner(language="ibmhlasm")
         self.test_file = Path("janus/language/alc/_tests/alc.asm")
@@ -20,7 +20,7 @@ class TestAlcSplitter(unittest.TestCase):
     def test_split(self):
         """Test the split method."""
         tree_root = self.splitter.split(self.test_file)
-        self.assertEqual(tree_root.n_descendents, 34)
+        self.assertAlmostEqual(tree_root.n_descendents, 32, delta=5)
         self.assertLessEqual(tree_root.max_tokens, self.splitter.max_tokens)
         self.assertFalse(tree_root.complete)
         self.combiner.combine_children(tree_root)