fix: optional dependency issue

pyproject.toml

@@ -19,7 +19,7 @@ python = ">=3.9,<3.13"
 pydantic = "^2.5.3"
 openai = "^1.10.0"
 cohere = "^4.32"
-mistralai= "^0.0.12"
+mistralai= {version = "^0.0.12", optional = true}
 numpy = "^1.25.2"
 colorlog = "^6.8.0"
 pyyaml = "^6.0.1"
@@ -44,6 +44,7 @@ local = ["torch", "transformers", "llama-cpp-python"]
 pinecone = ["pinecone-client"]
 vision = ["torch", "torchvision", "transformers", "pillow"]
 processing = ["matplotlib"]
+mistralai = ["mistralai"]
 
 [tool.poetry.group.dev.dependencies]
 ipykernel = "^6.25.0"
@@ -67,4 +68,4 @@ build-backend = "poetry.core.masonry.api"
 line-length = 88
 
 [tool.mypy]
-ignore_missing_imports = true
+ignore_missing_imports = true

semantic_router/encoders/mistral.py

@@ -1,20 +1,20 @@
 """This file contains the MistralEncoder class which is used to encode text using MistralAI"""
 import os
 from time import sleep
-from typing import List, Optional
+from typing import List, Optional, Any
 
-from mistralai.client import MistralClient
-from mistralai.exceptions import MistralException
-from mistralai.models.embeddings import EmbeddingResponse
 
 from semantic_router.encoders import BaseEncoder
 from semantic_router.utils.defaults import EncoderDefault
+from pydantic.v1 import PrivateAttr
 
 
 class MistralEncoder(BaseEncoder):
     """Class to encode text using MistralAI"""
 
