Merge pull request #5 from Zipstack/feat/make-request-header-part-of-…

…class feat: Make request header part of class
Zipstack · Jul 10, 2024 · bb384fd · bb384fd
2 parents aa0d613 + 57ea65b
commit bb384fd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/src/unstract/llmwhisperer/__init__.py b/src/unstract/llmwhisperer/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.2.0"
+__version__ = "0.2.1"
 
 from .client import LLMWhispererClient  # noqa: F401
 

diff --git a/src/unstract/llmwhisperer/client.py b/src/unstract/llmwhisperer/client.py
@@ -118,6 +118,8 @@ def __init__(
 
         self.api_timeout = api_timeout
 
+        self.headers = {"unstract-key": self.api_key}
+
     def get_usage_info(self) -> dict:
         """Retrieves the usage information of the LLMWhisperer API.
 
@@ -135,7 +137,7 @@ def get_usage_info(self) -> dict:
         self.logger.debug("get_usage_info called")
         url = f"{self.base_url}/get-usage-info"
         self.logger.debug("url: %s", url)
-        req = requests.Request("GET", url, headers={"unstract-key": self.api_key})
+        req = requests.Request("GET", url, headers=self.headers)
         prepared = req.prepare()
         s = requests.Session()
         response = s.send(prepared, timeout=self.api_timeout)
@@ -233,11 +235,11 @@ def whisper(
                 "POST",
                 api_url,
                 params=params,
-                headers={"unstract-key": self.api_key},
+                headers=self.headers,
                 data=data,
             )
         else:
-            req = requests.Request("POST", api_url, params=params, headers={"unstract-key": self.api_key})
+            req = requests.Request("POST", api_url, params=params, headers=self.headers)
         prepared = req.prepare()
         s = requests.Session()
         response = s.send(prepared, timeout=self.api_timeout)
@@ -279,7 +281,7 @@ def whisper_status(self, whisper_hash: str) -> dict:
         url = f"{self.base_url}/whisper-status"
         params = {"whisper-hash": whisper_hash}
         self.logger.debug("url: %s", url)
-        req = requests.Request("GET", url, headers={"unstract-key": self.api_key}, params=params)
+        req = requests.Request("GET", url, headers=self.headers, params=params)
         prepared = req.prepare()
         s = requests.Session()
         response = s.send(prepared, timeout=self.api_timeout)
@@ -314,7 +316,7 @@ def whisper_retrieve(self, whisper_hash: str) -> dict:
         url = f"{self.base_url}/whisper-retrieve"
         params = {"whisper-hash": whisper_hash}
         self.logger.debug("url: %s", url)
-        req = requests.Request("GET", url, headers={"unstract-key": self.api_key}, params=params)
+        req = requests.Request("GET", url, headers=self.headers, params=params)
         prepared = req.prepare()
         s = requests.Session()
         response = s.send(prepared, timeout=self.api_timeout)
@@ -357,7 +359,7 @@ def highlight_data(self, whisper_hash: str, search_text: str) -> dict:
         req = requests.Request(
             "POST",
             url,
-            headers={"unstract-key": self.api_key},
+            headers=self.headers,
             params=params,
             data=search_text,
         )