Merge pull request #23 from Zipstack/fix/custom-header-support

header updated

Author: jaseemjaskp

src/unstract/llmwhisperer/client_v2.py

@@ -25,6 +25,7 @@
 from typing import IO
 
 import requests
+from typing import Optional, Dict
 
 BASE_URL_V2 = "https://llmwhisperer-api.us-central.unstract.com/api/v2"
 
@@ -76,6 +77,7 @@ def __init__(
         base_url: str = "",
         api_key: str = "",
         logging_level: str = "",
+        custom_headers: Optional[Dict[str, str]] = None,
     ):
         """Initializes the LLMWhispererClient with the given parameters.
 
@@ -92,6 +94,12 @@ def __init__(
                                            value of the LLMWHISPERER_LOGGING_LEVEL
                                            environment variable, or "DEBUG" if the
                                            environment variable is not set.
+            custom_headers (Optional[Dict[str, str]], optional): Custom headers to add to
+                                                                every request. These will
+                                                                be merged with default
+                                                                headers, with custom
+                                                                headers taking precedence.
+                                                                Defaults to None.
         """
         if logging_level == "":
             logging_level = os.getenv("LLMWHISPERER_LOGGING_LEVEL", "DEBUG")
@@ -117,6 +125,8 @@ def __init__(
             self.api_key = api_key
 
         self.headers = {"unstract-key": self.api_key}
+        if custom_headers:
+            self.headers.update(custom_headers)
         # For test purpose
         # self.headers = {
         #     "Subscription-Id": "python-client",

tests/integration/client_v2_test.py

@@ -174,7 +174,7 @@ def test_whisper_v2_url_in_post(client_v2, data_dir, output_mode, mode, url, inp
     "url,token,webhook_name",
     [
         (
-            "https://webhook.site/b76ecc5f-8320-4410-b24f-66525d2c92cb",
+            "https://webhook.site/62bb38ac-408c-4fcf-b8f1-cb22adbf3f96", # need to find a clean solution
             "",
             "client_v2_test",
         ),