cohere-python

1
# This file was auto-generated by Fern from our API Definition.
2

3
import asyncio
4
import email.utils
5
import re
6
import time
7
import typing
8
from contextlib import asynccontextmanager, contextmanager
9
from functools import wraps
10
from random import random
11

12
import httpx
13

14
INITIAL_RETRY_DELAY_SECONDS = 0.5
15
MAX_RETRY_DELAY_SECONDS = 10
16
MAX_RETRY_DELAY_SECONDS_FROM_HEADER = 30
17

18

19
def _parse_retry_after(response_headers: httpx.Headers) -> typing.Optional[float]:
20
    """
21
    This function parses the `Retry-After` header in a HTTP response and returns the number of seconds to wait.
22

23
    Inspired by the urllib3 retry implementation.
24
    """
25
    retry_after_ms = response_headers.get("retry-after-ms")
26
    if retry_after_ms is not None:
27
        try:
28
            return int(retry_after_ms) / 1000 if retry_after_ms > 0 else 0
29
        except Exception:
30
            pass
31

32
    retry_after = response_headers.get("retry-after")
33
    if retry_after is None:
34
        return None
35

36
    # Attempt to parse the header as an int.
37
    if re.match(r"^\s*[0-9]+\s*$", retry_after):
38
        seconds = float(retry_after)
39
    # Fallback to parsing it as a date.
40
    else:
41
        retry_date_tuple = email.utils.parsedate_tz(retry_after)
42
        if retry_date_tuple is None:
43
            return None
44
        if retry_date_tuple[9] is None:  # Python 2
45
            # Assume UTC if no timezone was specified
46
            # On Python2.7, parsedate_tz returns None for a timezone offset
47
            # instead of 0 if no timezone is given, where mktime_tz treats
48
            # a None timezone offset as local time.
49
            retry_date_tuple = retry_date_tuple[:9] + (0,) + retry_date_tuple[10:]
50

51
        retry_date = email.utils.mktime_tz(retry_date_tuple)
52
        seconds = retry_date - time.time()
53

54
    if seconds < 0:
55
        seconds = 0
56

57
    return seconds
58

59

60
def _retry_timeout(response: httpx.Response, retries: int) -> float:
61
    """
62
    Determine the amount of time to wait before retrying a request.
63
    This function begins by trying to parse a retry-after header from the response, and then proceeds to use exponential backoff
64
    with a jitter to determine the number of seconds to wait.
65
    """
66

67
    # If the API asks us to wait a certain amount of time (and it's a reasonable amount), just do what it says.
68
    retry_after = _parse_retry_after(response.headers)
69
    if retry_after is not None and retry_after <= MAX_RETRY_DELAY_SECONDS_FROM_HEADER:
70
        return retry_after
71

72
    # Apply exponential backoff, capped at MAX_RETRY_DELAY_SECONDS.
73
    retry_delay = min(INITIAL_RETRY_DELAY_SECONDS * pow(2.0, retries), MAX_RETRY_DELAY_SECONDS)
74

75
    # Add a randomness / jitter to the retry delay to avoid overwhelming the server with retries.
76
    timeout = retry_delay * (1 - 0.25 * random())
77
    return timeout if timeout >= 0 else 0
78

79

80
def _should_retry(response: httpx.Response) -> bool:
81
    retriable_400s = [429, 408, 409]
82
    return response.status_code >= 500 or response.status_code in retriable_400s
83

84

85
class HttpClient:
86
    def __init__(self, *, httpx_client: httpx.Client):
87
        self.httpx_client = httpx_client
88

89
    # Ensure that the signature of the `request` method is the same as the `httpx.Client.request` method
90
    @wraps(httpx.Client.request)
91
    def request(
92
        self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any
93
    ) -> httpx.Response:
94
        response = self.httpx_client.request(*args, **kwargs)
95
        if _should_retry(response=response):
96
            if max_retries > retries:
97
                time.sleep(_retry_timeout(response=response, retries=retries))
98
                return self.request(max_retries=max_retries, retries=retries + 1, *args, **kwargs)
99
        return response
100

101
    @wraps(httpx.Client.stream)
102
    @contextmanager
103
    def stream(self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any) -> typing.Any:
104
        with self.httpx_client.stream(*args, **kwargs) as stream:
105
            yield stream
106

107

108
class AsyncHttpClient:
109
    def __init__(self, *, httpx_client: httpx.AsyncClient):
110
        self.httpx_client = httpx_client
111

112
    # Ensure that the signature of the `request` method is the same as the `httpx.Client.request` method
113
    @wraps(httpx.AsyncClient.request)
114
    async def request(
115
        self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any
116
    ) -> httpx.Response:
117
        response = await self.httpx_client.request(*args, **kwargs)
118
        if _should_retry(response=response):
119
            if max_retries > retries:
120
                await asyncio.sleep(_retry_timeout(response=response, retries=retries))
121
                return await self.request(max_retries=max_retries, retries=retries + 1, *args, **kwargs)
122
        return response
123

124
    @wraps(httpx.AsyncClient.stream)
125
    @asynccontextmanager
126
    async def stream(
127
        self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any
128
    ) -> typing.Any:
129
        async with self.httpx_client.stream(*args, **kwargs) as stream:
130
            yield stream
131