add max_retries to openai and azure encoders

semantic_router/encoders/openai.py

@@ -42,6 +42,7 @@ class OpenAIEncoder(BaseEncoder):
     token_limit: int = 8192  # default value, should be replaced by config
     _token_encoder: Any = PrivateAttr()
     type: str = "openai"
+    max_retries: int
 
     def __init__(
         self,
@@ -51,9 +52,13 @@ class OpenAIEncoder(BaseEncoder):
         openai_org_id: Optional[str] = None,
         score_threshold: Optional[float] = None,
         dimensions: Union[int, NotGiven] = NotGiven(),
+        max_retries: int | None = None,
     ):
         if name is None:
             name = EncoderDefault.OPENAI.value["embedding_model"]
+
+        max_retries = max_retries if max_retries is not None else 3
+
         if score_threshold is None and name in model_configs:
             set_score_threshold = model_configs[name].threshold
         elif score_threshold is None:
@@ -66,6 +71,7 @@ class OpenAIEncoder(BaseEncoder):
         super().__init__(
             name=name,
             score_threshold=set_score_threshold,
+            max_retries=max_retries,
         )
         api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
         base_url = openai_base_url or os.getenv("OPENAI_BASE_URL")
@@ -109,8 +115,9 @@ class OpenAIEncoder(BaseEncoder):
             docs = [self._truncate(doc) for doc in docs]
 
         # Exponential backoff
-        for j in range(1, 7):
+        for j in range(self.max_retries + 1):
             try:
