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

janus/converter.py

@@ -4,7 +4,6 @@ from typing import Any
 from langchain.schema.language_model import BaseLanguageModel
 
 from .language.binary import BinarySplitter
-from .language.combine import Combiner
 from .language.mumps import MumpsSplitter
 from .language.splitter import Splitter
 from .language.treesitter import TreeSitterSplitter
@@ -45,6 +44,8 @@ class Converter:
         self,
         source_language: str = "fortran",
         max_tokens: None | int = None,
+        protected_node_types: set[str] | list[str] | tuple[str] = (),
+        prune_node_types: set[str] | list[str] | tuple[str] = (),
     ) -> None:
         """Initialize a Converter instance.
 
@@ -59,13 +60,15 @@ class Converter:
 
         self._source_language: None | str
         self._source_glob: None | str
+        self._protected_node_types: tuple[str] = ()
+        self._prune_node_types: tuple[str] = ()
         self._splitter: None | Splitter
         self._llm: None | BaseLanguageModel = None
         self._max_tokens: None | int = max_tokens
 
-        self._combiner: Combiner = Combiner()
-
-        self.set_source_language(source_language=source_language)
+        self.set_source_language(source_language)
+        self.set_protected_node_types(protected_node_types)
+        self.set_prune_node_types(prune_node_types)
 
         # Child class must call this. Should we enforce somehow?
         # self._load_parameters()
@@ -86,7 +89,7 @@ class Converter:
     def set_source_language(self, source_language: str) -> None:
         """Validate and set the source language.
 
-        The affected objects will not be updated until translate() is called.
+        The affected objects will not be updated until _load_parameters() is called.
 
         Arguments:
             source_language: The source programming language.
@@ -101,27 +104,55 @@ class Converter:
         self._source_glob = f"**/*.{LANGUAGES[source_language]['suffix']}"
         self._source_language = source_language
 
-    @run_if_changed("_source_language", "_max_tokens", "_llm")
+    def set_protected_node_types(
+        self, protected_node_types: set[str] | list[str] | tuple[str]
+    ) -> None:
+        """Set the protected (non-mergeable) node types. This will often be structures
+        like functions, classes, or modules which you might want to keep separate
+
+        The affected objects will not be updated until _load_parameters() is called.
+
+        Arguments:
+            protected_node_types: A set of node types that aren't to be merged
+        """
+        self._protected_node_types = tuple(set(protected_node_types or []))
+
+    def set_prune_node_types(
+        self, prune_node_types: set[str] | list[str] | tuple[str]
+    ) -> None:
+        """Set the node types to prune. This will often be structures
+        like comments or whitespace which you might want to keep out of the LLM
+
+        The affected objects will not be updated until _load_parameters() is called.
+
+        Arguments:
+            prune_node_types: A set of node types which should be pruned
+        """
+        self._prune_node_types = tuple(set(prune_node_types or []))
+
+    @run_if_changed(
+        "_source_language",
+        "_max_tokens",
+        "_llm",
+        "_protected_node_types",
+        "_prune_node_types",
+    )
     def _load_splitter(self) -> None:
         """Load the splitter according to this instance's attributes.
 
         If the relevant fields have not been changed since the last time this method was
         called, nothing happens.
         """
+        kwargs = dict(
+            max_tokens=self._max_tokens,
+            model=self._llm,
+            protected_node_types=self._protected_node_types,
+            prune_node_types=self._prune_node_types,
+        )
         if self._source_language in CUSTOM_SPLITTERS:
             if self._source_language == "mumps":
-                self._splitter = MumpsSplitter(
-                    max_tokens=self._max_tokens,
-                    model=self._llm,
-                )
+                self._splitter = MumpsSplitter(**kwargs)
             elif self._source_language == "binary":
-                self._splitter = BinarySplitter(
-                    max_tokens=self._max_tokens,
-                    model=self._llm,
-                )
+                self._splitter = BinarySplitter(**kwargs)
         else:
-            self._splitter = TreeSitterSplitter(
-                language=self._source_language,
-                max_tokens=self._max_tokens,
-                model=self._llm,
-            )
+            self._splitter = TreeSitterSplitter(language=self._source_language, **kwargs)

