Skip to content

Commit 0c3aa60

Browse files
Khurdhula-HarshavardhanKhurdhula-Harshavardhan
committed
fix: pass user defined headers as config.
1 parent 8fa4631 commit 0c3aa60

18 files changed

+105
-137
lines changed

‎jigsawstack/async_request.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
class AsyncRequestConfig(TypedDict):
1616
api_url: str
1717
api_key: str
18-
disable_request_logging: Union[bool, None] = False
18+
headers: Union[Dict[str, str], None]
1919

2020

2121
class AsyncRequest(Generic[T]):
@@ -180,9 +180,6 @@ def __get_headers(self) -> Dict[str, str]:
180180
if not self.files and not self.data:
181181
h["Content-Type"] = "application/json"
182182

183-
if self.disable_request_logging:
184-
h["x-jigsaw-no-request-log"] = "true"
185-
186183
_headers = h.copy()
187184

188185
# don't override Content-Type if using multipart

‎jigsawstack/audio.py‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -55,14 +55,10 @@ def __init__(
5555
self,
5656
api_key: str,
5757
api_url: str,
58-
disable_request_logging: Union[bool, None] = False,
58+
headers: Union[Dict[str, str], None] = None,
5959
):
60-
super().__init__(api_key, api_url, disable_request_logging)
61-
self.config = RequestConfig(
62-
api_url=api_url,
63-
api_key=api_key,
64-
disable_request_logging=disable_request_logging,
65-
)
60+
super().__init__(api_key, api_url, headers)
61+
self.config = RequestConfig(api_url=api_url, api_key=api_key, headers=headers)
6662

