Skip to content

Commit 9dad1cd

Browse files
VisargDVisargD
authored
Merge pull request #413 from Portkey-AI/fix/extra-body-parameter-handling
fix: properly handle extra_body parameter in completions API
2 parents 80f7169 + 1d9c53e commit 9dad1cd

File tree

2 files changed

+68
-12
lines changed

2 files changed

+68
-12
lines changed

‎portkey_ai/api_resources/apis/chat_complete.py‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,11 @@ def stream_create( # type: ignore[return]
7474
store,
7575
**kwargs,
7676
) -> Union[ChatCompletions, Iterator[ChatCompletionChunk]]:
77-
extra_headers = kwargs.get("extra_headers", {})
77+
extra_headers = kwargs.pop("extra_headers", None)
78+
extra_query = kwargs.pop("extra_query", None)
79+
timeout = kwargs.pop("timeout", None)
80+
user_extra_body = kwargs.pop("extra_body", None) or {}
81+
merged_extra_body = {**user_extra_body, **kwargs}
7882
return self.openai_client.chat.completions.create(
7983
model=model,
8084
messages=messages,
@@ -90,7 +94,9 @@ def stream_create( # type: ignore[return]
9094
reasoning_effort=reasoning_effort,
9195
store=store,
9296
extra_headers=extra_headers,
93-
extra_body=kwargs,
97+
extra_query=extra_query,
98+
extra_body=merged_extra_body,
99+
timeout=timeout,
94100
)
95101

96102
def normal_create(
@@ -110,7 +116,11 @@ def normal_create(
110116
store,
111117
**kwargs,
112118
) -> ChatCompletions:
113-
extra_headers = kwargs.get("extra_headers", {})
119+
extra_headers = kwargs.pop("extra_headers", None)
120+
extra_query = kwargs.pop("extra_query", None)
121+
timeout = kwargs.pop("timeout", None)
122+
user_extra_body = kwargs.pop("extra_body", None) or {}
123+
merged_extra_body = {**user_extra_body, **kwargs}
114124
response = self.openai_client.with_raw_response.chat.completions.create(
115125
model=model,
116126
messages=messages,
@@ -126,7 +136,9 @@ def normal_create(
126136
reasoning_effort=reasoning_effort,
127137
store=store,
128138
extra_headers=extra_headers,
129-
extra_body=kwargs,
139+
extra_query=extra_query,
140+
extra_body=merged_extra_body,
141+
timeout=timeout,
130142
)
131143
data = ChatCompletions(**json.loads(response.text))
132144
data._headers = response.headers
@@ -472,7 +484,11 @@ async def stream_create(
472484
store,
473485
**kwargs,
474486
) -> Union[ChatCompletions, AsyncIterator[ChatCompletionChunk]]:
475-
extra_headers = kwargs.get("extra_headers", {})
487+
extra_headers = kwargs.pop("extra_headers", None)
488+
extra_query = kwargs.pop("extra_query", None)
489+
timeout = kwargs.pop("timeout", None)
490+
user_extra_body = kwargs.pop("extra_body", None) or {}
491+
merged_extra_body = {**user_extra_body, **kwargs}
476492
return await self.openai_client.chat.completions.create(
477493
model=model,
478494
messages=messages,
@@ -488,7 +504,9 @@ async def stream_create(
488504
reasoning_effort=reasoning_effort,
489505
store=store,
490506
extra_headers=extra_headers,
491-
extra_body=kwargs,
507+
extra_query=extra_query,
508+
extra_body=merged_extra_body,
509+
timeout=timeout,
492510
)
493511

494512
async def normal_create(
@@ -508,7 +526,11 @@ async def normal_create(
508526
store,
509527
**kwargs,
510528
) -> ChatCompletions:
511-
extra_headers = kwargs.get("extra_headers", {})
529+
extra_headers = kwargs.pop("extra_headers", None)
530+
extra_query = kwargs.pop("extra_query", None)
531+
timeout = kwargs.pop("timeout", None)
532+
user_extra_body = kwargs.pop("extra_body", None) or {}
533+
merged_extra_body = {**user_extra_body, **kwargs}
512534
response = await self.openai_client.with_raw_response.chat.completions.create(
513535
model=model,
514536
messages=messages,
@@ -524,7 +546,9 @@ async def normal_create(
524546
reasoning_effort=reasoning_effort,
525547
store=store,
526548
extra_headers=extra_headers,
527-
extra_body=kwargs,
549+
extra_query=extra_query,
550+
extra_body=merged_extra_body,
551+
timeout=timeout,
528552
)
529553
data = ChatCompletions(**json.loads(response.text))
530554
data._headers = response.headers

‎portkey_ai/api_resources/apis/complete.py‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,11 @@ def stream_create( # type: ignore[return]
3939
stream_options,
4040
**kwargs,
4141
) -> Union[TextCompletion, Iterator[TextCompletionChunk]]:
42+
extra_headers = kwargs.pop("extra_headers", None)
43+
extra_query = kwargs.pop("extra_query", None)
44+
timeout = kwargs.pop("timeout", None)
45+
user_extra_body = kwargs.pop("extra_body", None) or {}
46+
merged_extra_body = {**user_extra_body, **kwargs}
4247
return self.openai_client.completions.create(
4348
model=model,
4449
prompt=prompt,
@@ -58,7 +63,10 @@ def stream_create( # type: ignore[return]
5863
suffix=suffix,
5964
user=user,
6065
stream_options=stream_options,
61-
extra_body=kwargs,
66+
extra_headers=extra_headers,
67+
extra_query=extra_query,
68+
extra_body=merged_extra_body,
69+
timeout=timeout,
6270
)
6371

6472
def normal_create(
@@ -83,6 +91,11 @@ def normal_create(
8391
stream_options,
8492
**kwargs,
8593
) -> TextCompletion:
94+
extra_headers = kwargs.pop("extra_headers", None)
95+
extra_query = kwargs.pop("extra_query", None)
96+
timeout = kwargs.pop("timeout", None)
97+
user_extra_body = kwargs.pop("extra_body", None) or {}
98+
merged_extra_body = {**user_extra_body, **kwargs}
8699
response = self.openai_client.with_raw_response.completions.create(
87100
model=model,
88101
prompt=prompt,
@@ -102,7 +115,10 @@ def normal_create(
102115
suffix=suffix,
103116
user=user,
104117
stream_options=stream_options,
105-
extra_body=kwargs,
118+
extra_headers=extra_headers,
119+
extra_query=extra_query,
120+
extra_body=merged_extra_body,
121+
timeout=timeout,
106122
)
107123
data = TextCompletion(**json.loads(response.text))
108124
data._headers = response.headers
@@ -206,6 +222,11 @@ async def stream_create(
206222
stream_options,
207223
**kwargs,
208224
) -> Union[TextCompletion, AsyncIterator[TextCompletionChunk]]:
225+
extra_headers = kwargs.pop("extra_headers", None)
226+
extra_query = kwargs.pop("extra_query", None)
227+
timeout = kwargs.pop("timeout", None)
228+
user_extra_body = kwargs.pop("extra_body", None) or {}
229+
merged_extra_body = {**user_extra_body, **kwargs}
209230
return await self.openai_client.completions.create(
210231
model=model,
211232
prompt=prompt,
@@ -225,7 +246,10 @@ async def stream_create(
225246
suffix=suffix,
226247
user=user,
227248
stream_options=stream_options,
228-
extra_body=kwargs,
249+
extra_headers=extra_headers,
250+
extra_query=extra_query,
251+
extra_body=merged_extra_body,
252+
timeout=timeout,
229253
)
230254

231255
async def normal_create(
@@ -250,6 +274,11 @@ async def normal_create(
250274
stream_options,
251275
**kwargs,
252276
) -> TextCompletion:
277+
extra_headers = kwargs.pop("extra_headers", None)
278+
extra_query = kwargs.pop("extra_query", None)
279+
timeout = kwargs.pop("timeout", None)
280+
user_extra_body = kwargs.pop("extra_body", None) or {}
281+
merged_extra_body = {**user_extra_body, **kwargs}
253282
response = await self.openai_client.with_raw_response.completions.create(
254283
model=model,
255284
prompt=prompt,
@@ -269,7 +298,10 @@ async def normal_create(
269298
suffix=suffix,
270299
user=user,
271300
stream_options=stream_options,
272-
extra_body=kwargs,
301+
extra_headers=extra_headers,
302+
extra_query=extra_query,
303+
extra_body=merged_extra_body,
304+
timeout=timeout,
273305
)
274306
data = TextCompletion(**json.loads(response.text))
275307
data._headers = response.headers

0 commit comments

Comments
 (0)