janus/embedding/_tests/test_collections.py

@@ -14,9 +14,7 @@ class TestCollections(unittest.TestCase):
         self.collections = Collections(self._db)
 
     def test_creation(self):
-        self._db.create_collection.return_value = "foo"
-
-        result = self.collections.create(EmbeddingType.PSEUDO)
+        self.collections.create(EmbeddingType.PSEUDO)
 
         metadata = {
             "date_updated": datetime.datetime.now().date().isoformat(),
@@ -24,12 +22,9 @@ class TestCollections(unittest.TestCase):
         }
 
         self._db.create_collection.assert_called_with("pseudo_1", metadata=metadata)
-        self.assertEqual(result, "foo")
 
     def test_creation_triangulation(self):
-        self._db.create_collection.return_value = []
-
-        result = self.collections.create(EmbeddingType.REQUIREMENT)
+        self.collections.create(EmbeddingType.REQUIREMENT)
 
         metadata = {
             "date_updated": datetime.datetime.now().date().isoformat(),
@@ -37,7 +32,6 @@ class TestCollections(unittest.TestCase):
         }
 
         self._db.create_collection.assert_called_with("requirement_1", metadata=metadata)
-        self.assertEqual(result, [])
 
     def test_creation_of_existing_type(self):
         mock_collection = MagicMock()

janus/embedding/_tests/test_database.py

@@ -0,0 +1,32 @@
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from ..database import ChromaEmbeddingDatabase, uri_to_path
+
+
+class TestDatabase(unittest.TestCase):
+    def test_uri_to_path(self):
+        uri = (Path.home().expanduser() / "Documents" / "testfile.txt").as_uri()
+        expected_path = Path.home().expanduser() / "Documents" / "testfile.txt"
+        self.assertEqual(uri_to_path(uri), expected_path)
+
+    @patch("chromadb.PersistentClient", autospec=True)
+    def test_ChromaEmbeddingDatabase(self, mock_client):
+        # Test with default path
+        _ = ChromaEmbeddingDatabase()
+        mock_client.assert_called_once()
+
+        # Test with custom path
+        custom_path = "/custom/path/to/chroma-data"
+        _ = ChromaEmbeddingDatabase(custom_path)
+        mock_client.assert_called()
+
+        # Test with URL
+        url = "http://example.com/chroma-data"
+        _ = ChromaEmbeddingDatabase(url)
+        mock_client.assert_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

janus/embedding/_tests/test_vectorize.py

@@ -5,6 +5,7 @@ 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
 
@@ -22,7 +23,7 @@ class MockDBVectorizer(VectorizerFactory):
         model: None | str = "gpt4all",
         path: str | Path = None,
     ) -> Vectorizer:
-        return Vectorizer(self._db, source_language, max_tokens, model)
+        return Vectorizer(self._db)
 
 
 class TestVectorize(unittest.TestCase):
@@ -35,7 +36,11 @@ class TestVectorize(unittest.TestCase):
         self.database.list_collections = list_collections
         self.vectorizer = MockDBVectorizer(self.database).create_vectorizer()
         self.test_file = Path("janus/language/treesitter/_tests/languages/fortran.f90")
-        self.test_block = self.vectorizer._splitter.split(self.test_file)
+        splitter = TreeSitterSplitter(
+            language="fortran",
+            max_tokens=16_384,
+        )
+        self.test_block = splitter.split(self.test_file)
 
     def test_add_nodes_recursively(self):
         embedding_type = EmbeddingType.SOURCE
@@ -45,9 +50,9 @@ class TestVectorize(unittest.TestCase):
             "time_updated": datetime.datetime.now().time().isoformat("minutes"),
         }
         self.database.create_collection.assert_called_with("source_1", metadata=metadata)
