Skip to content

Commit

Permalink
fix issue when one of match groups is NoneType (#975)
Browse files Browse the repository at this point in the history
  • Loading branch information
LeonOstrez authored May 29, 2024
1 parent 9f225ba commit c9690ab
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 3 deletions.
9 changes: 6 additions & 3 deletions core/llm/openai_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,13 +104,16 @@ def rate_limit_sleep(self, err: RateLimitError) -> Optional[datetime.timedelta]:
match = re.search(time_regex, headers["x-ratelimit-reset-requests"])

if match:
seconds = int(match.group(1)) * 3600 + int(match.group(2)) * 60 + int(match.group(3))
hours = int(match.group(1)) if match.group(1) else 0
minutes = int(match.group(2)) if match.group(2) else 0
seconds = int(match.group(3)) if match.group(3) else 0
total_seconds = hours * 3600 + minutes * 60 + seconds
else:
# Not sure how this would happen, we would have to get a RateLimitError,
# but nothing (or invalid entry) in the `reset` field. Using a sane default.
seconds = 5
total_seconds = 5

return datetime.timedelta(seconds=seconds)
return datetime.timedelta(seconds=total_seconds)


__all__ = ["OpenAIClient"]
23 changes: 23 additions & 0 deletions tests/llm/test_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,3 +115,26 @@ async def test_openai_parser_fails(mock_AsyncOpenAI):

assert response is None
assert req_log.status == "error"


@pytest.mark.parametrize(
("remaining_tokens", "reset_tokens", "reset_requests", "expected"),
[
(0, "1h1m1s", "", 3661),
(0, "1h1s", "", 3601),
(0, "1m", "", 60),
(0, "", "1h1m1s", 0),
(1, "", "1h1m1s", 3661),
],
)
@patch("core.llm.openai_client.AsyncOpenAI")
def test_openai_rate_limit_parser(mock_AsyncOpenAI, remaining_tokens, reset_tokens, reset_requests, expected):
headers = {
"x-ratelimit-remaining-tokens": remaining_tokens,
"x-ratelimit-reset-tokens": reset_tokens,
"x-ratelimit-reset-requests": reset_requests,
}
err = MagicMock(response=MagicMock(headers=headers))

llm = OpenAIClient(LLMConfig(model="gpt-4"))
assert int(llm.rate_limit_sleep(err).total_seconds()) == expected

0 comments on commit c9690ab

Please sign in to comment.