janus-llm 3.2.0__py3-none-any.whl → 3.3.0__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.2.0"
+__version__ = "3.3.0"
 
 # 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

@@ -108,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()

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,30 +39,23 @@ 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-4o",
+            model="gpt-4o-mini",
             source_language="fortran",
             target_language="python",
             target_version="3.10",
+            splitter_type="ast-flex",
         )
         self.test_file = Path("janus/language/treesitter/_tests/languages/fortran.f90")
         self.TEST_FILE_EMBEDDING_COUNT = 14
 
         self.req_translator = RequirementsDocumenter(
-            model="gpt-4o",
+            model="gpt-4o-mini",
             source_language="fortran",
             prompt_template="requirements",
         )
@@ -105,200 +72,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(

janus/converter/converter.py

@@ -6,7 +6,6 @@ from pathlib import Path
 from typing import Any
 
 from langchain.output_parsers import RetryWithErrorOutputParser
-from langchain.output_parsers.fix import OutputFixingParser
 from langchain_core.exceptions import OutputParserException
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.output_parsers import BaseOutputParser
@@ -29,6 +28,8 @@ from janus.llm import load_model
 from janus.llm.model_callbacks import get_model_callback
 from janus.llm.models_info import MODEL_PROMPT_ENGINES
 from janus.parsers.code_parser import GenericParser
+from janus.parsers.refiner_parser import RefinerParser
+from janus.refiners.refiner import BasicRefiner, Refiner
 from janus.utils.enums import LANGUAGES
 from janus.utils.logger import create_logger
 
@@ -75,6 +76,7 @@ class Converter:
         protected_node_types: tuple[str, ...] = (),
         prune_node_types: tuple[str, ...] = (),
         splitter_type: str = "file",
+        refiner_type: str = "basic",
     ) -> None:
         """Initialize a Converter instance.
 
@@ -84,6 +86,17 @@ class Converter:
                 values are `"code"`, `"text"`, `"eval"`, and `None` (default). If `None`,
                 the `Converter` assumes you won't be parsing an output (i.e., adding to an
                 embedding DB).
+            max_prompts: The maximum number of prompts to try before giving up.
+            max_tokens: The maximum number of tokens to use in the LLM. If `None`, the
+                converter will use half the model's token limit.
+            prompt_template: The name of the prompt template to use.
+            db_path: The path to the database to use for vectorization.
+            db_config: The configuration for the database.
+            protected_node_types: A set of node types that aren't to be merged.
+            prune_node_types: A set of node types which should be pruned.
+            splitter_type: The type of splitter to use. Valid values are `"file"`,
+                `"tag"`, `"chunk"`, `"ast-strict"`, and `"ast-flex"`.
+            refiner_type: The type of refiner to use. Valid values are `"basic"`.
         """
         self._changed_attrs: set = set()
 
@@ -116,7 +129,11 @@ class Converter:
         self._parser: BaseOutputParser = GenericParser()
         self._combiner: Combiner = Combiner()
 
+        self._refiner_type: str
+        self._refiner: Refiner
+
         self.set_splitter(splitter_type=splitter_type)
+        self.set_refiner(refiner_type=refiner_type)
         self.set_model(model_name=model, **model_arguments)
         self.set_prompt(prompt_template=prompt_template)
         self.set_source_language(source_language)
@@ -142,6 +159,7 @@ class Converter:
         self._load_prompt()
         self._load_splitter()
         self._load_vectorizer()
+        self._load_refiner()
         self._changed_attrs.clear()
 
     def set_model(self, model_name: str, **custom_arguments: dict[str, Any]):
@@ -179,6 +197,16 @@ class Converter:
         """
         self._splitter_type = splitter_type
 
+    def set_refiner(self, refiner_type: str) -> None:
+        """Validate and set the refiner name
+
+        The affected objects will not be updated until translate is called
+
+        Arguments:
+            refiner_type: the name of the refiner to use
+        """
+        self._refiner_type = refiner_type
+
     def set_source_language(self, source_language: str) -> None:
         """Validate and set the source language.
 
@@ -249,10 +277,24 @@ class Converter:
         )
 
         if self._splitter_type == "tag":
-            kwargs["tag"] = "<ITMOD_ALC_SPLIT>"
+            kwargs["tag"] = "<ITMOD_ALC_SPLIT>"  # Hardcoded for now
 
         self._splitter = CUSTOM_SPLITTERS[self._splitter_type](**kwargs)
 