+                raise OpenAIError("Test")
                 embeds = self.client.embeddings.create(
                     input=docs,
                     model=self.name,
@@ -119,12 +126,14 @@ class OpenAIEncoder(BaseEncoder):
                 if embeds.data:
                     break
             except OpenAIError as e:
-                sleep(2**j)
-                error_message = str(e)
-                logger.warning(f"Retrying in {2**j} seconds...")
+                logger.error("Exception occurred", exc_info=True)
+                if self.max_retries != 0:
+                    sleep(2**j)
+                    logger.warning(f"Retrying in {2**j} seconds due to OpenAIError: {e}")
+
             except Exception as e:
-                logger.error(f"OpenAI API call failed. Error: {error_message}")
-                raise ValueError(f"OpenAI API call failed. Error: {e}") from e
+                logger.error(f"OpenAI API call failed. Error: {e}")
+                raise ValueError(f"OpenAI API call failed. Error: {str(e)}") from e
 
         if (
             not embeds
@@ -132,7 +141,7 @@ class OpenAIEncoder(BaseEncoder):
             or not embeds.data
         ):
             logger.info(f"Returned embeddings: {embeds}")
-            raise ValueError(f"No embeddings returned. Error: {error_message}")
+            raise ValueError(f"No embeddings returned.")
 
         embeddings = [embeds_obj.embedding for embeds_obj in embeds.data]
         return embeddings
@@ -161,8 +170,9 @@ class OpenAIEncoder(BaseEncoder):
             docs = [self._truncate(doc) for doc in docs]
 
         # Exponential backoff
-        for j in range(1, 7):
+        for j in range(self.max_retries + 1):
             try:
+                raise OpenAIError("Test")
                 embeds = await self.async_client.embeddings.create(
                     input=docs,
                     model=self.name,
@@ -171,11 +181,12 @@ class OpenAIEncoder(BaseEncoder):
                 if embeds.data:
                     break
             except OpenAIError as e:
-                await asleep(2**j)
-                error_message = str(e)
-                logger.warning(f"Retrying in {2**j} seconds...")
+                logger.error("Exception occurred", exc_info=True)
+                if self.max_retries != 0:
+                    await asleep(2**j)
+                    logger.warning(f"Retrying in {2**j} seconds due to OpenAIError: {e}")
             except Exception as e:
-                logger.error(f"OpenAI API call failed. Error: {error_message}")
+                logger.error(f"OpenAI API call failed. Error: {e}")
                 raise ValueError(f"OpenAI API call failed. Error: {e}") from e
 
         if (
@@ -184,7 +195,7 @@ class OpenAIEncoder(BaseEncoder):
             or not embeds.data
         ):
             logger.info(f"Returned embeddings: {embeds}")
-            raise ValueError(f"No embeddings returned. Error: {error_message}")
+            raise ValueError(f"No embeddings returned.")
 
         embeddings = [embeds_obj.embedding for embeds_obj in embeds.data]
         return embeddings

semantic_router/encoders/zure.py

@@ -23,6 +23,7 @@ class AzureOpenAIEncoder(BaseEncoder):
     azure_endpoint: Optional[str] = None
     api_version: Optional[str] = None
     model: Optional[str] = None
+    max_retries: int
 
     def __init__(
         self,
@@ -33,11 +34,15 @@ class AzureOpenAIEncoder(BaseEncoder):
         model: Optional[str] = None,  # TODO we should change to `name` JB
         score_threshold: float = 0.82,
         dimensions: Union[int, NotGiven] = NotGiven(),
+        max_retries: int | None = None,
     ):
         name = deployment_name
         if name is None:
             name = EncoderDefault.AZURE.value["embedding_model"]
-        super().__init__(name=name, score_threshold=score_threshold)
+
+        max_retries = max_retries if max_retries is not None else 3
+        
+        super().__init__(name=name, score_threshold=score_threshold, max_retries=max_retries)
         self.api_key = api_key
         self.deployment_name = deployment_name
         self.azure_endpoint = azure_endpoint
@@ -100,8 +105,9 @@ class AzureOpenAIEncoder(BaseEncoder):
         error_message = ""
 
         # Exponential backoff
-        for j in range(3):
+        for j in range(self.max_retries + 1):
             try:
+                raise OpenAIError("Test")
                 embeds = self.client.embeddings.create(
                     input=docs,
                     model=str(self.model),
@@ -110,15 +116,12 @@ class AzureOpenAIEncoder(BaseEncoder):
                 if embeds.data:
                     break
             except OpenAIError as e:
-                # print full traceback
-                import traceback
-
-                traceback.print_exc()
-                sleep(2**j)
-                error_message = str(e)
-                logger.warning(f"Retrying in {2**j} seconds...")
+                logger.error("Exception occurred", exc_info=True)
+                if self.max_retries != 0:
+                    sleep(2**j)
+                    logger.warning(f"Retrying in {2**j} seconds due to OpenAIError: {e}")
             except Exception as e:
-                logger.error(f"Azure OpenAI API call failed. Error: {error_message}")
+                logger.error(f"Azure OpenAI API call failed. Error: {e}")
                 raise ValueError(f"Azure OpenAI API call failed. Error: {e}") from e
 
         if (
@@ -126,7 +129,7 @@ class AzureOpenAIEncoder(BaseEncoder):
             or not isinstance(embeds, CreateEmbeddingResponse)
             or not embeds.data
         ):
-            raise ValueError(f"No embeddings returned. Error: {error_message}")
+            raise ValueError(f"No embeddings returned.")
 
         embeddings = [embeds_obj.embedding for embeds_obj in embeds.data]
         return embeddings
@@ -138,8 +141,9 @@ class AzureOpenAIEncoder(BaseEncoder):
         error_message = ""
 
         # Exponential backoff
-        for j in range(3):
+        for j in range(self.max_retries + 1):
             try:
+                raise OpenAIError("Test")
                 embeds = await self.async_client.embeddings.create(
                     input=docs,
                     model=str(self.model),
@@ -147,16 +151,14 @@ class AzureOpenAIEncoder(BaseEncoder):
                 )
                 if embeds.data:
                     break
-            except OpenAIError as e:
-                # print full traceback
-                import traceback
 
-                traceback.print_exc()
-                await asleep(2**j)
-                error_message = str(e)
-                logger.warning(f"Retrying in {2**j} seconds...")
+            except OpenAIError as e:
+                logger.error("Exception occurred", exc_info=True)
+                if self.max_retries != 0:
+                    await asleep(2**j)
+                    logger.warning(f"Retrying in {2**j} seconds due to OpenAIError: {e}")
             except Exception as e:
-                logger.error(f"Azure OpenAI API call failed. Error: {error_message}")
+                logger.error(f"Azure OpenAI API call failed. Error: {e}")
                 raise ValueError(f"Azure OpenAI API call failed. Error: {e}") from e
 
         if (
@@ -164,7 +166,7 @@ class AzureOpenAIEncoder(BaseEncoder):
             or not isinstance(embeds, CreateEmbeddingResponse)
             or not embeds.data
         ):
-            raise ValueError(f"No embeddings returned. Error: {error_message}")
+            raise ValueError(f"No embeddings returned.")
 
         embeddings = [embeds_obj.embedding for embeds_obj in embeds.data]
         return embeddings