-        self.vectorizer._add_nodes_recursively(
+        self.vectorizer.add_nodes_recursively(
             self.test_block, embedding_type, self.test_file.name
         )
         self.database.get_or_create_collection.assert_called_with(
-            "source_1", metadata=metadata
+            name="source_1", embedding_function=None, metadata=metadata
         )

janus/embedding/collections.py

@@ -1,10 +1,12 @@
 import datetime
 import os
-from typing import Sequence
+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
 
 # See https://docs.trychroma.com/telemetry#in-chromas-backend-using-environment-variables
 os.environ["ANONYMIZED_TELEMETRY"] = "False"
@@ -13,10 +15,16 @@ os.environ["ANONYMIZED_TELEMETRY"] = "False"
 class Collections:
     """Manage embedding collections"""
 
-    def __init__(self, client: Client):
+    def __init__(self, client: Client, config: Optional[Dict[str, str]] = None):
         self._client = client
+        if config is not None:
+            self._config = config
+        else:
+            self._config = {}
 
-    def create(self, name: EmbeddingType | str) -> Collection:
+    def create(
+        self, name: EmbeddingType | str, model_name: Optional[str] = None
+    ) -> Chroma:
         """Create a Chroma collection for the given embedding type.
 
         Arguments:
@@ -27,9 +35,23 @@ class Collections:
             "date_updated": datetime.datetime.now().date().isoformat(),
             "time_updated": datetime.datetime.now().time().isoformat("minutes"),
         }
-        return self._client.create_collection(collection_name, metadata=metadata)
+        if model_name is not None:
+            metadata["embedding_model"] = model_name
+            self._client.create_collection(collection_name, metadata=metadata)
+            self._config[collection_name] = model_name
+            model, _, _ = load_embedding_model(model_name)
+            return Chroma(
+                client=self._client,
+                collection_name=collection_name,
+                embedding_function=model,
+            )
+        else:
+            self._client.create_collection(collection_name, metadata=metadata)
+            return Chroma(client=self._client, collection_name=collection_name)
 
-    def get_or_create(self, name: EmbeddingType | str) -> Collection:
+    def get_or_create(
+        self, name: EmbeddingType | str, model_name: Optional[str] = None
+    ) -> Chroma:
         """Create a Chroma collection for the given embedding type.
 
         Arguments:
@@ -40,7 +62,26 @@ class Collections:
             "date_updated": datetime.datetime.now().date().isoformat(),
             "time_updated": datetime.datetime.now().time().isoformat("minutes"),
         }
