Add rate limit handling to embedders #425
Conversation
|
|
||
| def is_rate_limit_error(exception: Exception) -> bool: | ||
| """Check if an exception is a rate limit error from any LLM provider. | ||
| """Check if an exception is a rate limit error from any LLM provider or embedder. |
There was a problem hiding this comment.
Do you think we can move this file somewhere else now that it's not only used for LLMs?
| ) | ||
| embedder = OpenAIEmbeddings(api_key="my key") | ||
| with pytest.raises( | ||
| EmbeddingsGenerationError, match="Failed to generate embedding with OpenAI" |
There was a problem hiding this comment.
I don't understand this test, shouldn't this error be retried? I'm not sure this test captures it? I would add a check to make sure the function is called the expected number of times.
| self._rate_limit_handler = rate_limit_handler | ||
| else: | ||
| self._rate_limit_handler = DEFAULT_RATE_LIMIT_HANDLER | ||
|
|
There was a problem hiding this comment.
We could consider having the embed_query method in the base class, that deals with the common logic of retries, and calls an _embed_query (or another name) in the subclasses. In this way, the implementation in the subclasses doesn't have to change (just update the method name):
In short:
base class:
@rate_limit_handler
def embed_query(self, query):
try:
return self._embed_query(query)
except Exception as e:
raise EmbeddingGenerationError() from e
@abstractmethod
def _embed_query(self, query):
...
and in the children:
# no need to add the decorator
def _embed_query(self, query):
# rest of the code
(up to you if you want to do it this way, it's not a request for change now)
Description
This PR extends the rate limiting functionality (previously available only for LLMs) to all embedding providers, and ensures consistent error handling
Type of Change
Complexity
Low
Complexity:
How Has This Been Tested?
Checklist
The following requirements should have been met (depending on the changes in the branch):