-    client: Optional[MistralClient]
+    _client: Any = PrivateAttr()
+    _embedding_response: Any = PrivateAttr()
+    _mistral_exception: Any = PrivateAttr()
     type: str = "mistral"
 
     def __init__(
@@ -26,33 +26,58 @@ def __init__(
         if name is None:
             name = EncoderDefault.MISTRAL.value["embedding_model"]
         super().__init__(name=name, score_threshold=score_threshold)
-        api_key = mistralai_api_key or os.getenv("MISTRALAI_API_KEY")
+        (
+            self._client,
+            self._embedding_response,
+            self._mistral_exception,
+        ) = self._initialize_client(mistralai_api_key)
+
+    def _initialize_client(self, api_key):
+        try:
+            from mistralai.client import MistralClient
+            from mistralai.exceptions import MistralException
+            from mistralai.models.embeddings import EmbeddingResponse
+        except ImportError:
+            raise ImportError(
+                "Please install MistralAI to use MistralEncoder. "
+                "You can install it with: "
+                "`pip install 'semantic-router[mistralai]'`"
+            )
+
+        api_key = api_key or os.getenv("MISTRALAI_API_KEY")
         if api_key is None:
             raise ValueError("Mistral API key not provided")
         try:
-            self.client = MistralClient(api_key=api_key)
+            client = MistralClient(api_key=api_key)
+            embedding_response = EmbeddingResponse
+            mistral_exception = MistralException
         except Exception as e:
             raise ValueError(f"Unable to connect to MistralAI {e.args}: {e}") from e
+        return client, embedding_response, mistral_exception
 
     def __call__(self, docs: List[str]) -> List[List[float]]:
-        if self.client is None:
+        if self._client is None:
             raise ValueError("Mistral client not initialized")
         embeds = None
         error_message = ""
 
         # Exponential backoff
         for _ in range(3):
             try:
-                embeds = self.client.embeddings(model=self.name, input=docs)
+                embeds = self._client.embeddings(model=self.name, input=docs)
                 if embeds.data:
                     break
-            except MistralException as e:
+            except self._mistral_exception as e:
                 sleep(2**_)
                 error_message = str(e)
             except Exception as e:
                 raise ValueError(f"Unable to connect to MistralAI {e.args}: {e}") from e
 
-        if not embeds or not isinstance(embeds, EmbeddingResponse) or not embeds.data:
+        if (
+            not embeds
+            or not isinstance(embeds, self._embedding_response)
+            or not embeds.data
+        ):
             raise ValueError(f"No embeddings returned from MistralAI: {error_message}")
         embeddings = [embeds_obj.embedding for embeds_obj in embeds.data]
         return embeddings

semantic_router/llms/llamacpp.py

@@ -2,26 +2,29 @@
 from pathlib import Path
 from typing import Any, Optional
 
-from llama_cpp import Llama, LlamaGrammar
+# from llama_cpp import Llama, LlamaGrammar
 
 from semantic_router.llms.base import BaseLLM
 from semantic_router.schema import Message
 from semantic_router.utils.logger import logger
 
+from pydantic.v1 import PrivateAttr
+
 
 class LlamaCppLLM(BaseLLM):
-    llm: Llama
+    llm: Any
     temperature: float
     max_tokens: Optional[int] = 200
-    grammar: Optional[LlamaGrammar] = None
+    grammar: Optional[Any] = None
+    _llama_cpp: Any = PrivateAttr()
 
     def __init__(
         self,
-        llm: Llama,
+        llm: Any,
         name: str = "llama.cpp",
         temperature: float = 0.2,
         max_tokens: Optional[int] = 200,
-        grammar: Optional[LlamaGrammar] = None,
+        grammar: Optional[Any] = None,
     ):
         super().__init__(
             name=name,
@@ -30,6 +33,18 @@ def __init__(
             max_tokens=max_tokens,
             grammar=grammar,
         )
+
+        try:
+            import llama_cpp
+        except ImportError:
+            raise ImportError(
+                "Please install LlamaCPP to use Llama CPP llm. "
+                "You can install it with: "
+                "`pip install 'semantic-router[llama-cpp-python]'`"
+            )
+        self._llama_cpp = llama_cpp
+        llm = self._llama_cpp.Llama
+        grammar = Optional[self._llama_cpp.LlamaGrammar]
         self.llm = llm
         self.temperature = temperature
         self.max_tokens = max_tokens
@@ -62,7 +77,7 @@ def _grammar(self):
         grammar_path = Path(__file__).parent.joinpath("grammars", "json.gbnf")
         assert grammar_path.exists(), f"{grammar_path}\ndoes not exist"
         try:
-            self.grammar = LlamaGrammar.from_file(grammar_path)
+            self.grammar = self._llama_cpp.LlamaGrammar.from_file(grammar_path)
             yield
         finally:
             self.grammar = None

semantic_router/llms/mistral.py

@@ -1,16 +1,17 @@
 import os
-from typing import List, Optional
+from typing import List, Optional, Any
 
-from mistralai.client import MistralClient
 
 from semantic_router.llms import BaseLLM
 from semantic_router.schema import Message
 from semantic_router.utils.defaults import EncoderDefault
 from semantic_router.utils.logger import logger
 
+from pydantic.v1 import PrivateAttr
+
 
 class MistralAILLM(BaseLLM):
-    client: Optional[MistralClient]
+    _client: Any = PrivateAttr()
     temperature: Optional[float]
     max_tokens: Optional[int]
 
@@ -24,23 +25,35 @@
         if name is None:
             name = EncoderDefault.MISTRAL.value["language_model"]
         super().__init__(name=name)
-        api_key = mistralai_api_key or os.getenv("MISTRALAI_API_KEY")
+        self._client = self._initialize_client(mistralai_api_key)
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+
+    def _initialize_client(self, api_key):
+        try:
+            from mistralai.client import MistralClient
+        except ImportError:
+            raise ImportError(
+                "Please install MistralAI to use MistralAI LLM. "
+                "You can install it with: "
+                "`pip install 'semantic-router[mistralai]'`"
+            )
+        api_key = api_key or os.getenv("MISTRALAI_API_KEY")
         if api_key is None:
             raise ValueError("MistralAI API key cannot be 'None'.")
         try:
-            self.client = MistralClient(api_key=api_key)
+            client = MistralClient(api_key=api_key)
         except Exception as e:
             raise ValueError(
                 f"MistralAI API client failed to initialize. Error: {e}"
             ) from e
-        self.temperature = temperature
-        self.max_tokens = max_tokens
+        return client
 
     def __call__(self, messages: List[Message]) -> str:
-        if self.client is None:
+        if self._client is None:
             raise ValueError("MistralAI client is not initialized.")
         try:
-            completion = self.client.chat(
+            completion = self._client.chat(
                 model=self.name,
                 messages=[m.to_mistral() for m in messages],
                 temperature=self.temperature,

tests/unit/encoders/test_mistral.py

@@ -4,6 +4,8 @@
 
 from semantic_router.encoders import MistralEncoder
 
+from unittest.mock import patch
+
 
 @pytest.fixture
 def mistralai_encoder(mocker):
@@ -12,9 +14,20 @@ def mistralai_encoder(mocker):
 
 
 class TestMistralEncoder:
+    def test_mistral_encoder_import_errors(self):
+        with patch.dict("sys.modules", {"mistralai": None}):
+            with pytest.raises(ImportError) as error:
+                MistralEncoder()
+
+        assert (
+            "Please install MistralAI to use MistralEncoder. "
+            "You can install it with: "
+            "`pip install 'semantic-router[mistralai]'`" in str(error.value)
+        )
+
     def test_mistralai_encoder_init_success(self, mocker):
         encoder = MistralEncoder(mistralai_api_key="test_api_key")
-        assert encoder.client is not None
+        assert encoder._client is not None
 
     def test_mistralai_encoder_init_no_api_key(self, mocker):
         mocker.patch("os.getenv", return_value=None)
@@ -23,7 +36,7 @@ def test_mistralai_encoder_init_no_api_key(self, mocker):
 
     def test_mistralai_encoder_call_uninitialized_client(self, mistralai_encoder):
         # Set the client to None to simulate an uninitialized client
-        mistralai_encoder.client = None
+        mistralai_encoder._client = None
         with pytest.raises(ValueError) as e:
             mistralai_encoder(["test document"])
         assert "Mistral client not initialized" in str(e.value)
@@ -60,7 +73,7 @@ def test_mistralai_encoder_call_success(self, mistralai_encoder, mocker):
 
         responses = [MistralException("mistralai error"), mock_response]
         mocker.patch.object(
-            mistralai_encoder.client, "embeddings", side_effect=responses
+            mistralai_encoder._client, "embeddings", side_effect=responses
         )
         embeddings = mistralai_encoder(["test document"])
         assert embeddings == [[0.1, 0.2]]
@@ -69,7 +82,7 @@ def test_mistralai_encoder_call_with_retries(self, mistralai_encoder, mocker):
         mocker.patch("os.getenv", return_value="fake-api-key")
         mocker.patch("time.sleep", return_value=None)  # To speed up the test
         mocker.patch.object(
-            mistralai_encoder.client,
+            mistralai_encoder._client,
             "embeddings",
             side_effect=MistralException("Test error"),
         )
@@ -83,7 +96,7 @@ def test_mistralai_encoder_call_failure_non_mistralai_error(
         mocker.patch("os.getenv", return_value="fake-api-key")
         mocker.patch("time.sleep", return_value=None)  # To speed up the test
         mocker.patch.object(
-            mistralai_encoder.client,
+            mistralai_encoder._client,
             "embeddings",
             side_effect=Exception("Non-MistralException"),
         )
@@ -118,7 +131,7 @@ def test_mistralai_encoder_call_successful_retry(self, mistralai_encoder, mocker
 
         responses = [MistralException("mistralai error"), mock_response]
         mocker.patch.object(
-            mistralai_encoder.client, "embeddings", side_effect=responses
+            mistralai_encoder._client, "embeddings", side_effect=responses
         )
         embeddings = mistralai_encoder(["test document"])
         assert embeddings == [[0.1, 0.2]]

tests/unit/llms/test_llm_llamacpp.py

@@ -4,6 +4,8 @@
 from semantic_router.llms.llamacpp import LlamaCppLLM
 from semantic_router.schema import Message
 
+from unittest.mock import patch
+
 
 @pytest.fixture
 def llamacpp_llm(mocker):
@@ -13,6 +15,17 @@ def llamacpp_llm(mocker):
 
 
 class TestLlamaCppLLM:
+    def test_llama_cpp_import_errors(self, llamacpp_llm):
+        with patch.dict("sys.modules", {"llama_cpp": None}):
+            with pytest.raises(ImportError) as error:
+                LlamaCppLLM(llamacpp_llm.llm)
+
+        assert (
+            "Please install LlamaCPP to use Llama CPP llm. "
+            "You can install it with: "
+            "`pip install 'semantic-router[llama-cpp-python]'`" in str(error.value)
+        )
+
     def test_llamacpp_llm_init_success(self, llamacpp_llm):
         assert llamacpp_llm.name == "llama.cpp"
         assert llamacpp_llm.temperature == 0.2

tests/unit/llms/test_llm_mistral.py

@@ -11,14 +11,25 @@ def mistralai_llm(mocker):
 
 
 class TestMistralAILLM:
+    # def test_mistral_llm_import_errors(self):
+    #     with patch.dict("sys.modules", {"mistralai": None}):
+    #         with pytest.raises(ImportError) as error:
+    #             MistralAILLM()
+
+    #     assert (
+    #         "Please install MistralAI to use MistralAI LLM. "
+    #         "You can install it with: "
+    #         "`pip install 'semantic-router[mistralai]'`" in str(error.value)
+    #     )
+
     def test_mistralai_llm_init_with_api_key(self, mistralai_llm):
-        assert mistralai_llm.client is not None, "Client should be initialized"
+        assert mistralai_llm._client is not None, "Client should be initialized"
         assert mistralai_llm.name == "mistral-tiny", "Default name not set correctly"
 
     def test_mistralai_llm_init_success(self, mocker):
         mocker.patch("os.getenv", return_value="fake-api-key")
         llm = MistralAILLM()
-        assert llm.client is not None
+        assert llm._client is not None
 
     def test_mistralai_llm_init_without_api_key(self, mocker):
         mocker.patch("os.getenv", return_value=None)
@@ -27,7 +38,7 @@ def test_mistralai_llm_init_without_api_key(self, mocker):
 
     def test_mistralai_llm_call_uninitialized_client(self, mistralai_llm):
         # Set the client to None to simulate an uninitialized client
-        mistralai_llm.client = None
+        mistralai_llm._client = None
         with pytest.raises(ValueError) as e:
             llm_input = [Message(role="user", content="test")]
             mistralai_llm(llm_input)
@@ -48,7 +59,7 @@ def test_mistralai_llm_call_success(self, mistralai_llm, mocker):
 
         mocker.patch("os.getenv", return_value="fake-api-key")
         mocker.patch.object(
-            mistralai_llm.client,
+            mistralai_llm._client,
             "chat",
             return_value=mock_completion,
         )