Skip to content

Commit

Permalink
Adding bearer token in compass (#16)
Browse files Browse the repository at this point in the history
<!-- begin-generated-description -->

This pull request introduces a new `bearer_token` parameter to the
`__init__` method of the `compass_sdk/compass.py` file. The
`bearer_token` is an optional string that can be used for authentication
purposes.

The changes made in this PR are as follows:
- A new `bearer_token` parameter is added to the `__init__` method,
allowing for optional authentication using a bearer token.
- The `self.bearer_token` attribute is set to the value of the
`bearer_token` parameter.
- The `_send_request_with_retry` method is modified to include a check
for the `self.bearer_token`. If it exists, the request headers are
updated with the bearer token for authentication.

The PR enhances the authentication process by providing an alternative
to the existing username and password-based authentication. The bearer
token can be used to securely authenticate requests, offering a more
flexible and secure approach.

<!-- end-generated-description -->
  • Loading branch information
ankush-cohere authored Sep 25, 2024
1 parent 21c1690 commit 6f26265
Showing 1 changed file with 12 additions and 5 deletions.
17 changes: 12 additions & 5 deletions compass_sdk/compass.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ def __init__(
index_url: str,
username: Optional[str] = None,
password: Optional[str] = None,
bearer_token: Optional[str] = None,
logger_level: LoggerLevel = LoggerLevel.INFO,
):
"""
Expand All @@ -87,6 +88,7 @@ def __init__(
self.username = username or os.getenv("COHERE_COMPASS_USERNAME")
self.password = password or os.getenv("COHERE_COMPASS_PASSWORD")
self.session = requests.Session()
self.bearer_token = bearer_token

self.function_call = {
"create_index": self.session.put,
Expand Down Expand Up @@ -497,18 +499,23 @@ def _send_request(
)
def _send_request_with_retry():
nonlocal error

try:

data_dict = None
if data:
if isinstance(data, BaseModel):
data_dict = data.model_dump(mode="json")
elif isinstance(data, Dict):
data_dict = data

response = self.function_call[function](
target_path, json=data_dict, auth=(self.username, self.password)
)
else:
response = self.function_call[function](target_path, auth=(self.username, self.password))
headers = None
auth = (self.username, self.password)
if self.bearer_token:
headers = {"Authorization": f"Bearer {self.bearer_token}"}
auth = None

response = self.function_call[function](target_path, json=data_dict, auth=auth, headers=headers)

if response.ok:
error = None
Expand Down

0 comments on commit 6f26265

Please sign in to comment.