Skip to content

Commit

Permalink
[chore]: Added Retrying Mechanism in case of Request Rate Limit Error
Browse files Browse the repository at this point in the history
  • Loading branch information
keenborder786 committed Nov 1, 2024
1 parent b4cb208 commit 4fcc456
Showing 1 changed file with 17 additions and 4 deletions.
21 changes: 17 additions & 4 deletions libs/partners/mistralai/langchain_mistralai/embeddings.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from typing import Iterable, List

import httpx
from httpx import Response
from langchain_core.embeddings import Embeddings
from langchain_core.utils import (
secret_from_env,
Expand All @@ -15,6 +16,7 @@
SecretStr,
model_validator,
)
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
from tokenizers import Tokenizer # type: ignore
from typing_extensions import Self

Expand Down Expand Up @@ -209,16 +211,27 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]:
List of embeddings, one for each text.
"""
try:
batch_responses = (
self.client.post(
batch_responses = []

@retry(
retry=retry_if_exception_type(Exception),
wait=wait_fixed(30), # Wait 30 seconds between retries
stop=stop_after_attempt(5), # Stop after 5 attempts
)
def _embed_batch(batch: List[str]) -> Response:
response = self.client.post(
url="/embeddings",
json=dict(
model=self.model,
input=batch,
),
)
for batch in self._get_batches(texts)
)
if response.status_code == 429:
raise Exception("Requests rate limit exceeded")
return response

for batch in self._get_batches(texts):
batch_responses.append(_embed_batch(batch))
return [
list(map(float, embedding_obj["embedding"]))
for response in batch_responses
Expand Down

0 comments on commit 4fcc456

Please sign in to comment.