GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/openai-python

Public

mirrored from https://github.com/openai/openai-pythonAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
8fe0ab87e67eeb3cc27426b50093845229520f0e

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

openai-python/tests

tests/test_client.py

2952lines · modecode

RawDownload
Latest commit unavailable.
1# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
2
3from __future__ import annotations
4
5import gc
6import os
7import sys
8import json
9import asyncio
10import inspect
11import dataclasses
12import tracemalloc
13from typing import Any, Union, TypeVar, Callable, Iterable, Iterator, Optional, Coroutine, cast
14from unittest import mock
15from typing_extensions import Literal, AsyncIterator, override
16
17import httpx
18import pytest
19from respx import MockRouter
20from pydantic import ValidationError
21from respx.models import Call as MockRequestCall
22
23from openai import OpenAI, AsyncOpenAI, OpenAIError, APIResponseValidationError
24from openai.auth import WorkloadIdentity
25from openai._types import Omit
26from openai._utils import asyncify
27from openai._models import BaseModel, FinalRequestOptions
28from openai._streaming import Stream, AsyncStream
29from openai._exceptions import APIStatusError, APITimeoutError, APIResponseValidationError
30from openai._base_client import (
31 DEFAULT_TIMEOUT,
32 HTTPX_DEFAULT_TIMEOUT,
33 BaseClient,
34 OtherPlatform,
35 DefaultHttpxClient,
36 DefaultAsyncHttpxClient,
37 get_platform,
38 make_request_options,
39)
40
41from .utils import update_env
42
43T = TypeVar("T")
44base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
45api_key = "My API Key"
46admin_api_key = "My Admin API Key"
47workload_identity: WorkloadIdentity = {
48 "client_id": "client_123",
49 "identity_provider_id": "provider_123",
50 "service_account_id": "service_account_123",
51 "provider": {
52 "get_token": lambda: "external-subject-token",
53 "token_type": "jwt",
54 },
55}
56
57
58def _get_params(client: BaseClient[Any, Any]) -> dict[str, str]:
59 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
60 url = httpx.URL(request.url)
61 return dict(url.params)
62
63
64def _low_retry_timeout(*_args: Any, **_kwargs: Any) -> float:
65 return 0.1
66
67
68def mirror_request_content(request: httpx.Request) -> httpx.Response:
69 return httpx.Response(200, content=request.content)
70
71
72# note: we can't use the httpx.MockTransport class as it consumes the request
73# body itself, which means we can't test that the body is read lazily
74class MockTransport(httpx.BaseTransport, httpx.AsyncBaseTransport):
75 def __init__(
76 self,
77 handler: Callable[[httpx.Request], httpx.Response]
78 | Callable[[httpx.Request], Coroutine[Any, Any, httpx.Response]],
79 ) -> None:
80 self.handler = handler
81
82 @override
83 def handle_request(
84 self,
85 request: httpx.Request,
86 ) -> httpx.Response:
87 assert not inspect.iscoroutinefunction(self.handler), "handler must not be a coroutine function"
88 assert inspect.isfunction(self.handler), "handler must be a function"
89 return self.handler(request)
90
91 @override
92 async def handle_async_request(
93 self,
94 request: httpx.Request,
95 ) -> httpx.Response:
96 assert inspect.iscoroutinefunction(self.handler), "handler must be a coroutine function"
97 return await self.handler(request)
98
99
100@dataclasses.dataclass
101class Counter:
102 value: int = 0
103
104
105def _make_sync_iterator(iterable: Iterable[T], counter: Optional[Counter] = None) -> Iterator[T]:
106 for item in iterable:
107 if counter:
108 counter.value += 1
109 yield item
110
111
112async def _make_async_iterator(iterable: Iterable[T], counter: Optional[Counter] = None) -> AsyncIterator[T]:
113 for item in iterable:
114 if counter:
115 counter.value += 1
116 yield item
117
118
119def _get_open_connections(client: OpenAI | AsyncOpenAI) -> int:
120 transport = client._client._transport
121 assert isinstance(transport, httpx.HTTPTransport) or isinstance(transport, httpx.AsyncHTTPTransport)
122
123 pool = transport._pool
124 return len(pool._requests)
125
126
127class TestOpenAI:
128 @pytest.mark.respx(base_url=base_url)
129 def test_raw_response(self, respx_mock: MockRouter, client: OpenAI) -> None:
130 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
131
132 response = client.post("/foo", cast_to=httpx.Response)
133 assert response.status_code == 200
134 assert isinstance(response, httpx.Response)
135 assert response.json() == {"foo": "bar"}
136
137 @pytest.mark.respx(base_url=base_url)
138 def test_raw_response_for_binary(self, respx_mock: MockRouter, client: OpenAI) -> None:
139 respx_mock.post("/foo").mock(
140 return_value=httpx.Response(200, headers={"Content-Type": "application/binary"}, content='{"foo": "bar"}')
141 )
142
143 response = client.post("/foo", cast_to=httpx.Response)
144 assert response.status_code == 200
145 assert isinstance(response, httpx.Response)
146 assert response.json() == {"foo": "bar"}
147
148 def test_copy(self, client: OpenAI) -> None:
149 copied = client.copy()
150 assert id(copied) != id(client)
151
152 copied = client.copy(api_key="another My API Key")
153 assert copied.api_key == "another My API Key"
154 assert client.api_key == "My API Key"
155
156 copied = client.copy(admin_api_key="another My Admin API Key")
157 assert copied.admin_api_key == "another My Admin API Key"
158 assert client.admin_api_key == "My Admin API Key"
159
160 def test_copy_default_options(self, client: OpenAI) -> None:
161 # options that have a default are overridden correctly
162 copied = client.copy(max_retries=7)
163 assert copied.max_retries == 7
164 assert client.max_retries == 2
165
166 copied2 = copied.copy(max_retries=6)
167 assert copied2.max_retries == 6
168 assert copied.max_retries == 7
169
170 # timeout
171 assert isinstance(client.timeout, httpx.Timeout)
172 copied = client.copy(timeout=None)
173 assert copied.timeout is None
174 assert isinstance(client.timeout, httpx.Timeout)
175
176 def test_copy_default_headers(self) -> None:
177 client = OpenAI(
178 base_url=base_url,
179 api_key=api_key,
180 admin_api_key=admin_api_key,
181 _strict_response_validation=True,
182 default_headers={"X-Foo": "bar"},
183 )
184 assert client.default_headers["X-Foo"] == "bar"
185
186 # does not override the already given value when not specified
187 copied = client.copy()
188 assert copied.default_headers["X-Foo"] == "bar"
189
190 # merges already given headers
191 copied = client.copy(default_headers={"X-Bar": "stainless"})
192 assert copied.default_headers["X-Foo"] == "bar"
193 assert copied.default_headers["X-Bar"] == "stainless"
194
195 # uses new values for any already given headers
196 copied = client.copy(default_headers={"X-Foo": "stainless"})
197 assert copied.default_headers["X-Foo"] == "stainless"
198
199 # set_default_headers
200
201 # completely overrides already set values
202 copied = client.copy(set_default_headers={})
203 assert copied.default_headers.get("X-Foo") is None
204
205 copied = client.copy(set_default_headers={"X-Bar": "Robert"})
206 assert copied.default_headers["X-Bar"] == "Robert"
207
208 with pytest.raises(
209 ValueError,
210 match="`default_headers` and `set_default_headers` arguments are mutually exclusive",
211 ):
212 client.copy(set_default_headers={}, default_headers={"X-Foo": "Bar"})
213 client.close()
214
215 def test_copy_default_query(self) -> None:
216 client = OpenAI(
217 base_url=base_url,
218 api_key=api_key,
219 admin_api_key=admin_api_key,
220 _strict_response_validation=True,
221 default_query={"foo": "bar"},
222 )
223 assert _get_params(client)["foo"] == "bar"
224
225 # does not override the already given value when not specified
226 copied = client.copy()
227 assert _get_params(copied)["foo"] == "bar"
228
229 # merges already given params
230 copied = client.copy(default_query={"bar": "stainless"})
231 params = _get_params(copied)
232 assert params["foo"] == "bar"
233 assert params["bar"] == "stainless"
234
235 # uses new values for any already given headers
236 copied = client.copy(default_query={"foo": "stainless"})
237 assert _get_params(copied)["foo"] == "stainless"
238
239 # set_default_query
240
241 # completely overrides already set values
242 copied = client.copy(set_default_query={})
243 assert _get_params(copied) == {}
244
245 copied = client.copy(set_default_query={"bar": "Robert"})
246 assert _get_params(copied)["bar"] == "Robert"
247
248 with pytest.raises(
249 ValueError,
250 # TODO: update
251 match="`default_query` and `set_default_query` arguments are mutually exclusive",
252 ):
253 client.copy(set_default_query={}, default_query={"foo": "Bar"})
254
255 client.close()
256
257 def test_copy_signature(self, client: OpenAI) -> None:
258 # ensure the same parameters that can be passed to the client are defined in the `.copy()` method
259 init_signature = inspect.signature(
260 # mypy doesn't like that we access the `__init__` property.
261 client.__init__, # type: ignore[misc]
262 )
263 copy_signature = inspect.signature(client.copy)
264 exclude_params = {"transport", "proxies", "_strict_response_validation"}
265
266 for name in init_signature.parameters.keys():
267 if name in exclude_params:
268 continue
269
270 copy_param = copy_signature.parameters.get(name)
271 assert copy_param is not None, f"copy() signature is missing the {name} param"
272
273 @pytest.mark.skipif(sys.version_info >= (3, 10), reason="fails because of a memory leak that started from 3.12")
274 def test_copy_build_request(self, client: OpenAI) -> None:
275 options = FinalRequestOptions(method="get", url="/foo")
276
277 def build_request(options: FinalRequestOptions) -> None:
278 client_copy = client.copy()
279 client_copy._build_request(options)
280
281 # ensure that the machinery is warmed up before tracing starts.
282 build_request(options)
283 gc.collect()
284
285 tracemalloc.start(1000)
286
287 snapshot_before = tracemalloc.take_snapshot()
288
289 ITERATIONS = 10
290 for _ in range(ITERATIONS):
291 build_request(options)
292
293 gc.collect()
294 snapshot_after = tracemalloc.take_snapshot()
295
296 tracemalloc.stop()
297
298 def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
299 if diff.count == 0:
300 # Avoid false positives by considering only leaks (i.e. allocations that persist).
301 return
302
303 if diff.count % ITERATIONS != 0:
304 # Avoid false positives by considering only leaks that appear per iteration.
305 return
306
307 for frame in diff.traceback:
308 if any(
309 frame.filename.endswith(fragment)
310 for fragment in [
311 # to_raw_response_wrapper leaks through the @functools.wraps() decorator.
312 #
313 # removing the decorator fixes the leak for reasons we don't understand.
314 "openai/_legacy_response.py",
315 "openai/_response.py",
316 # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
317 "openai/_compat.py",
318 # Standard library leaks we don't care about.
319 "/logging/__init__.py",
320 ]
321 ):
322 return
323
324 leaks.append(diff)
325
326 leaks: list[tracemalloc.StatisticDiff] = []
327 for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
328 add_leak(leaks, diff)
329 if leaks:
330 for leak in leaks:
331 print("MEMORY LEAK:", leak)
332 for frame in leak.traceback:
333 print(frame)
334 raise AssertionError()
335
336 def test_request_timeout(self, client: OpenAI) -> None:
337 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
338 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
339 assert timeout == DEFAULT_TIMEOUT
340
341 request = client._build_request(FinalRequestOptions(method="get", url="/foo", timeout=httpx.Timeout(100.0)))
342 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
343 assert timeout == httpx.Timeout(100.0)
344
345 def test_client_timeout_option(self) -> None:
346 client = OpenAI(
347 base_url=base_url,
348 api_key=api_key,
349 admin_api_key=admin_api_key,
350 _strict_response_validation=True,
351 timeout=httpx.Timeout(0),
352 )
353
354 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
355 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
356 assert timeout == httpx.Timeout(0)
357
358 client.close()
359
360 def test_http_client_timeout_option(self) -> None:
361 # custom timeout given to the httpx client should be used
362 with httpx.Client(timeout=None) as http_client:
363 client = OpenAI(
364 base_url=base_url,
365 api_key=api_key,
366 admin_api_key=admin_api_key,
367 _strict_response_validation=True,
368 http_client=http_client,
369 )
370
371 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
372 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
373 assert timeout == httpx.Timeout(None)
374
375 client.close()
376
377 # no timeout given to the httpx client should not use the httpx default
378 with httpx.Client() as http_client:
379 client = OpenAI(
380 base_url=base_url,
381 api_key=api_key,
382 admin_api_key=admin_api_key,
383 _strict_response_validation=True,
384 http_client=http_client,
385 )
386
387 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
388 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
389 assert timeout == DEFAULT_TIMEOUT
390
391 client.close()
392
393 # explicitly passing the default timeout currently results in it being ignored
394 with httpx.Client(timeout=HTTPX_DEFAULT_TIMEOUT) as http_client:
395 client = OpenAI(
396 base_url=base_url,
397 api_key=api_key,
398 admin_api_key=admin_api_key,
399 _strict_response_validation=True,
400 http_client=http_client,
401 )
402
403 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
404 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
405 assert timeout == DEFAULT_TIMEOUT # our default
406
407 client.close()
408
409 async def test_invalid_http_client(self) -> None:
410 with pytest.raises(TypeError, match="Invalid `http_client` arg"):
411 async with httpx.AsyncClient() as http_client:
412 OpenAI(
413 base_url=base_url,
414 api_key=api_key,
415 admin_api_key=admin_api_key,
416 _strict_response_validation=True,
417 http_client=cast(Any, http_client),
418 )
419
420 def test_default_headers_option(self) -> None:
421 test_client = OpenAI(
422 base_url=base_url,
423 api_key=api_key,
424 admin_api_key=admin_api_key,
425 _strict_response_validation=True,
426 default_headers={"X-Foo": "bar"},
427 )
428 request = test_client._build_request(FinalRequestOptions(method="get", url="/foo"))
429 assert request.headers.get("x-foo") == "bar"
430 assert request.headers.get("x-stainless-lang") == "python"
431
432 test_client2 = OpenAI(
433 base_url=base_url,
434 api_key=api_key,
435 admin_api_key=admin_api_key,
436 _strict_response_validation=True,
437 default_headers={
438 "X-Foo": "stainless",
439 "X-Stainless-Lang": "my-overriding-header",
440 },
441 )
442 request = test_client2._build_request(FinalRequestOptions(method="get", url="/foo"))
443 assert request.headers.get("x-foo") == "stainless"
444 assert request.headers.get("x-stainless-lang") == "my-overriding-header"
445
446 test_client.close()
447 test_client2.close()
448
449 def test_validate_headers(self) -> None:
450 client = OpenAI(
451 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
452 )
453 options = client._prepare_options(FinalRequestOptions(method="get", url="/foo"))
454 request = client._build_request(options)
455
456 assert request.headers.get("Authorization") == f"Bearer {api_key}"
457
458 admin_request = client._build_request(
459 FinalRequestOptions(
460 method="get",
461 url="/organization/projects",
462 security={"admin_api_key_auth": True},
463 )
464 )
465 assert admin_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
466
467 with update_env(**{"OPENAI_API_KEY": Omit()}):
468 admin_only = OpenAI(
469 base_url=base_url,
470 api_key=None,
471 admin_api_key=admin_api_key,
472 _strict_response_validation=True,
473 )
474 admin_only_request = admin_only._build_request(
475 FinalRequestOptions(
476 method="get",
477 url="/organization/projects",
478 security={"admin_api_key_auth": True},
479 )
480 )
481 assert admin_only_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
482
483 with pytest.raises(
484 TypeError,
485 match="Could not resolve authentication method",
486 ):
487 admin_only._build_request(
488 FinalRequestOptions(
489 method="post",
490 url="/responses",
491 security={"bearer_auth": True},
492 )
493 )
494
495 with update_env(
496 **{
497 "OPENAI_API_KEY": Omit(),
498 "OPENAI_ADMIN_KEY": Omit(),
499 }
500 ):
501 no_credentials = OpenAI(
502 base_url=base_url,
503 api_key=None,
504 admin_api_key=None,
505 _enforce_credentials=False,
506 _strict_response_validation=True,
507 )
508 lowercase_auth_request = no_credentials._build_request(
509 FinalRequestOptions(method="get", url="/foo", headers={"authorization": "Bearer custom"})
510 )
511 assert lowercase_auth_request.headers.get("Authorization") == "Bearer custom"
512
513 omitted_auth_request = no_credentials._build_request(
514 FinalRequestOptions(method="get", url="/foo", headers={"authorization": Omit()})
515 )
516 assert "Authorization" not in omitted_auth_request.headers
517
518 with update_env(
519 **{
520 "OPENAI_API_KEY": Omit(),
521 "OPENAI_ADMIN_KEY": Omit(),
522 }
523 ):
524 with pytest.raises(OpenAIError, match="Missing credentials"):
525 OpenAI(base_url=base_url, api_key=None, admin_api_key=None, _strict_response_validation=True)
526
527 @pytest.mark.respx(base_url=base_url)
528 def test_api_key_provider_preserves_admin_auth(self, respx_mock: MockRouter) -> None:
529 respx_mock.get("/organization/projects").mock(return_value=httpx.Response(200, json={"ok": True}))
530
531 provider_called = False
532
533 def api_key_provider() -> str:
534 nonlocal provider_called
535 provider_called = True
536 return "dynamic-api-key"
537
538 client = OpenAI(base_url=base_url, api_key=api_key_provider, admin_api_key=admin_api_key)
539 response = client.get(
540 "/organization/projects",
541 cast_to=httpx.Response,
542 options={"security": {"admin_api_key_auth": True}},
543 )
544
545 assert response.request.headers.get("Authorization") == f"Bearer {admin_api_key}"
546 assert provider_called is False
547
548 def test_api_key_provider_does_not_fill_admin_auth(self) -> None:
549 provider_called = False
550
551 def api_key_provider() -> str:
552 nonlocal provider_called
553 provider_called = True
554 return "dynamic-api-key"
555
556 with update_env(OPENAI_ADMIN_KEY=Omit()):
557 client = OpenAI(base_url=base_url, api_key=api_key_provider, admin_api_key=None)
558 with pytest.raises(TypeError, match="Could not resolve authentication method"):
559 client.get(
560 "/organization/projects",
561 cast_to=httpx.Response,
562 options={"security": {"admin_api_key_auth": True}},
563 )
564
565 assert provider_called is False
566
567 @pytest.mark.respx(base_url=base_url)
568 def test_workload_identity_preserves_admin_auth(self, respx_mock: MockRouter) -> None:
569 respx_mock.get("/organization/projects").mock(return_value=httpx.Response(200, json={"ok": True}))
570
571 client = OpenAI(base_url=base_url, workload_identity=workload_identity, admin_api_key=admin_api_key)
572 response = client.get(
573 "/organization/projects",
574 cast_to=httpx.Response,
575 options={"security": {"admin_api_key_auth": True}},
576 )
577
578 assert response.request.headers.get("Authorization") == f"Bearer {admin_api_key}"
579
580 def test_workload_identity_is_mutually_exclusive_with_api_key(self) -> None:
581 with pytest.raises(
582 OpenAIError,
583 match="The `api_key` and `workload_identity` arguments are mutually exclusive",
584 ):
585 OpenAI(
586 base_url=base_url,
587 api_key=api_key,
588 workload_identity=workload_identity, # type: ignore[reportArgumentType]
589 organization="org_123",
590 _strict_response_validation=True,
591 )
592
593 def test_default_query_option(self) -> None:
594 client = OpenAI(
595 base_url=base_url,
596 api_key=api_key,
597 admin_api_key=admin_api_key,
598 _strict_response_validation=True,
599 default_query={"query_param": "bar"},
600 )
601 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
602 url = httpx.URL(request.url)
603 assert dict(url.params) == {"query_param": "bar"}
604
605 request = client._build_request(
606 FinalRequestOptions(
607 method="get",
608 url="/foo",
609 params={"foo": "baz", "query_param": "overridden"},
610 )
611 )
612 url = httpx.URL(request.url)
613 assert dict(url.params) == {"foo": "baz", "query_param": "overridden"}
614
615 client.close()
616
617 def test_hardcoded_query_params_in_url(self, client: OpenAI) -> None:
618 request = client._build_request(FinalRequestOptions(method="get", url="/foo?beta=true"))
619 url = httpx.URL(request.url)
620 assert dict(url.params) == {"beta": "true"}
621
622 request = client._build_request(
623 FinalRequestOptions(
624 method="get",
625 url="/foo?beta=true",
626 params={"limit": "10", "page": "abc"},
627 )
628 )
629 url = httpx.URL(request.url)
630 assert dict(url.params) == {"beta": "true", "limit": "10", "page": "abc"}
631
632 request = client._build_request(
633 FinalRequestOptions(
634 method="get",
635 url="/files/a%2Fb?beta=true",
636 params={"limit": "10"},
637 )
638 )
639 assert request.url.raw_path == b"/files/a%2Fb?beta=true&limit=10"
640
641 def test_request_extra_json(self, client: OpenAI) -> None:
642 request = client._build_request(
643 FinalRequestOptions(
644 method="post",
645 url="/foo",
646 json_data={"foo": "bar"},
647 extra_json={"baz": False},
648 ),
649 )
650 data = json.loads(request.content.decode("utf-8"))
651 assert data == {"foo": "bar", "baz": False}
652
653 request = client._build_request(
654 FinalRequestOptions(
655 method="post",
656 url="/foo",
657 extra_json={"baz": False},
658 ),
659 )
660 data = json.loads(request.content.decode("utf-8"))
661 assert data == {"baz": False}
662
663 # `extra_json` takes priority over `json_data` when keys clash
664 request = client._build_request(
665 FinalRequestOptions(
666 method="post",
667 url="/foo",
668 json_data={"foo": "bar", "baz": True},
669 extra_json={"baz": None},
670 ),
671 )
672 data = json.loads(request.content.decode("utf-8"))
673 assert data == {"foo": "bar", "baz": None}
674
675 def test_request_extra_headers(self, client: OpenAI) -> None:
676 request = client._build_request(
677 FinalRequestOptions(
678 method="post",
679 url="/foo",
680 **make_request_options(extra_headers={"X-Foo": "Foo"}),
681 ),
682 )
683 assert request.headers.get("X-Foo") == "Foo"
684
685 # `extra_headers` takes priority over `default_headers` when keys clash
686 request = client.with_options(default_headers={"X-Bar": "true"})._build_request(
687 FinalRequestOptions(
688 method="post",
689 url="/foo",
690 **make_request_options(
691 extra_headers={"X-Bar": "false"},
692 ),
693 ),
694 )
695 assert request.headers.get("X-Bar") == "false"
696
697 def test_request_extra_query(self, client: OpenAI) -> None:
698 request = client._build_request(
699 FinalRequestOptions(
700 method="post",
701 url="/foo",
702 **make_request_options(
703 extra_query={"my_query_param": "Foo"},
704 ),
705 ),
706 )
707 params = dict(request.url.params)
708 assert params == {"my_query_param": "Foo"}
709
710 # if both `query` and `extra_query` are given, they are merged
711 request = client._build_request(
712 FinalRequestOptions(
713 method="post",
714 url="/foo",
715 **make_request_options(
716 query={"bar": "1"},
717 extra_query={"foo": "2"},
718 ),
719 ),
720 )
721 params = dict(request.url.params)
722 assert params == {"bar": "1", "foo": "2"}
723
724 # `extra_query` takes priority over `query` when keys clash
725 request = client._build_request(
726 FinalRequestOptions(
727 method="post",
728 url="/foo",
729 **make_request_options(
730 query={"foo": "1"},
731 extra_query={"foo": "2"},
732 ),
733 ),
734 )
735 params = dict(request.url.params)
736 assert params == {"foo": "2"}
737
738 def test_multipart_repeating_array(self, client: OpenAI) -> None:
739 request = client._build_request(
740 FinalRequestOptions.construct(
741 method="post",
742 url="/foo",
743 headers={"Content-Type": "multipart/form-data; boundary=6b7ba517decee4a450543ea6ae821c82"},
744 json_data={"array": ["foo", "bar"]},
745 files=[("foo.txt", b"hello world")],
746 )
747 )
748
749 assert request.read().split(b"\r\n") == [
750 b"--6b7ba517decee4a450543ea6ae821c82",
751 b'Content-Disposition: form-data; name="array[]"',
752 b"",
753 b"foo",
754 b"--6b7ba517decee4a450543ea6ae821c82",
755 b'Content-Disposition: form-data; name="array[]"',
756 b"",
757 b"bar",
758 b"--6b7ba517decee4a450543ea6ae821c82",
759 b'Content-Disposition: form-data; name="foo.txt"; filename="upload"',
760 b"Content-Type: application/octet-stream",
761 b"",
762 b"hello world",
763 b"--6b7ba517decee4a450543ea6ae821c82--",
764 b"",
765 ]
766
767 @pytest.mark.respx(base_url=base_url)
768 def test_binary_content_upload(self, respx_mock: MockRouter, client: OpenAI) -> None:
769 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
770
771 file_content = b"Hello, this is a test file."
772
773 response = client.post(
774 "/upload",
775 content=file_content,
776 cast_to=httpx.Response,
777 options={"headers": {"Content-Type": "application/octet-stream"}},
778 )
779
780 assert response.status_code == 200
781 assert response.request.headers["Content-Type"] == "application/octet-stream"
782 assert response.content == file_content
783
784 def test_binary_content_upload_with_iterator(self) -> None:
785 file_content = b"Hello, this is a test file."
786 counter = Counter()
787 iterator = _make_sync_iterator([file_content], counter=counter)
788
789 def mock_handler(request: httpx.Request) -> httpx.Response:
790 assert counter.value == 0, "the request body should not have been read"
791 return httpx.Response(200, content=request.read())
792
793 with OpenAI(
794 base_url=base_url,
795 api_key=api_key,
796 admin_api_key=admin_api_key,
797 _strict_response_validation=True,
798 http_client=httpx.Client(transport=MockTransport(handler=mock_handler)),
799 ) as client:
800 response = client.post(
801 "/upload",
802 content=iterator,
803 cast_to=httpx.Response,
804 options={"headers": {"Content-Type": "application/octet-stream"}},
805 )
806
807 assert response.status_code == 200
808 assert response.request.headers["Content-Type"] == "application/octet-stream"
809 assert response.content == file_content
810 assert counter.value == 1
811
812 @pytest.mark.respx(base_url=base_url)
813 def test_binary_content_upload_with_body_is_deprecated(self, respx_mock: MockRouter, client: OpenAI) -> None:
814 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
815
816 file_content = b"Hello, this is a test file."
817
818 with pytest.deprecated_call(
819 match="Passing raw bytes as `body` is deprecated and will be removed in a future version. Please pass raw bytes via the `content` parameter instead."
820 ):
821 response = client.post(
822 "/upload",
823 body=file_content,
824 cast_to=httpx.Response,
825 options={"headers": {"Content-Type": "application/octet-stream"}},
826 )
827
828 assert response.status_code == 200
829 assert response.request.headers["Content-Type"] == "application/octet-stream"
830 assert response.content == file_content
831
832 @pytest.mark.respx(base_url=base_url)
833 def test_basic_union_response(self, respx_mock: MockRouter, client: OpenAI) -> None:
834 class Model1(BaseModel):
835 name: str
836
837 class Model2(BaseModel):
838 foo: str
839
840 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
841
842 response = client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
843 assert isinstance(response, Model2)
844 assert response.foo == "bar"
845
846 @pytest.mark.respx(base_url=base_url)
847 def test_union_response_different_types(self, respx_mock: MockRouter, client: OpenAI) -> None:
848 """Union of objects with the same field name using a different type"""
849
850 class Model1(BaseModel):
851 foo: int
852
853 class Model2(BaseModel):
854 foo: str
855
856 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
857
858 response = client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
859 assert isinstance(response, Model2)
860 assert response.foo == "bar"
861
862 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": 1}))
863
864 response = client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
865 assert isinstance(response, Model1)
866 assert response.foo == 1
867
868 @pytest.mark.respx(base_url=base_url)
869 def test_non_application_json_content_type_for_json_data(self, respx_mock: MockRouter, client: OpenAI) -> None:
870 """
871 Response that sets Content-Type to something other than application/json but returns json data
872 """
873
874 class Model(BaseModel):
875 foo: int
876
877 respx_mock.get("/foo").mock(
878 return_value=httpx.Response(
879 200,
880 content=json.dumps({"foo": 2}),
881 headers={"Content-Type": "application/text"},
882 )
883 )
884
885 response = client.get("/foo", cast_to=Model)
886 assert isinstance(response, Model)
887 assert response.foo == 2
888
889 def test_base_url_setter(self) -> None:
890 client = OpenAI(
891 base_url="https://example.com/from_init",
892 api_key=api_key,
893 admin_api_key=admin_api_key,
894 _strict_response_validation=True,
895 )
896 assert client.base_url == "https://example.com/from_init/"
897
898 client.base_url = "https://example.com/from_setter" # type: ignore[assignment]
899
900 assert client.base_url == "https://example.com/from_setter/"
901
902 client.close()
903
904 def test_base_url_env(self) -> None:
905 with update_env(OPENAI_BASE_URL="http://localhost:5000/from/env"):
906 client = OpenAI(api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True)
907 assert client.base_url == "http://localhost:5000/from/env/"
908
909 @pytest.mark.parametrize(
910 "client",
911 [
912 OpenAI(
913 base_url="http://localhost:5000/custom/path/",
914 api_key=api_key,
915 admin_api_key=admin_api_key,
916 _strict_response_validation=True,
917 ),
918 OpenAI(
919 base_url="http://localhost:5000/custom/path/",
920 api_key=api_key,
921 admin_api_key=admin_api_key,
922 _strict_response_validation=True,
923 http_client=httpx.Client(),
924 ),
925 ],
926 ids=["standard", "custom http client"],
927 )
928 def test_base_url_trailing_slash(self, client: OpenAI) -> None:
929 request = client._build_request(
930 FinalRequestOptions(
931 method="post",
932 url="/foo",
933 json_data={"foo": "bar"},
934 ),
935 )
936 assert request.url == "http://localhost:5000/custom/path/foo"
937 client.close()
938
939 @pytest.mark.parametrize(
940 "client",
941 [
942 OpenAI(
943 base_url="http://localhost:5000/custom/path/",
944 api_key=api_key,
945 admin_api_key=admin_api_key,
946 _strict_response_validation=True,
947 ),
948 OpenAI(
949 base_url="http://localhost:5000/custom/path/",
950 api_key=api_key,
951 admin_api_key=admin_api_key,
952 _strict_response_validation=True,
953 http_client=httpx.Client(),
954 ),
955 ],
956 ids=["standard", "custom http client"],
957 )
958 def test_base_url_no_trailing_slash(self, client: OpenAI) -> None:
959 request = client._build_request(
960 FinalRequestOptions(
961 method="post",
962 url="/foo",
963 json_data={"foo": "bar"},
964 ),
965 )
966 assert request.url == "http://localhost:5000/custom/path/foo"
967 client.close()
968
969 @pytest.mark.parametrize(
970 "client",
971 [
972 OpenAI(
973 base_url="http://localhost:5000/custom/path/",
974 api_key=api_key,
975 admin_api_key=admin_api_key,
976 _strict_response_validation=True,
977 ),
978 OpenAI(
979 base_url="http://localhost:5000/custom/path/",
980 api_key=api_key,
981 admin_api_key=admin_api_key,
982 _strict_response_validation=True,
983 http_client=httpx.Client(),
984 ),
985 ],
986 ids=["standard", "custom http client"],
987 )
988 def test_absolute_request_url(self, client: OpenAI) -> None:
989 request = client._build_request(
990 FinalRequestOptions(
991 method="post",
992 url="https://myapi.com/foo",
993 json_data={"foo": "bar"},
994 ),
995 )
996 assert request.url == "https://myapi.com/foo"
997 client.close()
998
999 def test_copied_client_does_not_close_http(self) -> None:
1000 test_client = OpenAI(
1001 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
1002 )
1003 assert not test_client.is_closed()
1004
1005 copied = test_client.copy()
1006 assert copied is not test_client
1007
1008 del copied
1009
1010 assert not test_client.is_closed()
1011
1012 def test_client_context_manager(self) -> None:
1013 test_client = OpenAI(
1014 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
1015 )
1016 with test_client as c2:
1017 assert c2 is test_client
1018 assert not c2.is_closed()
1019 assert not test_client.is_closed()
1020 assert test_client.is_closed()
1021
1022 @pytest.mark.respx(base_url=base_url)
1023 def test_client_response_validation_error(self, respx_mock: MockRouter, client: OpenAI) -> None:
1024 class Model(BaseModel):
1025 foo: str
1026
1027 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": {"invalid": True}}))
1028
1029 with pytest.raises(APIResponseValidationError) as exc:
1030 client.get("/foo", cast_to=Model)
1031
1032 assert isinstance(exc.value.__cause__, ValidationError)
1033
1034 def test_client_max_retries_validation(self) -> None:
1035 with pytest.raises(TypeError, match=r"max_retries cannot be None"):
1036 OpenAI(
1037 base_url=base_url,
1038 api_key=api_key,
1039 admin_api_key=admin_api_key,
1040 _strict_response_validation=True,
1041 max_retries=cast(Any, None),
1042 )
1043
1044 @pytest.mark.respx(base_url=base_url)
1045 def test_default_stream_cls(self, respx_mock: MockRouter, client: OpenAI) -> None:
1046 class Model(BaseModel):
1047 name: str
1048
1049 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
1050
1051 stream = client.post("/foo", cast_to=Model, stream=True, stream_cls=Stream[Model])
1052 assert isinstance(stream, Stream)
1053 stream.response.close()
1054
1055 @pytest.mark.respx(base_url=base_url)
1056 def test_received_text_for_expected_json(self, respx_mock: MockRouter) -> None:
1057 class Model(BaseModel):
1058 name: str
1059
1060 respx_mock.get("/foo").mock(return_value=httpx.Response(200, text="my-custom-format"))
1061
1062 strict_client = OpenAI(
1063 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
1064 )
1065
1066 with pytest.raises(APIResponseValidationError):
1067 strict_client.get("/foo", cast_to=Model)
1068
1069 non_strict_client = OpenAI(
1070 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=False
1071 )
1072
1073 response = non_strict_client.get("/foo", cast_to=Model)
1074 assert isinstance(response, str) # type: ignore[unreachable]
1075
1076 strict_client.close()
1077 non_strict_client.close()
1078
1079 @pytest.mark.parametrize(
1080 "remaining_retries,retry_after,timeout",
1081 [
1082 [3, "20", 20],
1083 [3, "0", 0.5],
1084 [3, "-10", 0.5],
1085 [3, "60", 60],
1086 [3, "61", 0.5],
1087 [3, "Fri, 29 Sep 2023 16:26:57 GMT", 20],
1088 [3, "Fri, 29 Sep 2023 16:26:37 GMT", 0.5],
1089 [3, "Fri, 29 Sep 2023 16:26:27 GMT", 0.5],
1090 [3, "Fri, 29 Sep 2023 16:27:37 GMT", 60],
1091 [3, "Fri, 29 Sep 2023 16:27:38 GMT", 0.5],
1092 [3, "99999999999999999999999999999999999", 0.5],
1093 [3, "Zun, 29 Sep 2023 16:26:27 GMT", 0.5],
1094 [3, "", 0.5],
1095 [2, "", 0.5 * 2.0],
1096 [1, "", 0.5 * 4.0],
1097 [-1100, "", 8], # test large number potentially overflowing
1098 ],
1099 )
1100 @mock.patch("time.time", mock.MagicMock(return_value=1696004797))
1101 def test_parse_retry_after_header(
1102 self, remaining_retries: int, retry_after: str, timeout: float, client: OpenAI
1103 ) -> None:
1104 headers = httpx.Headers({"retry-after": retry_after})
1105 options = FinalRequestOptions(method="get", url="/foo", max_retries=3)
1106 calculated = client._calculate_retry_timeout(remaining_retries, options, headers)
1107 assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType]
1108
1109 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1110 @pytest.mark.respx(base_url=base_url)
1111 def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: OpenAI) -> None:
1112 respx_mock.post("/chat/completions").mock(side_effect=httpx.TimeoutException("Test timeout error"))
1113
1114 with pytest.raises(APITimeoutError):
1115 client.chat.completions.with_streaming_response.create(
1116 messages=[
1117 {
1118 "content": "string",
1119 "role": "developer",
1120 }
1121 ],
1122 model="gpt-5.4",
1123 ).__enter__()
1124
1125 assert _get_open_connections(client) == 0
1126
1127 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1128 @pytest.mark.respx(base_url=base_url)
1129 def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: OpenAI) -> None:
1130 respx_mock.post("/chat/completions").mock(return_value=httpx.Response(500))
1131
1132 with pytest.raises(APIStatusError):
1133 client.chat.completions.with_streaming_response.create(
1134 messages=[
1135 {
1136 "content": "string",
1137 "role": "developer",
1138 }
1139 ],
1140 model="gpt-5.4",
1141 ).__enter__()
1142 assert _get_open_connections(client) == 0
1143
1144 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1145 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1146 @pytest.mark.respx(base_url=base_url)
1147 @pytest.mark.parametrize("failure_mode", ["status", "exception"])
1148 def test_retries_taken(
1149 self,
1150 client: OpenAI,
1151 failures_before_success: int,
1152 failure_mode: Literal["status", "exception"],
1153 respx_mock: MockRouter,
1154 ) -> None:
1155 client = client.with_options(max_retries=4)
1156
1157 nb_retries = 0
1158
1159 def retry_handler(_request: httpx.Request) -> httpx.Response:
1160 nonlocal nb_retries
1161 if nb_retries < failures_before_success:
1162 nb_retries += 1
1163 if failure_mode == "exception":
1164 raise RuntimeError("oops")
1165 return httpx.Response(500)
1166 return httpx.Response(200)
1167
1168 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1169
1170 response = client.chat.completions.with_raw_response.create(
1171 messages=[
1172 {
1173 "content": "string",
1174 "role": "developer",
1175 }
1176 ],
1177 model="gpt-5.4",
1178 )
1179
1180 assert response.retries_taken == failures_before_success
1181 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
1182
1183 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1184 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1185 @pytest.mark.respx(base_url=base_url)
1186 def test_omit_retry_count_header(
1187 self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter
1188 ) -> None:
1189 client = client.with_options(max_retries=4)
1190
1191 nb_retries = 0
1192
1193 def retry_handler(_request: httpx.Request) -> httpx.Response:
1194 nonlocal nb_retries
1195 if nb_retries < failures_before_success:
1196 nb_retries += 1
1197 return httpx.Response(500)
1198 return httpx.Response(200)
1199
1200 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1201
1202 response = client.chat.completions.with_raw_response.create(
1203 messages=[
1204 {
1205 "content": "string",
1206 "role": "developer",
1207 }
1208 ],
1209 model="gpt-5.4",
1210 extra_headers={"x-stainless-retry-count": Omit()},
1211 )
1212
1213 assert len(response.http_request.headers.get_list("x-stainless-retry-count")) == 0
1214
1215 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1216 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1217 @pytest.mark.respx(base_url=base_url)
1218 def test_overwrite_retry_count_header(
1219 self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter
1220 ) -> None:
1221 client = client.with_options(max_retries=4)
1222
1223 nb_retries = 0
1224
1225 def retry_handler(_request: httpx.Request) -> httpx.Response:
1226 nonlocal nb_retries
1227 if nb_retries < failures_before_success:
1228 nb_retries += 1
1229 return httpx.Response(500)
1230 return httpx.Response(200)
1231
1232 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1233
1234 response = client.chat.completions.with_raw_response.create(
1235 messages=[
1236 {
1237 "content": "string",
1238 "role": "developer",
1239 }
1240 ],
1241 model="gpt-5.4",
1242 extra_headers={"x-stainless-retry-count": "42"},
1243 )
1244
1245 assert response.http_request.headers.get("x-stainless-retry-count") == "42"
1246
1247 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1248 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1249 @pytest.mark.respx(base_url=base_url)
1250 def test_retries_taken_new_response_class(
1251 self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter
1252 ) -> None:
1253 client = client.with_options(max_retries=4)
1254
1255 nb_retries = 0
1256
1257 def retry_handler(_request: httpx.Request) -> httpx.Response:
1258 nonlocal nb_retries
1259 if nb_retries < failures_before_success:
1260 nb_retries += 1
1261 return httpx.Response(500)
1262 return httpx.Response(200)
1263
1264 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1265
1266 with client.chat.completions.with_streaming_response.create(
1267 messages=[
1268 {
1269 "content": "string",
1270 "role": "developer",
1271 }
1272 ],
1273 model="gpt-5.4",
1274 ) as response:
1275 assert response.retries_taken == failures_before_success
1276 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
1277
1278 def test_proxy_environment_variables(self, monkeypatch: pytest.MonkeyPatch) -> None:
1279 # Test that the proxy environment variables are set correctly
1280 monkeypatch.setenv("HTTPS_PROXY", "https://example.org")
1281 # Delete in case our environment has any proxy env vars set
1282 monkeypatch.delenv("HTTP_PROXY", raising=False)
1283 monkeypatch.delenv("ALL_PROXY", raising=False)
1284 monkeypatch.delenv("NO_PROXY", raising=False)
1285 monkeypatch.delenv("http_proxy", raising=False)
1286 monkeypatch.delenv("https_proxy", raising=False)
1287 monkeypatch.delenv("all_proxy", raising=False)
1288 monkeypatch.delenv("no_proxy", raising=False)
1289
1290 client = DefaultHttpxClient()
1291
1292 mounts = tuple(client._mounts.items())
1293 assert len(mounts) == 1
1294 assert mounts[0][0].pattern == "https://"
1295
1296 @pytest.mark.filterwarnings("ignore:.*deprecated.*:DeprecationWarning")
1297 def test_default_client_creation(self) -> None:
1298 # Ensure that the client can be initialized without any exceptions
1299 DefaultHttpxClient(
1300 verify=True,
1301 cert=None,
1302 trust_env=True,
1303 http1=True,
1304 http2=False,
1305 limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
1306 )
1307
1308 @pytest.mark.respx(base_url=base_url)
1309 def test_follow_redirects(self, respx_mock: MockRouter, client: OpenAI) -> None:
1310 # Test that the default follow_redirects=True allows following redirects
1311 respx_mock.post("/redirect").mock(
1312 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1313 )
1314 respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
1315
1316 response = client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
1317 assert response.status_code == 200
1318 assert response.json() == {"status": "ok"}
1319
1320 @pytest.mark.respx(base_url=base_url)
1321 def test_follow_redirects_disabled(self, respx_mock: MockRouter, client: OpenAI) -> None:
1322 # Test that follow_redirects=False prevents following redirects
1323 respx_mock.post("/redirect").mock(
1324 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1325 )
1326
1327 with pytest.raises(APIStatusError) as exc_info:
1328 client.post("/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response)
1329
1330 assert exc_info.value.response.status_code == 302
1331 assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
1332
1333 def test_api_key_before_after_refresh_provider(self) -> None:
1334 client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token")
1335
1336 assert client.api_key == ""
1337 assert "Authorization" not in client.auth_headers
1338
1339 client._refresh_api_key()
1340
1341 assert client.api_key == "test_bearer_token"
1342 assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
1343
1344 def test_api_key_before_after_refresh_str(self) -> None:
1345 client = OpenAI(base_url=base_url, api_key="test_api_key")
1346
1347 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
1348 client._refresh_api_key()
1349
1350 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
1351
1352 @pytest.mark.respx()
1353 def test_api_key_refresh_on_retry(self, respx_mock: MockRouter) -> None:
1354 respx_mock.post(base_url + "/chat/completions").mock(
1355 side_effect=[
1356 httpx.Response(500, json={"error": "server error"}),
1357 httpx.Response(200, json={"foo": "bar"}),
1358 ]
1359 )
1360
1361 counter = 0
1362
1363 def token_provider() -> str:
1364 nonlocal counter
1365
1366 counter += 1
1367
1368 if counter == 1:
1369 return "first"
1370
1371 return "second"
1372
1373 client = OpenAI(base_url=base_url, api_key=token_provider)
1374 client.chat.completions.create(messages=[], model="gpt-4")
1375
1376 calls = cast("list[MockRequestCall]", respx_mock.calls)
1377 assert len(calls) == 2
1378
1379 assert calls[0].request.headers.get("Authorization") == "Bearer first"
1380 assert calls[1].request.headers.get("Authorization") == "Bearer second"
1381
1382 def test_copy_auth(self) -> None:
1383 client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token_1").copy(
1384 api_key=lambda: "test_bearer_token_2"
1385 )
1386 client._refresh_api_key()
1387 assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}
1388
1389
1390class TestAsyncOpenAI:
1391 @pytest.mark.respx(base_url=base_url)
1392 async def test_raw_response(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
1393 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
1394
1395 response = await async_client.post("/foo", cast_to=httpx.Response)
1396 assert response.status_code == 200
1397 assert isinstance(response, httpx.Response)
1398 assert response.json() == {"foo": "bar"}
1399
1400 @pytest.mark.respx(base_url=base_url)
1401 async def test_raw_response_for_binary(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
1402 respx_mock.post("/foo").mock(
1403 return_value=httpx.Response(200, headers={"Content-Type": "application/binary"}, content='{"foo": "bar"}')
1404 )
1405
1406 response = await async_client.post("/foo", cast_to=httpx.Response)
1407 assert response.status_code == 200
1408 assert isinstance(response, httpx.Response)
1409 assert response.json() == {"foo": "bar"}
1410
1411 def test_copy(self, async_client: AsyncOpenAI) -> None:
1412 copied = async_client.copy()
1413 assert id(copied) != id(async_client)
1414
1415 copied = async_client.copy(api_key="another My API Key")
1416 assert copied.api_key == "another My API Key"
1417 assert async_client.api_key == "My API Key"
1418
1419 copied = async_client.copy(admin_api_key="another My Admin API Key")
1420 assert copied.admin_api_key == "another My Admin API Key"
1421 assert async_client.admin_api_key == "My Admin API Key"
1422
1423 def test_copy_default_options(self, async_client: AsyncOpenAI) -> None:
1424 # options that have a default are overridden correctly
1425 copied = async_client.copy(max_retries=7)
1426 assert copied.max_retries == 7
1427 assert async_client.max_retries == 2
1428
1429 copied2 = copied.copy(max_retries=6)
1430 assert copied2.max_retries == 6
1431 assert copied.max_retries == 7
1432
1433 # timeout
1434 assert isinstance(async_client.timeout, httpx.Timeout)
1435 copied = async_client.copy(timeout=None)
1436 assert copied.timeout is None
1437 assert isinstance(async_client.timeout, httpx.Timeout)
1438
1439 async def test_copy_default_headers(self) -> None:
1440 client = AsyncOpenAI(
1441 base_url=base_url,
1442 api_key=api_key,
1443 admin_api_key=admin_api_key,
1444 _strict_response_validation=True,
1445 default_headers={"X-Foo": "bar"},
1446 )
1447 assert client.default_headers["X-Foo"] == "bar"
1448
1449 # does not override the already given value when not specified
1450 copied = client.copy()
1451 assert copied.default_headers["X-Foo"] == "bar"
1452
1453 # merges already given headers
1454 copied = client.copy(default_headers={"X-Bar": "stainless"})
1455 assert copied.default_headers["X-Foo"] == "bar"
1456 assert copied.default_headers["X-Bar"] == "stainless"
1457
1458 # uses new values for any already given headers
1459 copied = client.copy(default_headers={"X-Foo": "stainless"})
1460 assert copied.default_headers["X-Foo"] == "stainless"
1461
1462 # set_default_headers
1463
1464 # completely overrides already set values
1465 copied = client.copy(set_default_headers={})
1466 assert copied.default_headers.get("X-Foo") is None
1467
1468 copied = client.copy(set_default_headers={"X-Bar": "Robert"})
1469 assert copied.default_headers["X-Bar"] == "Robert"
1470
1471 with pytest.raises(
1472 ValueError,
1473 match="`default_headers` and `set_default_headers` arguments are mutually exclusive",
1474 ):
1475 client.copy(set_default_headers={}, default_headers={"X-Foo": "Bar"})
1476 await client.close()
1477
1478 async def test_copy_default_query(self) -> None:
1479 client = AsyncOpenAI(
1480 base_url=base_url,
1481 api_key=api_key,
1482 admin_api_key=admin_api_key,
1483 _strict_response_validation=True,
1484 default_query={"foo": "bar"},
1485 )
1486 assert _get_params(client)["foo"] == "bar"
1487
1488 # does not override the already given value when not specified
1489 copied = client.copy()
1490 assert _get_params(copied)["foo"] == "bar"
1491
1492 # merges already given params
1493 copied = client.copy(default_query={"bar": "stainless"})
1494 params = _get_params(copied)
1495 assert params["foo"] == "bar"
1496 assert params["bar"] == "stainless"
1497
1498 # uses new values for any already given headers
1499 copied = client.copy(default_query={"foo": "stainless"})
1500 assert _get_params(copied)["foo"] == "stainless"
1501
1502 # set_default_query
1503
1504 # completely overrides already set values
1505 copied = client.copy(set_default_query={})
1506 assert _get_params(copied) == {}
1507
1508 copied = client.copy(set_default_query={"bar": "Robert"})
1509 assert _get_params(copied)["bar"] == "Robert"
1510
1511 with pytest.raises(
1512 ValueError,
1513 # TODO: update
1514 match="`default_query` and `set_default_query` arguments are mutually exclusive",
1515 ):
1516 client.copy(set_default_query={}, default_query={"foo": "Bar"})
1517
1518 await client.close()
1519
1520 def test_copy_signature(self, async_client: AsyncOpenAI) -> None:
1521 # ensure the same parameters that can be passed to the client are defined in the `.copy()` method
1522 init_signature = inspect.signature(
1523 # mypy doesn't like that we access the `__init__` property.
1524 async_client.__init__, # type: ignore[misc]
1525 )
1526 copy_signature = inspect.signature(async_client.copy)
1527 exclude_params = {"transport", "proxies", "_strict_response_validation"}
1528
1529 for name in init_signature.parameters.keys():
1530 if name in exclude_params:
1531 continue
1532
1533 copy_param = copy_signature.parameters.get(name)
1534 assert copy_param is not None, f"copy() signature is missing the {name} param"
1535
1536 @pytest.mark.skipif(sys.version_info >= (3, 10), reason="fails because of a memory leak that started from 3.12")
1537 def test_copy_build_request(self, async_client: AsyncOpenAI) -> None:
1538 options = FinalRequestOptions(method="get", url="/foo")
1539
1540 def build_request(options: FinalRequestOptions) -> None:
1541 client_copy = async_client.copy()
1542 client_copy._build_request(options)
1543
1544 # ensure that the machinery is warmed up before tracing starts.
1545 build_request(options)
1546 gc.collect()
1547
1548 tracemalloc.start(1000)
1549
1550 snapshot_before = tracemalloc.take_snapshot()
1551
1552 ITERATIONS = 10
1553 for _ in range(ITERATIONS):
1554 build_request(options)
1555
1556 gc.collect()
1557 snapshot_after = tracemalloc.take_snapshot()
1558
1559 tracemalloc.stop()
1560
1561 def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
1562 if diff.count == 0:
1563 # Avoid false positives by considering only leaks (i.e. allocations that persist).
1564 return
1565
1566 if diff.count % ITERATIONS != 0:
1567 # Avoid false positives by considering only leaks that appear per iteration.
1568 return
1569
1570 for frame in diff.traceback:
1571 if any(
1572 frame.filename.endswith(fragment)
1573 for fragment in [
1574 # to_raw_response_wrapper leaks through the @functools.wraps() decorator.
1575 #
1576 # removing the decorator fixes the leak for reasons we don't understand.
1577 "openai/_legacy_response.py",
1578 "openai/_response.py",
1579 # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
1580 "openai/_compat.py",
1581 # Standard library leaks we don't care about.
1582 "/logging/__init__.py",
1583 ]
1584 ):
1585 return
1586
1587 leaks.append(diff)
1588
1589 leaks: list[tracemalloc.StatisticDiff] = []
1590 for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
1591 add_leak(leaks, diff)
1592 if leaks:
1593 for leak in leaks:
1594 print("MEMORY LEAK:", leak)
1595 for frame in leak.traceback:
1596 print(frame)
1597 raise AssertionError()
1598
1599 async def test_request_timeout(self, async_client: AsyncOpenAI) -> None:
1600 request = async_client._build_request(FinalRequestOptions(method="get", url="/foo"))
1601 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1602 assert timeout == DEFAULT_TIMEOUT
1603
1604 request = async_client._build_request(
1605 FinalRequestOptions(method="get", url="/foo", timeout=httpx.Timeout(100.0))
1606 )
1607 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1608 assert timeout == httpx.Timeout(100.0)
1609
1610 async def test_client_timeout_option(self) -> None:
1611 client = AsyncOpenAI(
1612 base_url=base_url,
1613 api_key=api_key,
1614 admin_api_key=admin_api_key,
1615 _strict_response_validation=True,
1616 timeout=httpx.Timeout(0),
1617 )
1618
1619 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1620 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1621 assert timeout == httpx.Timeout(0)
1622
1623 await client.close()
1624
1625 async def test_http_client_timeout_option(self) -> None:
1626 # custom timeout given to the httpx client should be used
1627 async with httpx.AsyncClient(timeout=None) as http_client:
1628 client = AsyncOpenAI(
1629 base_url=base_url,
1630 api_key=api_key,
1631 admin_api_key=admin_api_key,
1632 _strict_response_validation=True,
1633 http_client=http_client,
1634 )
1635
1636 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1637 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1638 assert timeout == httpx.Timeout(None)
1639
1640 await client.close()
1641
1642 # no timeout given to the httpx client should not use the httpx default
1643 async with httpx.AsyncClient() as http_client:
1644 client = AsyncOpenAI(
1645 base_url=base_url,
1646 api_key=api_key,
1647 admin_api_key=admin_api_key,
1648 _strict_response_validation=True,
1649 http_client=http_client,
1650 )
1651
1652 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1653 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1654 assert timeout == DEFAULT_TIMEOUT
1655
1656 await client.close()
1657
1658 # explicitly passing the default timeout currently results in it being ignored
1659 async with httpx.AsyncClient(timeout=HTTPX_DEFAULT_TIMEOUT) as http_client:
1660 client = AsyncOpenAI(
1661 base_url=base_url,
1662 api_key=api_key,
1663 admin_api_key=admin_api_key,
1664 _strict_response_validation=True,
1665 http_client=http_client,
1666 )
1667
1668 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1669 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1670 assert timeout == DEFAULT_TIMEOUT # our default
1671
1672 await client.close()
1673
1674 def test_invalid_http_client(self) -> None:
1675 with pytest.raises(TypeError, match="Invalid `http_client` arg"):
1676 with httpx.Client() as http_client:
1677 AsyncOpenAI(
1678 base_url=base_url,
1679 api_key=api_key,
1680 admin_api_key=admin_api_key,
1681 _strict_response_validation=True,
1682 http_client=cast(Any, http_client),
1683 )
1684
1685 async def test_default_headers_option(self) -> None:
1686 test_client = AsyncOpenAI(
1687 base_url=base_url,
1688 api_key=api_key,
1689 admin_api_key=admin_api_key,
1690 _strict_response_validation=True,
1691 default_headers={"X-Foo": "bar"},
1692 )
1693 request = test_client._build_request(FinalRequestOptions(method="get", url="/foo"))
1694 assert request.headers.get("x-foo") == "bar"
1695 assert request.headers.get("x-stainless-lang") == "python"
1696
1697 test_client2 = AsyncOpenAI(
1698 base_url=base_url,
1699 api_key=api_key,
1700 admin_api_key=admin_api_key,
1701 _strict_response_validation=True,
1702 default_headers={
1703 "X-Foo": "stainless",
1704 "X-Stainless-Lang": "my-overriding-header",
1705 },
1706 )
1707 request = test_client2._build_request(FinalRequestOptions(method="get", url="/foo"))
1708 assert request.headers.get("x-foo") == "stainless"
1709 assert request.headers.get("x-stainless-lang") == "my-overriding-header"
1710
1711 await test_client.close()
1712 await test_client2.close()
1713
1714 async def test_validate_headers(self) -> None:
1715 client = AsyncOpenAI(
1716 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
1717 )
1718 options = await client._prepare_options(FinalRequestOptions(method="get", url="/foo"))
1719 request = client._build_request(options)
1720 assert request.headers.get("Authorization") == f"Bearer {api_key}"
1721
1722 admin_request = client._build_request(
1723 FinalRequestOptions(
1724 method="get",
1725 url="/organization/projects",
1726 security={"admin_api_key_auth": True},
1727 )
1728 )
1729 assert admin_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
1730
1731 with update_env(**{"OPENAI_API_KEY": Omit()}):
1732 admin_only = AsyncOpenAI(
1733 base_url=base_url,
1734 api_key=None,
1735 admin_api_key=admin_api_key,
1736 _strict_response_validation=True,
1737 )
1738 admin_only_request = admin_only._build_request(
1739 FinalRequestOptions(
1740 method="get",
1741 url="/organization/projects",
1742 security={"admin_api_key_auth": True},
1743 )
1744 )
1745 assert admin_only_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
1746
1747 with pytest.raises(
1748 TypeError,
1749 match="Could not resolve authentication method",
1750 ):
1751 admin_only._build_request(
1752 FinalRequestOptions(
1753 method="post",
1754 url="/responses",
1755 security={"bearer_auth": True},
1756 )
1757 )
1758
1759 with update_env(
1760 **{
1761 "OPENAI_API_KEY": Omit(),
1762 "OPENAI_ADMIN_KEY": Omit(),
1763 }
1764 ):
1765 no_credentials = AsyncOpenAI(
1766 base_url=base_url,
1767 api_key=None,
1768 admin_api_key=None,
1769 _enforce_credentials=False,
1770 _strict_response_validation=True,
1771 )
1772 lowercase_auth_request = no_credentials._build_request(
1773 FinalRequestOptions(method="get", url="/foo", headers={"authorization": "Bearer custom"})
1774 )
1775 assert lowercase_auth_request.headers.get("Authorization") == "Bearer custom"
1776
1777 omitted_auth_request = no_credentials._build_request(
1778 FinalRequestOptions(method="get", url="/foo", headers={"authorization": Omit()})
1779 )
1780 assert "Authorization" not in omitted_auth_request.headers
1781
1782 with update_env(
1783 **{
1784 "OPENAI_API_KEY": Omit(),
1785 "OPENAI_ADMIN_KEY": Omit(),
1786 }
1787 ):
1788 with pytest.raises(OpenAIError, match="Missing credentials"):
1789 AsyncOpenAI(base_url=base_url, api_key=None, admin_api_key=None, _strict_response_validation=True)
1790
1791 @pytest.mark.respx(base_url=base_url)
1792 async def test_api_key_provider_preserves_admin_auth(self, respx_mock: MockRouter) -> None:
1793 respx_mock.get("/organization/projects").mock(return_value=httpx.Response(200, json={"ok": True}))
1794
1795 provider_called = False
1796
1797 async def api_key_provider() -> str:
1798 nonlocal provider_called
1799 provider_called = True
1800 return "dynamic-api-key"
1801
1802 client = AsyncOpenAI(base_url=base_url, api_key=api_key_provider, admin_api_key=admin_api_key)
1803 response = await client.get(
1804 "/organization/projects",
1805 cast_to=httpx.Response,
1806 options={"security": {"admin_api_key_auth": True}},
1807 )
1808
1809 assert response.request.headers.get("Authorization") == f"Bearer {admin_api_key}"
1810 assert provider_called is False
1811
1812 async def test_api_key_provider_does_not_fill_admin_auth(self) -> None:
1813 provider_called = False
1814
1815 async def api_key_provider() -> str:
1816 nonlocal provider_called
1817 provider_called = True
1818 return "dynamic-api-key"
1819
1820 with update_env(OPENAI_ADMIN_KEY=Omit()):
1821 client = AsyncOpenAI(base_url=base_url, api_key=api_key_provider, admin_api_key=None)
1822 with pytest.raises(TypeError, match="Could not resolve authentication method"):
1823 await client.get(
1824 "/organization/projects",
1825 cast_to=httpx.Response,
1826 options={"security": {"admin_api_key_auth": True}},
1827 )
1828
1829 assert provider_called is False
1830
1831 @pytest.mark.respx(base_url=base_url)
1832 async def test_workload_identity_preserves_admin_auth(self, respx_mock: MockRouter) -> None:
1833 respx_mock.get("/organization/projects").mock(return_value=httpx.Response(200, json={"ok": True}))
1834
1835 client = AsyncOpenAI(base_url=base_url, workload_identity=workload_identity, admin_api_key=admin_api_key)
1836 response = await client.get(
1837 "/organization/projects",
1838 cast_to=httpx.Response,
1839 options={"security": {"admin_api_key_auth": True}},
1840 )
1841
1842 assert response.request.headers.get("Authorization") == f"Bearer {admin_api_key}"
1843
1844 async def test_default_query_option(self) -> None:
1845 client = AsyncOpenAI(
1846 base_url=base_url,
1847 api_key=api_key,
1848 admin_api_key=admin_api_key,
1849 _strict_response_validation=True,
1850 default_query={"query_param": "bar"},
1851 )
1852 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1853 url = httpx.URL(request.url)
1854 assert dict(url.params) == {"query_param": "bar"}
1855
1856 request = client._build_request(
1857 FinalRequestOptions(
1858 method="get",
1859 url="/foo",
1860 params={"foo": "baz", "query_param": "overridden"},
1861 )
1862 )
1863 url = httpx.URL(request.url)
1864 assert dict(url.params) == {"foo": "baz", "query_param": "overridden"}
1865
1866 await client.close()
1867
1868 async def test_hardcoded_query_params_in_url(self, async_client: AsyncOpenAI) -> None:
1869 request = async_client._build_request(FinalRequestOptions(method="get", url="/foo?beta=true"))
1870 url = httpx.URL(request.url)
1871 assert dict(url.params) == {"beta": "true"}
1872
1873 request = async_client._build_request(
1874 FinalRequestOptions(
1875 method="get",
1876 url="/foo?beta=true",
1877 params={"limit": "10", "page": "abc"},
1878 )
1879 )
1880 url = httpx.URL(request.url)
1881 assert dict(url.params) == {"beta": "true", "limit": "10", "page": "abc"}
1882
1883 request = async_client._build_request(
1884 FinalRequestOptions(
1885 method="get",
1886 url="/files/a%2Fb?beta=true",
1887 params={"limit": "10"},
1888 )
1889 )
1890 assert request.url.raw_path == b"/files/a%2Fb?beta=true&limit=10"
1891
1892 def test_request_extra_json(self, client: OpenAI) -> None:
1893 request = client._build_request(
1894 FinalRequestOptions(
1895 method="post",
1896 url="/foo",
1897 json_data={"foo": "bar"},
1898 extra_json={"baz": False},
1899 ),
1900 )
1901 data = json.loads(request.content.decode("utf-8"))
1902 assert data == {"foo": "bar", "baz": False}
1903
1904 request = client._build_request(
1905 FinalRequestOptions(
1906 method="post",
1907 url="/foo",
1908 extra_json={"baz": False},
1909 ),
1910 )
1911 data = json.loads(request.content.decode("utf-8"))
1912 assert data == {"baz": False}
1913
1914 # `extra_json` takes priority over `json_data` when keys clash
1915 request = client._build_request(
1916 FinalRequestOptions(
1917 method="post",
1918 url="/foo",
1919 json_data={"foo": "bar", "baz": True},
1920 extra_json={"baz": None},
1921 ),
1922 )
1923 data = json.loads(request.content.decode("utf-8"))
1924 assert data == {"foo": "bar", "baz": None}
1925
1926 def test_request_extra_headers(self, client: OpenAI) -> None:
1927 request = client._build_request(
1928 FinalRequestOptions(
1929 method="post",
1930 url="/foo",
1931 **make_request_options(extra_headers={"X-Foo": "Foo"}),
1932 ),
1933 )
1934 assert request.headers.get("X-Foo") == "Foo"
1935
1936 # `extra_headers` takes priority over `default_headers` when keys clash
1937 request = client.with_options(default_headers={"X-Bar": "true"})._build_request(
1938 FinalRequestOptions(
1939 method="post",
1940 url="/foo",
1941 **make_request_options(
1942 extra_headers={"X-Bar": "false"},
1943 ),
1944 ),
1945 )
1946 assert request.headers.get("X-Bar") == "false"
1947
1948 def test_request_extra_query(self, client: OpenAI) -> None:
1949 request = client._build_request(
1950 FinalRequestOptions(
1951 method="post",
1952 url="/foo",
1953 **make_request_options(
1954 extra_query={"my_query_param": "Foo"},
1955 ),
1956 ),
1957 )
1958 params = dict(request.url.params)
1959 assert params == {"my_query_param": "Foo"}
1960
1961 # if both `query` and `extra_query` are given, they are merged
1962 request = client._build_request(
1963 FinalRequestOptions(
1964 method="post",
1965 url="/foo",
1966 **make_request_options(
1967 query={"bar": "1"},
1968 extra_query={"foo": "2"},
1969 ),
1970 ),
1971 )
1972 params = dict(request.url.params)
1973 assert params == {"bar": "1", "foo": "2"}
1974
1975 # `extra_query` takes priority over `query` when keys clash
1976 request = client._build_request(
1977 FinalRequestOptions(
1978 method="post",
1979 url="/foo",
1980 **make_request_options(
1981 query={"foo": "1"},
1982 extra_query={"foo": "2"},
1983 ),
1984 ),
1985 )
1986 params = dict(request.url.params)
1987 assert params == {"foo": "2"}
1988
1989 def test_multipart_repeating_array(self, async_client: AsyncOpenAI) -> None:
1990 request = async_client._build_request(
1991 FinalRequestOptions.construct(
1992 method="post",
1993 url="/foo",
1994 headers={"Content-Type": "multipart/form-data; boundary=6b7ba517decee4a450543ea6ae821c82"},
1995 json_data={"array": ["foo", "bar"]},
1996 files=[("foo.txt", b"hello world")],
1997 )
1998 )
1999
2000 assert request.read().split(b"\r\n") == [
2001 b"--6b7ba517decee4a450543ea6ae821c82",
2002 b'Content-Disposition: form-data; name="array[]"',
2003 b"",
2004 b"foo",
2005 b"--6b7ba517decee4a450543ea6ae821c82",
2006 b'Content-Disposition: form-data; name="array[]"',
2007 b"",
2008 b"bar",
2009 b"--6b7ba517decee4a450543ea6ae821c82",
2010 b'Content-Disposition: form-data; name="foo.txt"; filename="upload"',
2011 b"Content-Type: application/octet-stream",
2012 b"",
2013 b"hello world",
2014 b"--6b7ba517decee4a450543ea6ae821c82--",
2015 b"",
2016 ]
2017
2018 @pytest.mark.respx(base_url=base_url)
2019 async def test_binary_content_upload(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2020 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
2021
2022 file_content = b"Hello, this is a test file."
2023
2024 response = await async_client.post(
2025 "/upload",
2026 content=file_content,
2027 cast_to=httpx.Response,
2028 options={"headers": {"Content-Type": "application/octet-stream"}},
2029 )
2030
2031 assert response.status_code == 200
2032 assert response.request.headers["Content-Type"] == "application/octet-stream"
2033 assert response.content == file_content
2034
2035 async def test_binary_content_upload_with_asynciterator(self) -> None:
2036 file_content = b"Hello, this is a test file."
2037 counter = Counter()
2038 iterator = _make_async_iterator([file_content], counter=counter)
2039
2040 async def mock_handler(request: httpx.Request) -> httpx.Response:
2041 assert counter.value == 0, "the request body should not have been read"
2042 return httpx.Response(200, content=await request.aread())
2043
2044 async with AsyncOpenAI(
2045 base_url=base_url,
2046 api_key=api_key,
2047 admin_api_key=admin_api_key,
2048 _strict_response_validation=True,
2049 http_client=httpx.AsyncClient(transport=MockTransport(handler=mock_handler)),
2050 ) as client:
2051 response = await client.post(
2052 "/upload",
2053 content=iterator,
2054 cast_to=httpx.Response,
2055 options={"headers": {"Content-Type": "application/octet-stream"}},
2056 )
2057
2058 assert response.status_code == 200
2059 assert response.request.headers["Content-Type"] == "application/octet-stream"
2060 assert response.content == file_content
2061 assert counter.value == 1
2062
2063 @pytest.mark.respx(base_url=base_url)
2064 async def test_binary_content_upload_with_body_is_deprecated(
2065 self, respx_mock: MockRouter, async_client: AsyncOpenAI
2066 ) -> None:
2067 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
2068
2069 file_content = b"Hello, this is a test file."
2070
2071 with pytest.deprecated_call(
2072 match="Passing raw bytes as `body` is deprecated and will be removed in a future version. Please pass raw bytes via the `content` parameter instead."
2073 ):
2074 response = await async_client.post(
2075 "/upload",
2076 body=file_content,
2077 cast_to=httpx.Response,
2078 options={"headers": {"Content-Type": "application/octet-stream"}},
2079 )
2080
2081 assert response.status_code == 200
2082 assert response.request.headers["Content-Type"] == "application/octet-stream"
2083 assert response.content == file_content
2084
2085 @pytest.mark.respx(base_url=base_url)
2086 async def test_basic_union_response(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2087 class Model1(BaseModel):
2088 name: str
2089
2090 class Model2(BaseModel):
2091 foo: str
2092
2093 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
2094
2095 response = await async_client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
2096 assert isinstance(response, Model2)
2097 assert response.foo == "bar"
2098
2099 @pytest.mark.respx(base_url=base_url)
2100 async def test_union_response_different_types(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2101 """Union of objects with the same field name using a different type"""
2102
2103 class Model1(BaseModel):
2104 foo: int
2105
2106 class Model2(BaseModel):
2107 foo: str
2108
2109 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
2110
2111 response = await async_client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
2112 assert isinstance(response, Model2)
2113 assert response.foo == "bar"
2114
2115 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": 1}))
2116
2117 response = await async_client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
2118 assert isinstance(response, Model1)
2119 assert response.foo == 1
2120
2121 @pytest.mark.respx(base_url=base_url)
2122 async def test_non_application_json_content_type_for_json_data(
2123 self, respx_mock: MockRouter, async_client: AsyncOpenAI
2124 ) -> None:
2125 """
2126 Response that sets Content-Type to something other than application/json but returns json data
2127 """
2128
2129 class Model(BaseModel):
2130 foo: int
2131
2132 respx_mock.get("/foo").mock(
2133 return_value=httpx.Response(
2134 200,
2135 content=json.dumps({"foo": 2}),
2136 headers={"Content-Type": "application/text"},
2137 )
2138 )
2139
2140 response = await async_client.get("/foo", cast_to=Model)
2141 assert isinstance(response, Model)
2142 assert response.foo == 2
2143
2144 async def test_base_url_setter(self) -> None:
2145 client = AsyncOpenAI(
2146 base_url="https://example.com/from_init",
2147 api_key=api_key,
2148 admin_api_key=admin_api_key,
2149 _strict_response_validation=True,
2150 )
2151 assert client.base_url == "https://example.com/from_init/"
2152
2153 client.base_url = "https://example.com/from_setter" # type: ignore[assignment]
2154
2155 assert client.base_url == "https://example.com/from_setter/"
2156
2157 await client.close()
2158
2159 async def test_base_url_env(self) -> None:
2160 with update_env(OPENAI_BASE_URL="http://localhost:5000/from/env"):
2161 client = AsyncOpenAI(api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True)
2162 assert client.base_url == "http://localhost:5000/from/env/"
2163
2164 @pytest.mark.parametrize(
2165 "client",
2166 [
2167 AsyncOpenAI(
2168 base_url="http://localhost:5000/custom/path/",
2169 api_key=api_key,
2170 admin_api_key=admin_api_key,
2171 _strict_response_validation=True,
2172 ),
2173 AsyncOpenAI(
2174 base_url="http://localhost:5000/custom/path/",
2175 api_key=api_key,
2176 admin_api_key=admin_api_key,
2177 _strict_response_validation=True,
2178 http_client=httpx.AsyncClient(),
2179 ),
2180 ],
2181 ids=["standard", "custom http client"],
2182 )
2183 async def test_base_url_trailing_slash(self, client: AsyncOpenAI) -> None:
2184 request = client._build_request(
2185 FinalRequestOptions(
2186 method="post",
2187 url="/foo",
2188 json_data={"foo": "bar"},
2189 ),
2190 )
2191 assert request.url == "http://localhost:5000/custom/path/foo"
2192 await client.close()
2193
2194 @pytest.mark.parametrize(
2195 "client",
2196 [
2197 AsyncOpenAI(
2198 base_url="http://localhost:5000/custom/path/",
2199 api_key=api_key,
2200 admin_api_key=admin_api_key,
2201 _strict_response_validation=True,
2202 ),
2203 AsyncOpenAI(
2204 base_url="http://localhost:5000/custom/path/",
2205 api_key=api_key,
2206 admin_api_key=admin_api_key,
2207 _strict_response_validation=True,
2208 http_client=httpx.AsyncClient(),
2209 ),
2210 ],
2211 ids=["standard", "custom http client"],
2212 )
2213 async def test_base_url_no_trailing_slash(self, client: AsyncOpenAI) -> None:
2214 request = client._build_request(
2215 FinalRequestOptions(
2216 method="post",
2217 url="/foo",
2218 json_data={"foo": "bar"},
2219 ),
2220 )
2221 assert request.url == "http://localhost:5000/custom/path/foo"
2222 await client.close()
2223
2224 @pytest.mark.parametrize(
2225 "client",
2226 [
2227 AsyncOpenAI(
2228 base_url="http://localhost:5000/custom/path/",
2229 api_key=api_key,
2230 admin_api_key=admin_api_key,
2231 _strict_response_validation=True,
2232 ),
2233 AsyncOpenAI(
2234 base_url="http://localhost:5000/custom/path/",
2235 api_key=api_key,
2236 admin_api_key=admin_api_key,
2237 _strict_response_validation=True,
2238 http_client=httpx.AsyncClient(),
2239 ),
2240 ],
2241 ids=["standard", "custom http client"],
2242 )
2243 async def test_absolute_request_url(self, client: AsyncOpenAI) -> None:
2244 request = client._build_request(
2245 FinalRequestOptions(
2246 method="post",
2247 url="https://myapi.com/foo",
2248 json_data={"foo": "bar"},
2249 ),
2250 )
2251 assert request.url == "https://myapi.com/foo"
2252 await client.close()
2253
2254 async def test_copied_client_does_not_close_http(self) -> None:
2255 test_client = AsyncOpenAI(
2256 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
2257 )
2258 assert not test_client.is_closed()
2259
2260 copied = test_client.copy()
2261 assert copied is not test_client
2262
2263 del copied
2264
2265 await asyncio.sleep(0.2)
2266 assert not test_client.is_closed()
2267
2268 async def test_client_context_manager(self) -> None:
2269 test_client = AsyncOpenAI(
2270 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
2271 )
2272 async with test_client as c2:
2273 assert c2 is test_client
2274 assert not c2.is_closed()
2275 assert not test_client.is_closed()
2276 assert test_client.is_closed()
2277
2278 @pytest.mark.respx(base_url=base_url)
2279 async def test_client_response_validation_error(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2280 class Model(BaseModel):
2281 foo: str
2282
2283 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": {"invalid": True}}))
2284
2285 with pytest.raises(APIResponseValidationError) as exc:
2286 await async_client.get("/foo", cast_to=Model)
2287
2288 assert isinstance(exc.value.__cause__, ValidationError)
2289
2290 async def test_client_max_retries_validation(self) -> None:
2291 with pytest.raises(TypeError, match=r"max_retries cannot be None"):
2292 AsyncOpenAI(
2293 base_url=base_url,
2294 api_key=api_key,
2295 admin_api_key=admin_api_key,
2296 _strict_response_validation=True,
2297 max_retries=cast(Any, None),
2298 )
2299
2300 @pytest.mark.respx(base_url=base_url)
2301 async def test_default_stream_cls(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2302 class Model(BaseModel):
2303 name: str
2304
2305 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
2306
2307 stream = await async_client.post("/foo", cast_to=Model, stream=True, stream_cls=AsyncStream[Model])
2308 assert isinstance(stream, AsyncStream)
2309 await stream.response.aclose()
2310
2311 @pytest.mark.respx(base_url=base_url)
2312 async def test_received_text_for_expected_json(self, respx_mock: MockRouter) -> None:
2313 class Model(BaseModel):
2314 name: str
2315
2316 respx_mock.get("/foo").mock(return_value=httpx.Response(200, text="my-custom-format"))
2317
2318 strict_client = AsyncOpenAI(
2319 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
2320 )
2321
2322 with pytest.raises(APIResponseValidationError):
2323 await strict_client.get("/foo", cast_to=Model)
2324
2325 non_strict_client = AsyncOpenAI(
2326 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=False
2327 )
2328
2329 response = await non_strict_client.get("/foo", cast_to=Model)
2330 assert isinstance(response, str) # type: ignore[unreachable]
2331
2332 await strict_client.close()
2333 await non_strict_client.close()
2334
2335 @pytest.mark.parametrize(
2336 "remaining_retries,retry_after,timeout",
2337 [
2338 [3, "20", 20],
2339 [3, "0", 0.5],
2340 [3, "-10", 0.5],
2341 [3, "60", 60],
2342 [3, "61", 0.5],
2343 [3, "Fri, 29 Sep 2023 16:26:57 GMT", 20],
2344 [3, "Fri, 29 Sep 2023 16:26:37 GMT", 0.5],
2345 [3, "Fri, 29 Sep 2023 16:26:27 GMT", 0.5],
2346 [3, "Fri, 29 Sep 2023 16:27:37 GMT", 60],
2347 [3, "Fri, 29 Sep 2023 16:27:38 GMT", 0.5],
2348 [3, "99999999999999999999999999999999999", 0.5],
2349 [3, "Zun, 29 Sep 2023 16:26:27 GMT", 0.5],
2350 [3, "", 0.5],
2351 [2, "", 0.5 * 2.0],
2352 [1, "", 0.5 * 4.0],
2353 [-1100, "", 8], # test large number potentially overflowing
2354 ],
2355 )
2356 @mock.patch("time.time", mock.MagicMock(return_value=1696004797))
2357 async def test_parse_retry_after_header(
2358 self, remaining_retries: int, retry_after: str, timeout: float, async_client: AsyncOpenAI
2359 ) -> None:
2360 headers = httpx.Headers({"retry-after": retry_after})
2361 options = FinalRequestOptions(method="get", url="/foo", max_retries=3)
2362 calculated = async_client._calculate_retry_timeout(remaining_retries, options, headers)
2363 assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType]
2364
2365 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2366 @pytest.mark.respx(base_url=base_url)
2367 async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2368 respx_mock.post("/chat/completions").mock(side_effect=httpx.TimeoutException("Test timeout error"))
2369
2370 with pytest.raises(APITimeoutError):
2371 await async_client.chat.completions.with_streaming_response.create(
2372 messages=[
2373 {
2374 "content": "string",
2375 "role": "developer",
2376 }
2377 ],
2378 model="gpt-5.4",
2379 ).__aenter__()
2380
2381 assert _get_open_connections(async_client) == 0
2382
2383 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2384 @pytest.mark.respx(base_url=base_url)
2385 async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2386 respx_mock.post("/chat/completions").mock(return_value=httpx.Response(500))
2387
2388 with pytest.raises(APIStatusError):
2389 await async_client.chat.completions.with_streaming_response.create(
2390 messages=[
2391 {
2392 "content": "string",
2393 "role": "developer",
2394 }
2395 ],
2396 model="gpt-5.4",
2397 ).__aenter__()
2398 assert _get_open_connections(async_client) == 0
2399
2400 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2401 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2402 @pytest.mark.respx(base_url=base_url)
2403 @pytest.mark.parametrize("failure_mode", ["status", "exception"])
2404 async def test_retries_taken(
2405 self,
2406 async_client: AsyncOpenAI,
2407 failures_before_success: int,
2408 failure_mode: Literal["status", "exception"],
2409 respx_mock: MockRouter,
2410 ) -> None:
2411 client = async_client.with_options(max_retries=4)
2412
2413 nb_retries = 0
2414
2415 def retry_handler(_request: httpx.Request) -> httpx.Response:
2416 nonlocal nb_retries
2417 if nb_retries < failures_before_success:
2418 nb_retries += 1
2419 if failure_mode == "exception":
2420 raise RuntimeError("oops")
2421 return httpx.Response(500)
2422 return httpx.Response(200)
2423
2424 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2425
2426 response = await client.chat.completions.with_raw_response.create(
2427 messages=[
2428 {
2429 "content": "string",
2430 "role": "developer",
2431 }
2432 ],
2433 model="gpt-5.4",
2434 )
2435
2436 assert response.retries_taken == failures_before_success
2437 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
2438
2439 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2440 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2441 @pytest.mark.respx(base_url=base_url)
2442 async def test_omit_retry_count_header(
2443 self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
2444 ) -> None:
2445 client = async_client.with_options(max_retries=4)
2446
2447 nb_retries = 0
2448
2449 def retry_handler(_request: httpx.Request) -> httpx.Response:
2450 nonlocal nb_retries
2451 if nb_retries < failures_before_success:
2452 nb_retries += 1
2453 return httpx.Response(500)
2454 return httpx.Response(200)
2455
2456 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2457
2458 response = await client.chat.completions.with_raw_response.create(
2459 messages=[
2460 {
2461 "content": "string",
2462 "role": "developer",
2463 }
2464 ],
2465 model="gpt-5.4",
2466 extra_headers={"x-stainless-retry-count": Omit()},
2467 )
2468
2469 assert len(response.http_request.headers.get_list("x-stainless-retry-count")) == 0
2470
2471 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2472 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2473 @pytest.mark.respx(base_url=base_url)
2474 async def test_overwrite_retry_count_header(
2475 self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
2476 ) -> None:
2477 client = async_client.with_options(max_retries=4)
2478
2479 nb_retries = 0
2480
2481 def retry_handler(_request: httpx.Request) -> httpx.Response:
2482 nonlocal nb_retries
2483 if nb_retries < failures_before_success:
2484 nb_retries += 1
2485 return httpx.Response(500)
2486 return httpx.Response(200)
2487
2488 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2489
2490 response = await client.chat.completions.with_raw_response.create(
2491 messages=[
2492 {
2493 "content": "string",
2494 "role": "developer",
2495 }
2496 ],
2497 model="gpt-5.4",
2498 extra_headers={"x-stainless-retry-count": "42"},
2499 )
2500
2501 assert response.http_request.headers.get("x-stainless-retry-count") == "42"
2502
2503 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2504 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2505 @pytest.mark.respx(base_url=base_url)
2506 async def test_retries_taken_new_response_class(
2507 self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
2508 ) -> None:
2509 client = async_client.with_options(max_retries=4)
2510
2511 nb_retries = 0
2512
2513 def retry_handler(_request: httpx.Request) -> httpx.Response:
2514 nonlocal nb_retries
2515 if nb_retries < failures_before_success:
2516 nb_retries += 1
2517 return httpx.Response(500)
2518 return httpx.Response(200)
2519
2520 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2521
2522 async with client.chat.completions.with_streaming_response.create(
2523 messages=[
2524 {
2525 "content": "string",
2526 "role": "developer",
2527 }
2528 ],
2529 model="gpt-5.4",
2530 ) as response:
2531 assert response.retries_taken == failures_before_success
2532 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
2533
2534 async def test_get_platform(self) -> None:
2535 platform = await asyncify(get_platform)()
2536 assert isinstance(platform, (str, OtherPlatform))
2537
2538 async def test_proxy_environment_variables(self, monkeypatch: pytest.MonkeyPatch) -> None:
2539 # Test that the proxy environment variables are set correctly
2540 monkeypatch.setenv("HTTPS_PROXY", "https://example.org")
2541 # Delete in case our environment has any proxy env vars set
2542 monkeypatch.delenv("HTTP_PROXY", raising=False)
2543 monkeypatch.delenv("ALL_PROXY", raising=False)
2544 monkeypatch.delenv("NO_PROXY", raising=False)
2545 monkeypatch.delenv("http_proxy", raising=False)
2546 monkeypatch.delenv("https_proxy", raising=False)
2547 monkeypatch.delenv("all_proxy", raising=False)
2548 monkeypatch.delenv("no_proxy", raising=False)
2549
2550 client = DefaultAsyncHttpxClient()
2551
2552 mounts = tuple(client._mounts.items())
2553 assert len(mounts) == 1
2554 assert mounts[0][0].pattern == "https://"
2555
2556 @pytest.mark.filterwarnings("ignore:.*deprecated.*:DeprecationWarning")
2557 async def test_default_client_creation(self) -> None:
2558 # Ensure that the client can be initialized without any exceptions
2559 DefaultAsyncHttpxClient(
2560 verify=True,
2561 cert=None,
2562 trust_env=True,
2563 http1=True,
2564 http2=False,
2565 limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
2566 )
2567
2568 @pytest.mark.respx(base_url=base_url)
2569 async def test_follow_redirects(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2570 # Test that the default follow_redirects=True allows following redirects
2571 respx_mock.post("/redirect").mock(
2572 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
2573 )
2574 respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
2575
2576 response = await async_client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
2577 assert response.status_code == 200
2578 assert response.json() == {"status": "ok"}
2579
2580 @pytest.mark.respx(base_url=base_url)
2581 async def test_follow_redirects_disabled(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2582 # Test that follow_redirects=False prevents following redirects
2583 respx_mock.post("/redirect").mock(
2584 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
2585 )
2586
2587 with pytest.raises(APIStatusError) as exc_info:
2588 await async_client.post(
2589 "/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
2590 )
2591
2592 assert exc_info.value.response.status_code == 302
2593 assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
2594
2595 async def test_api_key_before_after_refresh_provider(self) -> None:
2596 async def mock_api_key_provider():
2597 return "test_bearer_token"
2598
2599 client = AsyncOpenAI(base_url=base_url, api_key=mock_api_key_provider)
2600
2601 assert client.api_key == ""
2602 assert "Authorization" not in client.auth_headers
2603
2604 await client._refresh_api_key()
2605
2606 assert client.api_key == "test_bearer_token"
2607 assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
2608
2609 async def test_api_key_before_after_refresh_str(self) -> None:
2610 client = AsyncOpenAI(base_url=base_url, api_key="test_api_key")
2611
2612 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
2613 await client._refresh_api_key()
2614
2615 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
2616
2617 @pytest.mark.respx()
2618 async def test_bearer_token_refresh_async(self, respx_mock: MockRouter) -> None:
2619 respx_mock.post(base_url + "/chat/completions").mock(
2620 side_effect=[
2621 httpx.Response(500, json={"error": "server error"}),
2622 httpx.Response(200, json={"foo": "bar"}),
2623 ]
2624 )
2625
2626 counter = 0
2627
2628 async def token_provider() -> str:
2629 nonlocal counter
2630
2631 counter += 1
2632
2633 if counter == 1:
2634 return "first"
2635
2636 return "second"
2637
2638 client = AsyncOpenAI(base_url=base_url, api_key=token_provider)
2639 await client.chat.completions.create(messages=[], model="gpt-4")
2640
2641 calls = cast("list[MockRequestCall]", respx_mock.calls)
2642 assert len(calls) == 2
2643
2644 assert calls[0].request.headers.get("Authorization") == "Bearer first"
2645 assert calls[1].request.headers.get("Authorization") == "Bearer second"
2646
2647 async def test_copy_auth(self) -> None:
2648 async def token_provider_1() -> str:
2649 return "test_bearer_token_1"
2650
2651 async def token_provider_2() -> str:
2652 return "test_bearer_token_2"
2653
2654 client = AsyncOpenAI(base_url=base_url, api_key=token_provider_1).copy(api_key=token_provider_2)
2655 await client._refresh_api_key()
2656 assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}
2657
2658
2659class TestWorkloadIdentity401Retry:
2660 @pytest.mark.respx()
2661 def test_workload_identity_401_retry(self, respx_mock: MockRouter) -> None:
2662 provider_call_count = 0
2663
2664 def provider() -> str:
2665 nonlocal provider_call_count
2666 provider_call_count += 1
2667 return f"external-subject-token-{provider_call_count}"
2668
2669 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2670 side_effect=[
2671 httpx.Response(
2672 200,
2673 json={
2674 "access_token": "openai-access-token-1",
2675 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2676 "token_type": "Bearer",
2677 "expires_in": 3600,
2678 },
2679 ),
2680 httpx.Response(
2681 200,
2682 json={
2683 "access_token": "openai-access-token-2",
2684 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2685 "token_type": "Bearer",
2686 "expires_in": 3600,
2687 },
2688 ),
2689 ]
2690 )
2691
2692 respx_mock.post(base_url + "/chat/completions").mock(
2693 side_effect=[
2694 httpx.Response(401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}),
2695 httpx.Response(
2696 200,
2697 json={
2698 "id": "chatcmpl-123",
2699 "object": "chat.completion",
2700 "created": 1234567890,
2701 "model": "gpt-4",
2702 "choices": [],
2703 },
2704 ),
2705 ]
2706 )
2707
2708 with OpenAI(
2709 base_url=base_url,
2710 workload_identity={
2711 **workload_identity,
2712 "provider": {
2713 "get_token": provider,
2714 "token_type": "jwt",
2715 },
2716 },
2717 organization="org_123",
2718 project="proj_123",
2719 _strict_response_validation=True,
2720 ) as client:
2721 client.chat.completions.create(messages=[], model="gpt-4")
2722
2723 calls = cast("list[MockRequestCall]", respx_mock.calls)
2724 assert len(calls) == 4
2725
2726 assert calls[0].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2727 assert calls[1].request.url == httpx.URL(base_url + "/chat/completions")
2728 assert calls[1].request.headers.get("Authorization") == "Bearer openai-access-token-1"
2729
2730 assert calls[2].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2731
2732 assert calls[3].request.url == httpx.URL(base_url + "/chat/completions")
2733 assert calls[3].request.headers.get("Authorization") == "Bearer openai-access-token-2"
2734
2735 assert provider_call_count == 2
2736
2737 @pytest.mark.respx()
2738 def test_401_without_workload_identity_no_retry(self, respx_mock: MockRouter) -> None:
2739 respx_mock.post(base_url + "/chat/completions").mock(
2740 return_value=httpx.Response(
2741 401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}
2742 )
2743 )
2744
2745 with OpenAI(
2746 base_url=base_url,
2747 api_key="test-api-key",
2748 _strict_response_validation=True,
2749 ) as client:
2750 with pytest.raises(APIStatusError) as exc_info:
2751 client.chat.completions.create(messages=[], model="gpt-4")
2752
2753 assert exc_info.value.status_code == 401
2754
2755 calls = cast("list[MockRequestCall]", respx_mock.calls)
2756 assert len(calls) == 1
2757
2758 @pytest.mark.respx()
2759 def test_non_401_errors_no_retry(self, respx_mock: MockRouter) -> None:
2760 provider_call_count = 0
2761
2762 def provider() -> str:
2763 nonlocal provider_call_count
2764 provider_call_count += 1
2765 return "external-subject-token"
2766
2767 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2768 return_value=httpx.Response(
2769 200,
2770 json={
2771 "access_token": "openai-access-token-1",
2772 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2773 "token_type": "Bearer",
2774 "expires_in": 3600,
2775 },
2776 )
2777 )
2778
2779 respx_mock.post(base_url + "/chat/completions").mock(
2780 return_value=httpx.Response(403, json={"error": {"message": "Forbidden", "type": "invalid_request_error"}})
2781 )
2782
2783 with OpenAI(
2784 base_url=base_url,
2785 workload_identity={
2786 **workload_identity,
2787 "provider": {
2788 "get_token": provider,
2789 "token_type": "jwt",
2790 },
2791 },
2792 organization="org_123",
2793 project="proj_123",
2794 _strict_response_validation=True,
2795 ) as client:
2796 with pytest.raises(APIStatusError) as exc_info:
2797 client.chat.completions.create(messages=[], model="gpt-4")
2798
2799 assert exc_info.value.status_code == 403
2800
2801 calls = cast("list[MockRequestCall]", respx_mock.calls)
2802 assert len(calls) == 2
2803
2804 assert provider_call_count == 1
2805
2806
2807class TestAsyncWorkloadIdentity401Retry:
2808 @pytest.mark.respx()
2809 async def test_workload_identity_401_retry(self, respx_mock: MockRouter) -> None:
2810 provider_call_count = 0
2811
2812 def provider() -> str:
2813 nonlocal provider_call_count
2814 provider_call_count += 1
2815 return f"external-subject-token-{provider_call_count}"
2816
2817 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2818 side_effect=[
2819 httpx.Response(
2820 200,
2821 json={
2822 "access_token": "openai-access-token-1",
2823 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2824 "token_type": "Bearer",
2825 "expires_in": 3600,
2826 },
2827 ),
2828 httpx.Response(
2829 200,
2830 json={
2831 "access_token": "openai-access-token-2",
2832 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2833 "token_type": "Bearer",
2834 "expires_in": 3600,
2835 },
2836 ),
2837 ]
2838 )
2839
2840 respx_mock.post(base_url + "/chat/completions").mock(
2841 side_effect=[
2842 httpx.Response(401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}),
2843 httpx.Response(
2844 200,
2845 json={
2846 "id": "chatcmpl-123",
2847 "object": "chat.completion",
2848 "created": 1234567890,
2849 "model": "gpt-4",
2850 "choices": [],
2851 },
2852 ),
2853 ]
2854 )
2855
2856 async with AsyncOpenAI(
2857 base_url=base_url,
2858 workload_identity={
2859 **workload_identity,
2860 "provider": {
2861 "get_token": provider,
2862 "token_type": "jwt",
2863 },
2864 },
2865 organization="org_123",
2866 project="proj_123",
2867 _strict_response_validation=True,
2868 ) as client:
2869 await client.chat.completions.create(messages=[], model="gpt-4")
2870
2871 calls = cast("list[MockRequestCall]", respx_mock.calls)
2872 assert len(calls) == 4
2873
2874 assert calls[0].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2875 assert calls[1].request.url == httpx.URL(base_url + "/chat/completions")
2876 assert calls[1].request.headers.get("Authorization") == "Bearer openai-access-token-1"
2877
2878 assert calls[2].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2879
2880 assert calls[3].request.url == httpx.URL(base_url + "/chat/completions")
2881 assert calls[3].request.headers.get("Authorization") == "Bearer openai-access-token-2"
2882
2883 assert provider_call_count == 2
2884
2885 @pytest.mark.respx()
2886 async def test_401_without_workload_identity_no_retry(self, respx_mock: MockRouter) -> None:
2887 respx_mock.post(base_url + "/chat/completions").mock(
2888 return_value=httpx.Response(
2889 401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}
2890 )
2891 )
2892
2893 async with AsyncOpenAI(
2894 base_url=base_url,
2895 api_key="test-api-key",
2896 _strict_response_validation=True,
2897 ) as client:
2898 with pytest.raises(APIStatusError) as exc_info:
2899 await client.chat.completions.create(messages=[], model="gpt-4")
2900
2901 assert exc_info.value.status_code == 401
2902
2903 calls = cast("list[MockRequestCall]", respx_mock.calls)
2904 assert len(calls) == 1
2905
2906 @pytest.mark.respx()
2907 async def test_non_401_errors_no_retry(self, respx_mock: MockRouter) -> None:
2908 provider_call_count = 0
2909
2910 def provider() -> str:
2911 nonlocal provider_call_count
2912 provider_call_count += 1
2913 return "external-subject-token"
2914
2915 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2916 return_value=httpx.Response(
2917 200,
2918 json={
2919 "access_token": "openai-access-token-1",
2920 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2921 "token_type": "Bearer",
2922 "expires_in": 3600,
2923 },
2924 )
2925 )
2926
2927 respx_mock.post(base_url + "/chat/completions").mock(
2928 return_value=httpx.Response(403, json={"error": {"message": "Forbidden", "type": "invalid_request_error"}})
2929 )
2930
2931 async with AsyncOpenAI(
2932 base_url=base_url,
2933 workload_identity={
2934 **workload_identity,
2935 "provider": {
2936 "get_token": provider,
2937 "token_type": "jwt",
2938 },
2939 },
2940 organization="org_123",
2941 project="proj_123",
2942 _strict_response_validation=True,
2943 ) as client:
2944 with pytest.raises(APIStatusError) as exc_info:
2945 await client.chat.completions.create(messages=[], model="gpt-4")
2946
2947 assert exc_info.value.status_code == 403
2948
2949 calls = cast("list[MockRequestCall]", respx_mock.calls)
2950 assert len(calls) == 2
2951
2952 assert provider_call_count == 1