-        return self._client.get_or_create_collection(collection_name, metadata=metadata)
+        if collection_name in self._config:
+            model_name = self._config[collection_name]
+        if model_name is not None:
+            metadata["embedding_model"] = model_name
+            self._config[collection_name] = model_name
+            model, _, _ = load_embedding_model(model_name)
+            self._client.get_or_create_collection(collection_name, metadata=metadata)
+            return Chroma(
+                client=self._client,
+                collection_name=collection_name,
+                embedding_function=model,
+                collection_metadata=metadata,
+            )
+        else:
+            self._client.get_or_create_collection(collection_name, metadata=metadata)
+            return Chroma(
+                client=self._client,
+                collection_name=collection_name,
+                collection_metadata=metadata,
+            )
 
     def get(self, name: None | EmbeddingType | str = None) -> Sequence[Collection]:
         """Get the Chroma collections.
@@ -61,6 +102,8 @@ class Collections:
             collection_name = name.name.lower()
         else:
             collection_name = name
+        if collection_name in self._config:
+            del self._config[collection_name]
         self._client.delete_collection(collection_name)
 
     def _set_collection_name(self, name: EmbeddingType | str) -> str:

janus/embedding/embedding_models_info.py

@@ -0,0 +1,120 @@
+import json
+from pathlib import Path
+from typing import Any, Callable, Dict, Tuple
+
+from aenum import MultiValueEnum
+from dotenv import load_dotenv
+from langchain_community.embeddings.huggingface import (
+    HuggingFaceEmbeddings,
+    HuggingFaceInferenceAPIEmbeddings,
+)
+from langchain_core.embeddings import Embeddings
+from langchain_openai import OpenAIEmbeddings
+
+from janus.utils.logger import create_logger
+
+load_dotenv()
+
+log = create_logger(__name__)
+
+
+class EmbeddingModelType(MultiValueEnum):
+    OpenAI = "OpenAI", "openai", "open-ai", "oai"
+    HuggingFaceLocal = "HuggingFaceLocal", "huggingfacelocal", "huggingface-local", "hfl"
+    HuggingFaceInferenceAPI = (
+        "HuggingFaceInferenceAPI",
+        "huggingfaceinferenceapi",
+        "huggingface-inference-api",
+        "hfia",
+    )
+
+
+EMBEDDING_MODEL_TYPE_CONSTRUCTORS: Dict[
+    EmbeddingModelType, Callable[[Any], Embeddings]
+] = {}
+
+for model_type in EmbeddingModelType:
+    for value in model_type.values:
+        if model_type == EmbeddingModelType.OpenAI:
+            EMBEDDING_MODEL_TYPE_CONSTRUCTORS[value] = OpenAIEmbeddings
+        elif model_type == EmbeddingModelType.HuggingFaceLocal:
+            EMBEDDING_MODEL_TYPE_CONSTRUCTORS[value] = HuggingFaceEmbeddings
+        elif model_type == EmbeddingModelType.HuggingFaceInferenceAPI:
+            EMBEDDING_MODEL_TYPE_CONSTRUCTORS[value] = HuggingFaceInferenceAPIEmbeddings
+
+EMBEDDING_MODEL_TYPE_DEFAULT_IDS: Dict[EmbeddingModelType, Dict[str, Any]] = {
+    EmbeddingModelType.OpenAI.value: "text-embedding-3-small",
+    EmbeddingModelType.HuggingFaceLocal.value: "all-MiniLM-L6-v2",
+    EmbeddingModelType.HuggingFaceInferenceAPI.value: "",
+}
+
+EMBEDDING_MODEL_DEFAULT_ARGUMENTS: Dict[str, Dict[str, Any]] = {
+    "text-embedding-3-small": dict(model="text-embedding-3-small"),
+    "text-embedding-3-large": dict(model="text-embedding-3-large"),
+    "text-embedding-ada-002": dict(model="text-embedding-ada-002"),
+}
+
+EMBEDDING_MODEL_CONFIG_DIR = Path.home().expanduser() / ".janus" / "embeddings"
+
+EMBEDDING_TOKEN_LIMITS: Dict[str, int] = {
+    "text-embedding-3-small": 8191,
+    "text-embedding-3-large": 8191,
+    "text-embedding-ada-002": 8191,
+}
+
+EMBEDDING_COST_PER_MODEL: Dict[str, float] = {
+    "text-embedding-3-small": {"input": 0.02 / 1e6, "output": 0.0},
+    "text-embedding-3-large": {"input": 0.13 / 1e6, "output": 0.0},
+    "text-embedding-ada-002": {"input": 0.10 / 1e6, "output": 0.0},
+}
+
+
+def load_embedding_model(
+    model_name: str,
+) -> Tuple[Embeddings, int, Dict[str, float]]:
+    """Load an embedding model from the configuration file or create a new one
+
+    Arguments:
+        model_name: The user-given name of the model to load.
+        model_type: The type of the model to load.
+        identifier: The identifier for the model (e.g. the name, URL, or HuggingFace
+            path).
+    """
+    if not EMBEDDING_MODEL_CONFIG_DIR.exists():
+        EMBEDDING_MODEL_CONFIG_DIR.mkdir(parents=True)
+    model_config_file = EMBEDDING_MODEL_CONFIG_DIR / f"{model_name}.json"
+
+    if not model_config_file.exists():
+        # The default model type is HuggingFaceLocal because that's the default for Chroma
+        model_type = EmbeddingModelType.HuggingFaceLocal.value
+        identifier = EMBEDDING_MODEL_TYPE_DEFAULT_IDS[model_type]
+        model_config = {
+            "model_type": model_type,
+            "model_identifier": identifier,
+            "model_args": EMBEDDING_MODEL_DEFAULT_ARGUMENTS.get(identifier, {}),
+            "token_limit": EMBEDDING_TOKEN_LIMITS.get(identifier, 8191),
+            "model_cost": EMBEDDING_COST_PER_MODEL.get(
+                identifier, {"input": 0, "output": 0}
+            ),
+        }
+        log.info(
+            f"WARNING: Creating new model config file: \
+                {model_config_file} with default config"
+        )
+        with open(model_config_file, "w") as f:
+            json.dump(model_config, f, indent=2)
+    else:
+        with open(model_config_file, "r") as f:
+            model_config = json.load(f)
+    model_constructor = EMBEDDING_MODEL_TYPE_CONSTRUCTORS[model_config["model_type"]]
+    model_args = model_config["model_args"]
+    if model_config["model_type"] in EmbeddingModelType.HuggingFaceInferenceAPI.values:
+        model_args.update({"api_url": model_config["model_identifier"]})
+    elif model_config["model_type"] in EmbeddingModelType.HuggingFaceLocal.values:
+        model_args.update({"model_name": model_config["model_identifier"]})
+    model = model_constructor(**model_args)
+    return (
+        model,
+        model_config["token_limit"],
+        model_config["model_cost"],
+    )

janus/embedding/vectorize.py

@@ -1,67 +1,52 @@
 import uuid
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Sequence
+from typing import Any, Dict, Optional, Sequence
 
 from chromadb import Client, Collection
+from langchain_community.vectorstores import Chroma
 
-from ..converter import Converter
-from ..language.block import CodeBlock
-from ..llm.models_info import TOKEN_LIMITS
+from ..language.block import CodeBlock, TranslatedCodeBlock
 from ..utils.enums import EmbeddingType
 from .collections import Collections
 from .database import ChromaEmbeddingDatabase
 
 
-class Vectorizer(Converter):
+class Vectorizer(object):
     """Class for creating embeddings/vectors in a specified ChromaDB"""
 
-    def __init__(
-        self,
-        client: Client,
-        source_language: str,
-        max_tokens: None | int,
-        model: None | str,
-    ) -> None:
+    def __init__(self, client: Client, config: Optional[Dict[str, Any]] = None) -> None:
         """Initializes the Vectorizer class
 
         Arguments:
             client: ChromaDB client instance
