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

janus/__init__.py

@@ -1,3 +1,11 @@
+import warnings
+
+from langchain_core._api.deprecation import LangChainDeprecationWarning
+
+from .metrics import *  # noqa: F403
 from .translate import Translator
 
-__version__ = "1.0.0"
+__version__ = "2.0.1"
+
+# Ignoring a deprecation warning from langchain_core that I can't seem to hunt down
+warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

janus/__main__.py

@@ -0,0 +1,4 @@
+from .cli import app
+
+if __name__ == "__main__":
+    app()

janus/_tests/test_cli.py

@@ -0,0 +1,128 @@
+import traceback
+import unittest
+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
+
+
+class TestCli(unittest.TestCase):
+    def setUp(self):
+        self.runner = CliRunner()
+
+    def test_help(self):
+        result = self.runner.invoke(app, ["--help"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_version(self):
+        result = self.runner.invoke(app, ["--version"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_invalid_command(self):
+        result = self.runner.invoke(app, ["invalid_command"])
+        self.assertNotEqual(result.exit_code, 0)
+
+    def test_no_arguments(self):
+        result = self.runner.invoke(app)
+        self.assertEqual(result.exit_code, 0)
+
+    def test_translate_help(self):
+        result = self.runner.invoke(app, ["translate", "-h"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_document_help(self):
+        result = self.runner.invoke(app, ["document", "-h"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_db_help(self):
+        result = self.runner.invoke(app, ["db", "-h"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_llm_help(self):
+        result = self.runner.invoke(app, ["llm", "-h"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_llm_add(self):
+        llm_model_path = MODEL_CONFIG_DIR / "test-model-name.json"
+        if llm_model_path.exists():
+            llm_model_path.unlink()
+        result = self.runner.invoke(app, ["llm", "add", "test-model-name"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertTrue(llm_model_path.exists())
+        llm_model_path.unlink()
+
+    def test_db_ls(self):
+        result = self.runner.invoke(app, ["db", "ls"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_db_status(self):
+        result = self.runner.invoke(app, ["db", "status"])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_embedding_add(self):
+        embedding_model_path = (
+            EMBEDDING_MODEL_CONFIG_DIR / "test-embedding-model-name.json"
+        )
+        if embedding_model_path.exists():
+            embedding_model_path.unlink()
+        result = self.runner.invoke(
+            app, ["embedding", "add", "test-embedding-model-name"]
+        )
+        self.assertEqual(result.exit_code, 0)
+        self.assertTrue(embedding_model_path.exists())
+        embedding_model_path.unlink()
+
+    def test_db_add_and_rm(self):
+        embedding_model_path = (
+            EMBEDDING_MODEL_CONFIG_DIR / "test-embedding-model-name.json"
+        )
+        if embedding_model_path.exists():
+            embedding_model_path.unlink()
+        result = self.runner.invoke(
+            app,
+            ["embedding", "add", "test-embedding-model-name", "-t", "HuggingFaceLocal"],
+        )
+        self.assertEqual(result.exit_code, 0)
+        result = self.runner.invoke(
+            app,
+            [
+                "db",
+                "add",
+                "test-db-name",
+                "test-embedding-model-name",
+                "-i",
+                "janus/language/mumps",
+            ],
+        )
+        traceback.print_exception(result.exception)
+        embedding_model_path.unlink()
+        self.assertEqual(result.exit_code, 0)
+
+        result = self.runner.invoke(app, ["db", "rm", "test-db-name", "-y"])
+        self.assertEqual(result.exit_code, 0)
+
+    @patch("janus.translate.Translator.translate", autospec=True)
+    def test_translate(self, mock_translate):
+        # Arrange
+        mock_instance = mock_translate.return_value
+        mock_instance.translate.return_value = None  # or whatever you expect
+
+        # Act
+        translate(
+            source_lang="matlab",
+            target_lang="python",
+            input_dir="janus/",
+            output_dir="janus/",
+            overwrite=True,
+            temp=0.7,
+            prompt_template="simple",
+            parser_type="code",
+            collection=None,
+        )
+
+        # Assert
+        mock_translate.assert_called_once()
+        mock_translate.assert_called_once_with(ANY, "janus/", "janus/", True, None)

janus/_tests/test_translate.py

@@ -7,7 +7,8 @@ from langchain.schema import Document
 from langchain.schema.embeddings import Embeddings
 from langchain.schema.vectorstore import VST, VectorStore
 
-from ..translate import Translator
+from ..language.block import CodeBlock, TranslatedCodeBlock
+from ..translate import DiagramGenerator, Translator
 
 # from langchain.vectorstores import Chroma
 
@@ -75,7 +76,7 @@ class TestTranslator(unittest.TestCase):
     def setUp(self):
         """Set up the tests."""
         self.translator = Translator(
-            model="gpt-3.5-turbo",
+            model="gpt-3.5-turbo-0125",
             source_language="fortran",
             target_language="python",
             target_version="3.10",
@@ -84,7 +85,7 @@ class TestTranslator(unittest.TestCase):
         self.TEST_FILE_EMBEDDING_COUNT = 14
 
         self.req_translator = Translator(
-            model="gpt-3.5-turbo",
+            model="gpt-3.5-turbo-0125",
             # embeddings_override=MockEmbeddingsFactory(),
             source_language="fortran",
             target_language="text",
@@ -317,6 +318,47 @@ class TestTranslator(unittest.TestCase):
         self.assertRaises(ValueError, self.translator._load_parameters)
 
 
+class TestDiagramGenerator(unittest.TestCase):
+    """Tests for the DiagramGenerator class."""
+
+    def setUp(self):
+        """Set up the tests."""
+        self.diagram_generator = DiagramGenerator(
+            model="gpt-3.5-turbo-0125",
+            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._source_language, "fortran")
+        self.assertEqual(self.diagram_generator._diagram_type, "Activity")
+
+    def test_add_translation(self):
+        """Test _add_translation method."""
+        block = TranslatedCodeBlock(
+            original=CodeBlock(
+                id="test",
+                name="Test Block",
+                node_type="function",
+                language="python",
+                text="print('Hello, World!')",
+                start_point=(0, 0),
+                end_point=(1, 0),
+                start_byte=0,
+                end_byte=1,
+                tokens=5,
+                children=[],
+            ),
+            language="python",
+        )
+        self.diagram_generator._add_translation(block)
+        self.assertTrue(block.translated)
+        self.assertIsNotNone(block.text)
+        self.assertIsNotNone(block.tokens)
+
+
 @pytest.mark.parametrize(
     "source_language,prompt_template,expected_target_language,expected_target_version,"
     "parser_type",
@@ -337,8 +379,8 @@ def test_language_combinations(
     """Tests that translator target language settings are consistent
     with prompt template expectations.
     """
-    translator = Translator(model="gpt-3.5-turbo")
-    translator.set_model("gpt-3.5-turbo-16k")
+    translator = Translator(model="gpt-3.5-turbo-0125")
+    translator.set_model("gpt-3.5-turbo-0125")
     translator.set_source_language(source_language)
     translator.set_target_language(expected_target_language, expected_target_version)
     translator.set_parser_type(parser_type)
@@ -348,5 +390,5 @@ def test_language_combinations(
     assert translator._target_version == expected_target_version  # nosec
     assert translator._parser_type == parser_type  # nosec
     assert translator._splitter.language == source_language  # nosec
-    assert translator._splitter.model.model_name == "gpt-3.5-turbo-16k"  # nosec
-    assert translator._prompt_engine._template_name == prompt_template  # nosec
+    assert translator._splitter.model.model_name == "gpt-3.5-turbo-0125"  # nosec
+    assert translator._prompt_template_name == prompt_template  # nosec