diff --git a/src/unstract/llmwhisperer/client_v2.py b/src/unstract/llmwhisperer/client_v2.py index 70f6ec5..27b78a8 100644 --- a/src/unstract/llmwhisperer/client_v2.py +++ b/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", diff --git a/tests/integration/client_v2_test.py b/tests/integration/client_v2_test.py index 261ed96..a92bd02 100644 --- a/tests/integration/client_v2_test.py +++ b/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", ),