6763
@overload
6864
def speech_to_text(
@@ -108,13 +104,13 @@ def __init__(
108104
self,
109105
api_key: str,
110106
api_url: str,
111-
disable_request_logging: Union[bool, None] = False,
107+
headers: Union[Dict[str, str], None] = None,
112108
):
113-
super().__init__(api_key, api_url, disable_request_logging)
109+
super().__init__(api_key, api_url, headers)
114110
self.config = AsyncRequestConfig(
115111
api_url=api_url,
116112
api_key=api_key,
117-
disable_request_logging=disable_request_logging,
113+
headers=headers,
118114
)
119115

120116
@overload

‎jigsawstack/classification.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -68,13 +68,13 @@ def __init__(
6868
self,
6969
api_key: str,
7070
api_url: str,
71-
disable_request_logging: Union[bool, None] = False,
71+
headers: Union[Dict[str, str], None] = None,
7272
):
73-
super().__init__(api_key, api_url, disable_request_logging)
73+
super().__init__(api_key, api_url, headers)
7474
self.config = RequestConfig(
7575
api_url=api_url,
7676
api_key=api_key,
77-
disable_request_logging=disable_request_logging,
77+
headers=headers,
7878
)
7979

8080
def classify(self, params: ClassificationParams) -> ClassificationResponse:
@@ -95,13 +95,13 @@ def __init__(
9595
self,
9696
api_key: str,
9797
api_url: str,
98-
disable_request_logging: Union[bool, None] = False,
98+
headers: Union[Dict[str, str], None] = None,
9999
):
100-
super().__init__(api_key, api_url, disable_request_logging)
100+
super().__init__(api_key, api_url, headers)
101101
self.config = AsyncRequestConfig(
102102
api_url=api_url,
103103
api_key=api_key,
104-
disable_request_logging=disable_request_logging,
104+
headers=headers,
105105
)
106106

107107
async def classify(self, params: ClassificationParams) -> ClassificationResponse:

‎jigsawstack/embedding.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,13 @@ def __init__(
3434
self,
3535
api_key: str,
3636
api_url: str,
37-
disable_request_logging: Union[bool, None] = False,
37+
headers: Union[Dict[str, str], None] = None,
3838
):
39-
super().__init__(api_key, api_url, disable_request_logging)
39+
super().__init__(api_key, api_url, headers)
4040
self.config = RequestConfig(
4141
api_url=api_url,
4242
api_key=api_key,
43-
disable_request_logging=disable_request_logging,
43+
headers=headers,
4444
)
4545

4646
@overload
@@ -82,13 +82,13 @@ def __init__(
8282
self,
8383
api_key: str,
8484
api_url: str,
85-
disable_request_logging: Union[bool, None] = False,
85+
headers: Union[Dict[str, str], None] = None,
8686
):
87-
super().__init__(api_key, api_url, disable_request_logging)
87+
super().__init__(api_key, api_url, headers)
8888
self.config = RequestConfig(
8989
api_url=api_url,
9090
api_key=api_key,
91-
disable_request_logging=disable_request_logging,
91+
headers=headers,
9292
)
9393

9494
@overload

‎jigsawstack/embedding_v2.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,13 @@ def __init__(
3232
self,
3333
api_key: str,
3434
api_url: str,
35-
disable_request_logging: Union[bool, None] = False,
35+
headers: Union[Dict[str, str], None] = None,
3636
):
37-
super().__init__(api_key, api_url, disable_request_logging)
37+
super().__init__(api_key, api_url, headers)
3838
self.config = RequestConfig(
3939
api_url=api_url,
4040
api_key=api_key,
41-
disable_request_logging=disable_request_logging,
41+
headers=headers,
4242
)
4343

4444
@overload
@@ -80,13 +80,13 @@ def __init__(
8080
self,
8181
api_key: str,
8282
api_url: str,
83-
disable_request_logging: Union[bool, None] = False,
83+
headers: Union[Dict[str, str], None] = None,
8484
):
85-
super().__init__(api_key, api_url, disable_request_logging)
85+
super().__init__(api_key, api_url, headers)
8686
self.config = RequestConfig(
8787
api_url=api_url,
8888
api_key=api_key,
89-
disable_request_logging=disable_request_logging,
89+
headers=headers,
9090
)
9191

9292
@overload

‎jigsawstack/image_generation.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -90,13 +90,13 @@ def __init__(
9090
self,
9191
api_key: str,
9292
api_url: str,
93-
disable_request_logging: Union[bool, None] = False,
93+
headers: Union[Dict[str, str], None] = None,
9494
):
95-
super().__init__(api_key, api_url, disable_request_logging=disable_request_logging)
95+
super().__init__(api_key, api_url, headers)
9696
self.config = RequestConfig(
9797
api_url=api_url,
9898
api_key=api_key,
99-
disable_request_logging=disable_request_logging,
99+
headers=headers,
100100
)
101101

102102
def image_generation(
@@ -119,13 +119,13 @@ def __init__(
119119
self,
120120
api_key: str,
121121
api_url: str,
122-
disable_request_logging: Union[bool, None] = False,
122+
headers: Union[Dict[str, str], None] = None,
123123
):
124-
super().__init__(api_key, api_url, disable_request_logging=disable_request_logging)
124+
super().__init__(api_key, api_url, headers)
125125
self.config = RequestConfig(
126126
api_url=api_url,
127127
api_key=api_key,
128-
disable_request_logging=disable_request_logging,
128+
headers=headers,
129129
)
130130

131131
async def image_generation(

‎jigsawstack/prediction.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -49,13 +49,13 @@ def __init__(
4949
self,
5050
api_key: str,
5151
api_url: str,
52-
disable_request_logging: Union[bool, None] = False,
52+
headers: Union[Dict[str, str], None] = None,
5353
):
54-
super().__init__(api_key, api_url, disable_request_logging)
54+
super().__init__(api_key, api_url, headers)
5555
self.config = RequestConfig(
5656
api_url=api_url,
5757
api_key=api_key,
58-
disable_request_logging=disable_request_logging,
58+
headers=headers,
5959
)
6060

6161
def predict(self, params: PredictionParams) -> PredictionResponse:
@@ -76,13 +76,13 @@ def __init__(
7676
self,
7777
api_key: str,
7878
api_url: str,
79-
disable_request_logging: Union[bool, None] = False,
79+
headers: Union[Dict[str, str], None] = None,
8080
):
81-
super().__init__(api_key, api_url, disable_request_logging)
81+
super().__init__(api_key, api_url, headers)
8282
self.config = RequestConfig(
8383
api_url=api_url,
8484
api_key=api_key,
85-
disable_request_logging=disable_request_logging,
85+
headers=headers,
8686
)
8787

8888
async def predict(self, params: PredictionParams) -> PredictionResponse:

‎jigsawstack/prompt_engine.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,13 +98,13 @@ def __init__(
9898
self,
9999
api_key: str,
100100
api_url: str,
101-
disable_request_logging: Union[bool, None] = False,
101+
headers: Union[Dict[str, str], None] = None,
102102
):
103-
super().__init__(api_key, api_url, disable_request_logging)
103+
super().__init__(api_key, api_url, headers)
104104
self.config = RequestConfig(
105105
api_url=api_url,
106106
api_key=api_key,
107-
disable_request_logging=disable_request_logging,
107+
headers=headers,
108108
)
109109

110110
def create(self, params: PromptEngineCreateParams) -> PromptEngineCreateResponse:

‎jigsawstack/request.py‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
class RequestConfig(TypedDict):
1515
api_url: str
1616
api_key: str
17-
disable_request_logging: Union[bool, None] = False
17+
headers: Union[Dict[str, str], None]
1818

1919

2020
# This class wraps the HTTP request creation logic
@@ -37,7 +37,6 @@ def __init__(
3737
self.api_key = config.get("api_key")
3838
self.data = data
3939
self.headers = headers or {"Content-Type": "application/json"}
40-
self.disable_request_logging = config.get("disable_request_logging")
4140
self.stream = stream
4241
self.files = files
4342

@@ -162,9 +161,6 @@ def __get_headers(self) -> Dict[Any, Any]:
162161
if not self.files and not self.data:
163162
h["Content-Type"] = "application/json"
164163

165-
if self.disable_request_logging:
166-
h["x-jigsaw-no-request-log"] = "true"
167-
168164
_headers = h.copy()
169165

170166
# Don't override Content-Type if using multipart

‎jigsawstack/search.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -226,13 +226,13 @@ def __init__(
226226
self,
227227
api_key: str,
228228
api_url: str,
229-
disable_request_logging: Union[bool, None] = False,
229+
headers: Union[Dict[str, str], None] = None,
230230
):
231-
super().__init__(api_key, api_url, disable_request_logging)
231+
super().__init__(api_key, api_url, headers)
232232
self.config = RequestConfig(
233233
api_url=api_url,
234234
api_key=api_key,
235-
disable_request_logging=disable_request_logging,
235+
headers=headers,
236236
)
237237

238238
def search(self, params: SearchParams) -> SearchResponse:
@@ -288,13 +288,13 @@ def __init__(
288288
self,
289289
api_key: str,
290290
api_url: str,
291-
disable_request_logging: Union[bool, None] = False,
291+
headers: Union[Dict[str, str], None] = None,
292292
):
293-
super().__init__(api_key, api_url, disable_request_logging)
293+
super().__init__(api_key, api_url, headers)
294294
self.config = AsyncRequestConfig(
295295
api_url=api_url,
296296
api_key=api_key,
297-
disable_request_logging=disable_request_logging,
297+
headers=headers,
298298
)
299299

300300
async def search(self, params: SearchParams) -> SearchResponse:

0 commit comments

Comments
 (0)