From 52303f8811eb781eba6a0a8b95f4c44342ada796 Mon Sep 17 00:00:00 2001
From: Ankush Khanna <ankush@cohere.com>
Date: Wed, 18 Sep 2024 17:39:29 +0200
Subject: [PATCH] Adding bearer token in compass

---
 compass_sdk/compass.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/compass_sdk/compass.py b/compass_sdk/compass.py
index 23c5307..0c4bcdb 100644
--- a/compass_sdk/compass.py
+++ b/compass_sdk/compass.py
@@ -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,
     ):
         """
@@ -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,
@@ -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