Adding bearer token in compass

cohere-ai · Sep 18, 2024 · 5ecdff6 · 5ecdff6
1 parent 3960096
commit 5ecdff6
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/compass_sdk/compass.py b/compass_sdk/compass.py
@@ -70,6 +70,7 @@ def __init__(
         index_url: str,
         username: Optional[str] = None,
         password: Optional[str] = None,
+        bearer_token: Optional[str] = None,
         logger_level: LoggerLevel = LoggerLevel.INFO,
     ):
         """
@@ -82,6 +83,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,
@@ -439,18 +441,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()
                     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