-            source_language: The source programming language.
-            max_tokens: The maximum number of tokens to send to the embedding model at
-                once. If `None`, the `Vectorizer` will use the default value for the
-                `model`.
-            model: The name of the model to use. This will also determine the `max_tokens`
-                if that variable is not set.
         """
-        if max_tokens is None:
-            max_tokens = TOKEN_LIMITS[model]
-
-        super().__init__(
-            source_language=source_language,
-            max_tokens=max_tokens,
-        )
         self._db = client
-        self._collections = Collections(self._db)
+        self._collections = Collections(self._db, config)
 
-        super()._load_parameters()
+    def get_or_create_collection(
+        self, name: EmbeddingType | str, model_name: Optional[str] = None
+    ) -> Chroma:
+        return self._collections.get_or_create(name, model_name=model_name)
 
-    def create_collection(self, embedding_type: EmbeddingType) -> Collection:
-        return self._collections.create(embedding_type)
+    def create_collection(
+        self, embedding_type: EmbeddingType, model_name: Optional[str] = None
+    ) -> Chroma:
+        return self._collections.create(embedding_type, model_name=model_name)
 
     def collections(
         self, name: None | EmbeddingType | str = None
     ) -> Sequence[Collection]:
         return self._collections.get(name)
 
-    def _add_nodes_recursively(
+    def add_nodes_recursively(
         self, code_block: CodeBlock, collection_name: EmbeddingType | str, file_name: str
     ) -> None:
         """Embed all nodes in the tree rooted at `code_block`
 
         Arguments:
             code_block: CodeBlock to embed
-            embedding_type: EmbeddingType to use
+            collection_name: Collection to add to
             file_name: Name of file containing `code_block`
         """
         nodes = [code_block]
@@ -74,41 +59,55 @@ class Vectorizer(Converter):
         self,
         code_block: CodeBlock,
         collection_name: EmbeddingType | str,
-        file_name: 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
-    ) -> bool:
+    ) -> None:
         """Calculate `code_block` embedding, returning success & storing in `embedding_id`
 
         Arguments:
             code_block: CodeBlock to embed
