Skip to content

Commit

Permalink
fixing pytest issues for optional dependency issues
Browse files Browse the repository at this point in the history
  • Loading branch information
zahid-syed committed Mar 11, 2024
1 parent 363ae33 commit 3b0de2d
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 13 deletions.
18 changes: 9 additions & 9 deletions semantic_router/encoders/mistral.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@ class MistralEncoder(BaseEncoder):
"""Class to encode text using MistralAI"""

_client: Any = PrivateAttr()
embedding_response: Any = PrivateAttr()
mistral_exception: Any = PrivateAttr()
_embedding_response: Any = PrivateAttr()
_mistral_exception: Any = PrivateAttr()
type: str = "mistral"

def __init__(
Expand Down Expand Up @@ -51,33 +51,33 @@ def _initialize_client(self, api_key):
)

try:
self.client = MistralClient(api_key=api_key)
self.embedding_response = EmbeddingResponse
self.mistral_exception = MistralException
self._client = MistralClient(api_key=api_key)
self._embedding_response = EmbeddingResponse
self._mistral_exception = MistralException
except Exception as e:
raise ValueError(f"Unable to connect to MistralAI {e.args}: {e}") from e

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 self.mistral_exception 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, self.embedding_response)
or not isinstance(embeds, self._embedding_response)
or not embeds.data
):
raise ValueError(f"No embeddings returned from MistralAI: {error_message}")
Expand Down
8 changes: 4 additions & 4 deletions semantic_router/llms/mistral.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@


class MistralAILLM(BaseLLM):
client: Any = PrivateAttr()
_client: Any = PrivateAttr()
temperature: Optional[float]
max_tokens: Optional[int]

Expand Down Expand Up @@ -42,17 +42,17 @@ def _initialize_client(self, api_key):
"`pip install 'semantic-router[mistralai]'`"
)
try:
self.client = MistralClient(api_key=api_key)
self._client = MistralClient(api_key=api_key)
except Exception as e:
raise ValueError(
f"MistralAI API client failed to initialize. Error: {e}"
) from e

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,
Expand Down

0 comments on commit 3b0de2d

Please sign in to comment.