Skip to content

Commit

Permalink
fix: Allows clients to pass in an HTTP Session (#59)
Browse files Browse the repository at this point in the history
<!-- begin-generated-description -->

This pull request introduces changes to the `__init__` method of the
`compass.py` module, which is responsible for initializing a compass
client to interact with the Compass API. The modifications primarily
involve updating the handling of the `default_timeout` and introducing a
new `http_session` parameter.

## Changes:
- The `default_timeout` parameter's type hint has been updated from
`int` to `int | None`, indicating that it can now accept `None` as a
valid value.
- A new parameter, `http_session`, has been added to the `__init__`
method, allowing for the specification of an HTTP session with a custom
timeout.
- The instantiation of the `SessionWithDefaultTimeout` class has been
replaced with a conditional assignment, setting `self.session` to either
the provided `http_session` or a new instance of `requests.Session()`.
- A warning message is logged if `default_timeout` is not `None`,
indicating that the variable is deprecated and will not have any effect.
This message also provides guidance on using the `http_session`
parameter for specifying HTTP request timeouts.

<!-- end-generated-description -->
  • Loading branch information
coessiane authored Dec 10, 2024
1 parent 95dbd13 commit 8126e73
Showing 1 changed file with 2 additions and 15 deletions.
17 changes: 2 additions & 15 deletions cohere/compass/clients/compass.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,22 +65,9 @@ class RetryResult:
error: Optional[str] = None


_DEFAULT_TIMEOUT = 30


logger = logging.getLogger(__name__)


class SessionWithDefaultTimeout(requests.Session):
def __init__(self, timeout: int):
self._timeout = timeout
super().__init__()

def request(self, *args: Any, **kwargs: Any):
kwargs.setdefault("timeout", self._timeout)
return super().request(*args, **kwargs)


class CompassClient:
def __init__(
self,
Expand All @@ -89,7 +76,7 @@ def __init__(
username: Optional[str] = None,
password: Optional[str] = None,
bearer_token: Optional[str] = None,
default_timeout: int = _DEFAULT_TIMEOUT,
http_session: Optional[requests.Session] = None,
):
"""
A compass client to interact with the Compass API
Expand All @@ -100,7 +87,7 @@ def __init__(
self.index_url = index_url
self.username = username or os.getenv("COHERE_COMPASS_USERNAME")
self.password = password or os.getenv("COHERE_COMPASS_PASSWORD")
self.session = SessionWithDefaultTimeout(default_timeout)
self.session = http_session or requests.Session()
self.bearer_token = bearer_token

self.api_method = {
Expand Down

0 comments on commit 8126e73

Please sign in to comment.