-            embedding_type: EmbeddingType to use
-            file_name: Name of file containing `code_block`
-
-        Returns:
-            True if embedding was successful, False otherwise
+            collection_name: Collection to add to
+            filename: Name of file containing `code_block`
         """
         if code_block.text:
             metadatas = [
                 {
-                    "type": code_block.type,
-                    "original_filename": file_name,
+                    "type": code_block.node_type,
+                    "id": code_block.id,
+                    "name": code_block.name,
+                    "language": code_block.language,
+                    "filename": filename,
                     "tokens": code_block.tokens,
                     "cost": 0,  # TranslatedCodeBlock has cost
                 },
             ]
+            if collection_name in self.config:
+                metadatas[0]["embedding_model"] = self.config[collection_name]
             # for now, dealing with missing metadata by skipping it
+            if isinstance(code_block, TranslatedCodeBlock):
+                self._add(
+                    code_block=code_block.original,
+                    collection_name=collection_name,
+                    filename=filename,
+                )
+                if code_block.original.embedding_id is not None:
+                    metadatas[0][
+                        "original_embedding_id"
+                    ] = code_block.original.embedding_id
+                metadatas[0]["cost"] = code_block.cost
             if code_block.text is not None:
                 metadatas[0]["hash"] = hash(code_block.text)
             if code_block.start_point is not None:
                 metadatas[0]["start_line"] = code_block.start_point[0]
             if code_block.end_point is not None:
                 metadatas[0]["end_line"] = code_block.end_point[0]
+            # TODO: Add metadata about translation parameters (e.g. model)
             the_text = [code_block.text]
-            code_block.embedding_id = self.add_text(collection_name, the_text, metadatas)[
-                0
-            ]
-            return True
-        return False
+            code_block.embedding_id = self.add_text(
+                collection_name,
+                the_text,
+                metadatas,
+            )[0]
 
     def add_text(
         self,
@@ -121,7 +120,7 @@ class Vectorizer(Converter):
         metadatas, returning the embedding id
 
         Arguments:
-            embedding_type: EmbeddingType to use
+            collection_name: Collection to add to
             texts: list of texts to store
             metadatas: list of metadatas to store
             ids: list of embedding ids (must match lengh of texts),
@@ -137,20 +136,20 @@ class Vectorizer(Converter):
             # based on the text.
             ids = [str(uuid.uuid3(uuid.NAMESPACE_DNS, text)) for text in texts]
         collection = self._collections.get_or_create(collection_name)
-        collection.upsert(ids=ids, documents=texts, metadatas=metadatas)
+        collection.add_texts(ids=ids, texts=texts, metadatas=metadatas)
         return ids
 
+    @property
+    def config(self):
+        return self._collections._config
+
 
 class VectorizerFactory(ABC):
     """Interface for creating a Vectorizer independent of type of ChromaDB client"""
 
     @abstractmethod
     def create_vectorizer(
-        self,
-        source_language: str,
-        max_tokens: None | int,
-        model: None | str,
-        path: str | Path,
+        self, path: str | Path, config: Dict[str, Any] = {}
     ) -> Vectorizer:
         """Factory method"""
 
@@ -160,19 +159,11 @@ class ChromaDBVectorizer(VectorizerFactory):
 
     def create_vectorizer(
         self,
-        source_language: str = "fortran",
-        max_tokens: None | int = None,
-        model: None | str = "gpt4all",
         path: str | Path = Path.home() / ".janus" / "chroma" / "chroma-data",
+        config: Optional[Dict[str, Any]] = None,
     ) -> Vectorizer:
         """
         Arguments:
-            source_language: The source programming language.
-            max_tokens: The maximum number of tokens to send to the embedding model at
-                once. If `None`, the `Vectorizer` will use the default value for the
-                `model`.
-            model: The name of the model to use. This will also determine the `max_tokens`
-                if that variable is not set.
             path: The path to the ChromaDB. Can be either a string of a URL or path or a
                 Path object
 
@@ -180,4 +171,4 @@ class ChromaDBVectorizer(VectorizerFactory):
                 Vectorizer
         """
         database = ChromaEmbeddingDatabase(path)
-        return Vectorizer(database, source_language, max_tokens, model)
+        return Vectorizer(database, config)