+    @run_if_changed("_refiner_type", "_model_name")
+    def _load_refiner(self) -> None:
+        """Load the refiner according to this instance's attributes.
+
+        If the relevant fields have not been changed since the last time this method was
+        called, nothing happens.
+        """
+        if self._refiner_type == "basic":
+            self._refiner = BasicRefiner(
+                "basic_refinement", self._model_name, self._source_language
+            )
+        else:
+            raise ValueError(f"Error: unknown refiner type {self._refiner_type}")
+
     @run_if_changed("_model_name", "_custom_model_arguments")
     def _load_model(self) -> None:
         """Load the model according to this instance's attributes.
@@ -561,22 +603,22 @@ class Converter:
         # Retries with just the input
         n3 = math.ceil(self.max_prompts / (n1 * n2))
 
-        fix_format = OutputFixingParser.from_llm(
-            llm=self._llm,
+        refine_output = RefinerParser(
             parser=self._parser,
+            initial_prompt=self._prompt.format(**{"SOURCE_CODE": block.original.text}),
+            refiner=self._refiner,
             max_retries=n1,
+            llm=self._llm,
         )
         retry = RetryWithErrorOutputParser.from_llm(
             llm=self._llm,
-            parser=fix_format,
+            parser=refine_output,
             max_retries=n2,
         )
-
         completion_chain = self._prompt | self._llm
         chain = RunnableParallel(
             completion=completion_chain, prompt_value=self._prompt
         ) | RunnableLambda(lambda x: retry.parse_with_prompt(**x))
-
         for _ in range(n3):
             try:
                 return chain.invoke({"SOURCE_CODE": block.original.text})

janus/converter/diagram.py

@@ -1,10 +1,14 @@
-import json
-from copy import deepcopy
+import math
+
+from langchain.output_parsers import RetryWithErrorOutputParser
+from langchain_core.exceptions import OutputParserException
+from langchain_core.runnables import RunnableLambda, RunnableParallel
 
 from janus.converter.converter import run_if_changed
 from janus.converter.document import Documenter
 from janus.language.block import TranslatedCodeBlock
 from janus.llm.models_info import MODEL_PROMPT_ENGINES
+from janus.parsers.refiner_parser import RefinerParser
 from janus.parsers.uml import UMLSyntaxParser
 from janus.utils.logger import create_logger
 
@@ -47,65 +51,74 @@ class DiagramGenerator(Documenter):
             self._diagram_prompt_template_name = "diagram"
         self._load_diagram_prompt_engine()
 
-    def _add_translation(self, block: TranslatedCodeBlock) -> None:
-        """Given an "empty" `TranslatedCodeBlock`, translate the code represented in
-        `block.original`, setting the relevant fields in the translated block. The
-        `TranslatedCodeBlock` is updated in-pace, nothing is returned. Note that this
-        translates *only* the code for this block, not its children.
-
-        Arguments:
-            block: An empty `TranslatedCodeBlock`
-        """
-        if block.translated:
-            return
-
-        if block.original.text is None:
-            block.translated = True
-            return
-
-        if self._add_documentation:
-            documentation_block = deepcopy(block)
-            super()._add_translation(documentation_block)
-            if not documentation_block.translated:
-                message = "Error: unable to produce documentation for code block"
-                log.info(message)
-                raise ValueError(message)
-            documentation = json.loads(documentation_block.text)["docstring"]
-
-        if self._llm is None:
-            message = (
-                "Model not configured correctly, cannot translate. Try setting "
-                "the model"
-            )
-            log.error(message)
-            raise ValueError(message)
-
-        log.debug(f"[{block.name}] Translating...")
-        log.debug(f"[{block.name}] Input text:\n{block.original.text}")
-
+    def _run_chain(self, block: TranslatedCodeBlock) -> str:
         self._parser.set_reference(block.original)
+        n1 = round(self.max_prompts ** (1 / 3))
 
-        query_and_parse = self.diagram_prompt | self._llm | self._diagram_parser
+        # Retries with the input, output, and error
+        n2 = round((self.max_prompts // n1) ** (1 / 2))
+
+        # Retries with just the input
+        n3 = math.ceil(self.max_prompts / (n1 * n2))
 
         if self._add_documentation:
-            block.text = query_and_parse.invoke(
-                {
-                    "SOURCE_CODE": block.original.text,
-                    "DIAGRAM_TYPE": self._diagram_type,
-                    "DOCUMENTATION": documentation,
-                }
+            documentation_text = super()._run_chain(block)
+            refine_output = RefinerParser(
+                parser=self._diagram_parser,
+                initial_prompt=self._diagram_prompt.format(
+                    **{
+                        "SOURCE_CODE": block.original.text,
+                        "DOCUMENTATION": documentation_text,
+                        "DIAGRAM_TYPE": self._diagram_type,
+                    }
+                ),
+                refiner=self._refiner,
+                max_retries=n1,
+                llm=self._llm,
             )
         else:
-            block.text = query_and_parse.invoke(
-                {
-                    "SOURCE_CODE": block.original.text,
-                    "DIAGRAM_TYPE": self._diagram_type,
-                }
+            refine_output = RefinerParser(
+                parser=self._diagram_parser,
+                initial_prompt=self._diagram_prompt.format(
+                    **{
+                        "SOURCE_CODE": block.original.text,
+                        "DIAGRAM_TYPE": self._diagram_type,
+                    }
+                ),
+                refiner=self._refiner,
+                max_retries=n1,
+                llm=self._llm,
             )
-        block.tokens = self._llm.get_num_tokens(block.text)
-        block.translated = True
-
-        log.debug(f"[{block.name}] Output code:\n{block.text}")
+        retry = RetryWithErrorOutputParser.from_llm(
+            llm=self._llm,
+            parser=refine_output,
+            max_retries=n2,
+        )
+        completion_chain = self._prompt | self._llm
+        chain = RunnableParallel(
+            completion=completion_chain, prompt_value=self._diagram_prompt
+        ) | RunnableLambda(lambda x: retry.parse_with_prompt(**x))
+        for _ in range(n3):
+            try:
+                if self._add_documentation:
+                    return chain.invoke(
+                        {
+                            "SOURCE_CODE": block.original.text,
+                            "DOCUMENTATION": documentation_text,
+                            "DIAGRAM_TYPE": self._diagram_type,
+                        }
+                    )
+                else:
+                    return chain.invoke(
+                        {
+                            "SOURCE_CODE": block.original.text,
+                            "DIAGRAM_TYPE": self._diagram_type,
+                        }
+                    )
+            except OutputParserException:
+                pass
+
+        raise OutputParserException(f"Failed to parse after {n1*n2*n3} retries")
 
     @run_if_changed(
         "_diagram_prompt_template_name",
@@ -123,4 +136,4 @@ class DiagramGenerator(Documenter):
             target_version=None,
             prompt_template=self._diagram_prompt_template_name,
         )
-        self.diagram_prompt = self._diagram_prompt_engine.prompt
+        self._diagram_prompt = self._diagram_prompt_engine.prompt

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()