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
98ece5783a061011ceede7dc2f94fa081666fdcf

Branches

Tags

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

Clone

HTTPS

Download ZIP

openai-python/src/openai

src/openai/_response.py

832lines · modecode

RawDownload
Latest commit unavailable.
1from __future__ import annotations
2
3import os
4import inspect
5import logging
6import datetime
7import functools
8from types import TracebackType
9from typing import (
10 TYPE_CHECKING,
11 Any,
12 Union,
13 Generic,
14 TypeVar,
15 Callable,
16 Iterator,
17 AsyncIterator,
18 cast,
19 overload,
20)
21from typing_extensions import Awaitable, ParamSpec, override, get_origin
22
23import anyio
24import httpx
25import pydantic
26
27from ._types import NoneType
28from ._utils import is_given, extract_type_arg, is_annotated_type, extract_type_var_from_base
29from ._models import BaseModel, is_basemodel
30from ._constants import RAW_RESPONSE_HEADER, OVERRIDE_CAST_TO_HEADER
31from ._streaming import Stream, AsyncStream, is_stream_class_type, extract_stream_chunk_type
32from ._exceptions import OpenAIError, APIResponseValidationError
33
34if TYPE_CHECKING:
35 from ._models import FinalRequestOptions
36 from ._base_client import BaseClient
37
38
39P = ParamSpec("P")
40R = TypeVar("R")
41_T = TypeVar("_T")
42_APIResponseT = TypeVar("_APIResponseT", bound="APIResponse[Any]")
43_AsyncAPIResponseT = TypeVar("_AsyncAPIResponseT", bound="AsyncAPIResponse[Any]")
44
45log: logging.Logger = logging.getLogger(__name__)
46
47
48class BaseAPIResponse(Generic[R]):
49 _cast_to: type[R]
50 _client: BaseClient[Any, Any]
51 _parsed_by_type: dict[type[Any], Any]
52 _is_sse_stream: bool
53 _stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None
54 _options: FinalRequestOptions
55
56 http_response: httpx.Response
57
58 def __init__(
59 self,
60 *,
61 raw: httpx.Response,
62 cast_to: type[R],
63 client: BaseClient[Any, Any],
64 stream: bool,
65 stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
66 options: FinalRequestOptions,
67 ) -> None:
68 self._cast_to = cast_to
69 self._client = client
70 self._parsed_by_type = {}
71 self._is_sse_stream = stream
72 self._stream_cls = stream_cls
73 self._options = options
74 self.http_response = raw
75
76 @property
77 def headers(self) -> httpx.Headers:
78 return self.http_response.headers
79
80 @property
81 def http_request(self) -> httpx.Request:
82 """Returns the httpx Request instance associated with the current response."""
83 return self.http_response.request
84
85 @property
86 def status_code(self) -> int:
87 return self.http_response.status_code
88
89 @property
90 def url(self) -> httpx.URL:
91 """Returns the URL for which the request was made."""
92 return self.http_response.url
93
94 @property
95 def method(self) -> str:
96 return self.http_request.method
97
98 @property
99 def http_version(self) -> str:
100 return self.http_response.http_version
101
102 @property
103 def elapsed(self) -> datetime.timedelta:
104 """The time taken for the complete request/response cycle to complete."""
105 return self.http_response.elapsed
106
107 @property
108 def is_closed(self) -> bool:
109 """Whether or not the response body has been closed.
110
111 If this is False then there is response data that has not been read yet.
112 You must either fully consume the response body or call `.close()`
113 before discarding the response to prevent resource leaks.
114 """
115 return self.http_response.is_closed
116
117 @override
118 def __repr__(self) -> str:
119 return (
120 f"<{self.__class__.__name__} [{self.status_code} {self.http_response.reason_phrase}] type={self._cast_to}>"
121 )
122
123 def _parse(self, *, to: type[_T] | None = None) -> R | _T:
124 # unwrap `Annotated[T, ...]` -> `T`
125 if to and is_annotated_type(to):
126 to = extract_type_arg(to, 0)
127
128 if self._is_sse_stream:
129 if to:
130 if not is_stream_class_type(to):
131 raise TypeError(f"Expected custom parse type to be a subclass of {Stream} or {AsyncStream}")
132
133 return cast(
134 _T,
135 to(
136 cast_to=extract_stream_chunk_type(
137 to,
138 failure_message="Expected custom stream type to be passed with a type argument, e.g. Stream[ChunkType]",
139 ),
140 response=self.http_response,
141 client=cast(Any, self._client),
142 ),
143 )
144
145 if self._stream_cls:
146 return cast(
147 R,
148 self._stream_cls(
149 cast_to=extract_stream_chunk_type(self._stream_cls),
150 response=self.http_response,
151 client=cast(Any, self._client),
152 ),
153 )
154
155 stream_cls = cast("type[Stream[Any]] | type[AsyncStream[Any]] | None", self._client._default_stream_cls)
156 if stream_cls is None:
157 raise MissingStreamClassError()
158
159 return cast(
160 R,
161 stream_cls(
162 cast_to=self._cast_to,
163 response=self.http_response,
164 client=cast(Any, self._client),
165 ),
166 )
167
168 cast_to = to if to is not None else self._cast_to
169
170 # unwrap `Annotated[T, ...]` -> `T`
171 if is_annotated_type(cast_to):
172 cast_to = extract_type_arg(cast_to, 0)
173
174 if cast_to is NoneType:
175 return cast(R, None)
176
177 response = self.http_response
178 if cast_to == str:
179 return cast(R, response.text)
180
181 if cast_to == bytes:
182 return cast(R, response.content)
183
184 if cast_to == int:
185 return cast(R, int(response.text))
186
187 if cast_to == float:
188 return cast(R, float(response.text))
189
190 origin = get_origin(cast_to) or cast_to
191
192 # handle the legacy binary response case
193 if inspect.isclass(cast_to) and cast_to.__name__ == "HttpxBinaryResponseContent":
194 return cast(R, cast_to(response)) # type: ignore
195
196 if origin == APIResponse:
197 raise RuntimeError("Unexpected state - cast_to is `APIResponse`")
198
199 if inspect.isclass(origin) and issubclass(origin, httpx.Response):
200 # Because of the invariance of our ResponseT TypeVar, users can subclass httpx.Response
201 # and pass that class to our request functions. We cannot change the variance to be either
202 # covariant or contravariant as that makes our usage of ResponseT illegal. We could construct
203 # the response class ourselves but that is something that should be supported directly in httpx
204 # as it would be easy to incorrectly construct the Response object due to the multitude of arguments.
205 if cast_to != httpx.Response:
206 raise ValueError(f"Subclasses of httpx.Response cannot be passed to `cast_to`")
207 return cast(R, response)
208
209 if inspect.isclass(origin) and not issubclass(origin, BaseModel) and issubclass(origin, pydantic.BaseModel):
210 raise TypeError("Pydantic models must subclass our base model type, e.g. `from openai import BaseModel`")
211
212 if (
213 cast_to is not object
214 and not origin is list
215 and not origin is dict
216 and not origin is Union
217 and not issubclass(origin, BaseModel)
218 ):
219 raise RuntimeError(
220 f"Unsupported type, expected {cast_to} to be a subclass of {BaseModel}, {dict}, {list}, {Union}, {NoneType}, {str} or {httpx.Response}."
221 )
222
223 # split is required to handle cases where additional information is included
224 # in the response, e.g. application/json; charset=utf-8
225 content_type, *_ = response.headers.get("content-type", "*").split(";")
226 if content_type != "application/json":
227 if is_basemodel(cast_to):
228 try:
229 data = response.json()
230 except Exception as exc:
231 log.debug("Could not read JSON from response data due to %s - %s", type(exc), exc)
232 else:
233 return self._client._process_response_data(
234 data=data,
235 cast_to=cast_to, # type: ignore
236 response=response,
237 )
238
239 if self._client._strict_response_validation:
240 raise APIResponseValidationError(
241 response=response,
242 message=f"Expected Content-Type response header to be `application/json` but received `{content_type}` instead.",
243 body=response.text,
244 )
245
246 # If the API responds with content that isn't JSON then we just return
247 # the (decoded) text without performing any parsing so that you can still
248 # handle the response however you need to.
249 return response.text # type: ignore
250
251 data = response.json()
252
253 return self._client._process_response_data(
254 data=data,
255 cast_to=cast_to, # type: ignore
256 response=response,
257 )
258
259
260class APIResponse(BaseAPIResponse[R]):
261 @property
262 def request_id(self) -> str | None:
263 return self.http_response.headers.get("x-request-id") # type: ignore[no-any-return]
264
265 @overload
266 def parse(self, *, to: type[_T]) -> _T:
267 ...
268
269 @overload
270 def parse(self) -> R:
271 ...
272
273 def parse(self, *, to: type[_T] | None = None) -> R | _T:
274 """Returns the rich python representation of this response's data.
275
276 For lower-level control, see `.read()`, `.json()`, `.iter_bytes()`.
277
278 You can customise the type that the response is parsed into through
279 the `to` argument, e.g.
280
281 ```py
282 from openai import BaseModel
283
284
285 class MyModel(BaseModel):
286 foo: str
287
288
289 obj = response.parse(to=MyModel)
290 print(obj.foo)
291 ```
292
293 We support parsing:
294 - `BaseModel`
295 - `dict`
296 - `list`
297 - `Union`
298 - `str`
299 - `int`
300 - `float`
301 - `httpx.Response`
302 """
303 cache_key = to if to is not None else self._cast_to
304 cached = self._parsed_by_type.get(cache_key)
305 if cached is not None:
306 return cached # type: ignore[no-any-return]
307
308 if not self._is_sse_stream:
309 self.read()
310
311 parsed = self._parse(to=to)
312 if is_given(self._options.post_parser):
313 parsed = self._options.post_parser(parsed)
314
315 self._parsed_by_type[cache_key] = parsed
316 return parsed
317
318 def read(self) -> bytes:
319 """Read and return the binary response content."""
320 try:
321 return self.http_response.read()
322 except httpx.StreamConsumed as exc:
323 # The default error raised by httpx isn't very
324 # helpful in our case so we re-raise it with
325 # a different error message.
326 raise StreamAlreadyConsumed() from exc
327
328 def text(self) -> str:
329 """Read and decode the response content into a string."""
330 self.read()
331 return self.http_response.text
332
333 def json(self) -> object:
334 """Read and decode the JSON response content."""
335 self.read()
336 return self.http_response.json()
337
338 def close(self) -> None:
339 """Close the response and release the connection.
340
341 Automatically called if the response body is read to completion.
342 """
343 self.http_response.close()
344
345 def iter_bytes(self, chunk_size: int | None = None) -> Iterator[bytes]:
346 """
347 A byte-iterator over the decoded response content.
348
349 This automatically handles gzip, deflate and brotli encoded responses.
350 """
351 for chunk in self.http_response.iter_bytes(chunk_size):
352 yield chunk
353
354 def iter_text(self, chunk_size: int | None = None) -> Iterator[str]:
355 """A str-iterator over the decoded response content
356 that handles both gzip, deflate, etc but also detects the content's
357 string encoding.
358 """
359 for chunk in self.http_response.iter_text(chunk_size):
360 yield chunk
361
362 def iter_lines(self) -> Iterator[str]:
363 """Like `iter_text()` but will only yield chunks for each line"""
364 for chunk in self.http_response.iter_lines():
365 yield chunk
366
367
368class AsyncAPIResponse(BaseAPIResponse[R]):
369 @property
370 def request_id(self) -> str | None:
371 return self.http_response.headers.get("x-request-id") # type: ignore[no-any-return]
372
373 @overload
374 async def parse(self, *, to: type[_T]) -> _T:
375 ...
376
377 @overload
378 async def parse(self) -> R:
379 ...
380
381 async def parse(self, *, to: type[_T] | None = None) -> R | _T:
382 """Returns the rich python representation of this response's data.
383
384 For lower-level control, see `.read()`, `.json()`, `.iter_bytes()`.
385
386 You can customise the type that the response is parsed into through
387 the `to` argument, e.g.
388
389 ```py
390 from openai import BaseModel
391
392
393 class MyModel(BaseModel):
394 foo: str
395
396
397 obj = response.parse(to=MyModel)
398 print(obj.foo)
399 ```
400
401 We support parsing:
402 - `BaseModel`
403 - `dict`
404 - `list`
405 - `Union`
406 - `str`
407 - `httpx.Response`
408 """
409 cache_key = to if to is not None else self._cast_to
410 cached = self._parsed_by_type.get(cache_key)
411 if cached is not None:
412 return cached # type: ignore[no-any-return]
413
414 if not self._is_sse_stream:
415 await self.read()
416
417 parsed = self._parse(to=to)
418 if is_given(self._options.post_parser):
419 parsed = self._options.post_parser(parsed)
420
421 self._parsed_by_type[cache_key] = parsed
422 return parsed
423
424 async def read(self) -> bytes:
425 """Read and return the binary response content."""
426 try:
427 return await self.http_response.aread()
428 except httpx.StreamConsumed as exc:
429 # the default error raised by httpx isn't very
430 # helpful in our case so we re-raise it with
431 # a different error message
432 raise StreamAlreadyConsumed() from exc
433
434 async def text(self) -> str:
435 """Read and decode the response content into a string."""
436 await self.read()
437 return self.http_response.text
438
439 async def json(self) -> object:
440 """Read and decode the JSON response content."""
441 await self.read()
442 return self.http_response.json()
443
444 async def close(self) -> None:
445 """Close the response and release the connection.
446
447 Automatically called if the response body is read to completion.
448 """
449 await self.http_response.aclose()
450
451 async def iter_bytes(self, chunk_size: int | None = None) -> AsyncIterator[bytes]:
452 """
453 A byte-iterator over the decoded response content.
454
455 This automatically handles gzip, deflate and brotli encoded responses.
456 """
457 async for chunk in self.http_response.aiter_bytes(chunk_size):
458 yield chunk
459
460 async def iter_text(self, chunk_size: int | None = None) -> AsyncIterator[str]:
461 """A str-iterator over the decoded response content
462 that handles both gzip, deflate, etc but also detects the content's
463 string encoding.
464 """
465 async for chunk in self.http_response.aiter_text(chunk_size):
466 yield chunk
467
468 async def iter_lines(self) -> AsyncIterator[str]:
469 """Like `iter_text()` but will only yield chunks for each line"""
470 async for chunk in self.http_response.aiter_lines():
471 yield chunk
472
473
474class BinaryAPIResponse(APIResponse[bytes]):
475 """Subclass of APIResponse providing helpers for dealing with binary data.
476
477 Note: If you want to stream the response data instead of eagerly reading it
478 all at once then you should use `.with_streaming_response` when making
479 the API request, e.g. `.with_streaming_response.get_binary_response()`
480 """
481
482 def write_to_file(
483 self,
484 file: str | os.PathLike[str],
485 ) -> None:
486 """Write the output to the given file.
487
488 Accepts a filename or any path-like object, e.g. pathlib.Path
489
490 Note: if you want to stream the data to the file instead of writing
491 all at once then you should use `.with_streaming_response` when making
492 the API request, e.g. `.with_streaming_response.get_binary_response()`
493 """
494 with open(file, mode="wb") as f:
495 for data in self.iter_bytes():
496 f.write(data)
497
498
499class AsyncBinaryAPIResponse(AsyncAPIResponse[bytes]):
500 """Subclass of APIResponse providing helpers for dealing with binary data.
501
502 Note: If you want to stream the response data instead of eagerly reading it
503 all at once then you should use `.with_streaming_response` when making
504 the API request, e.g. `.with_streaming_response.get_binary_response()`
505 """
506
507 async def write_to_file(
508 self,
509 file: str | os.PathLike[str],
510 ) -> None:
511 """Write the output to the given file.
512
513 Accepts a filename or any path-like object, e.g. pathlib.Path
514
515 Note: if you want to stream the data to the file instead of writing
516 all at once then you should use `.with_streaming_response` when making
517 the API request, e.g. `.with_streaming_response.get_binary_response()`
518 """
519 path = anyio.Path(file)
520 async with await path.open(mode="wb") as f:
521 async for data in self.iter_bytes():
522 await f.write(data)
523
524
525class StreamedBinaryAPIResponse(APIResponse[bytes]):
526 def stream_to_file(
527 self,
528 file: str | os.PathLike[str],
529 *,
530 chunk_size: int | None = None,
531 ) -> None:
532 """Streams the output to the given file.
533
534 Accepts a filename or any path-like object, e.g. pathlib.Path
535 """
536 with open(file, mode="wb") as f:
537 for data in self.iter_bytes(chunk_size):
538 f.write(data)
539
540
541class AsyncStreamedBinaryAPIResponse(AsyncAPIResponse[bytes]):
542 async def stream_to_file(
543 self,
544 file: str | os.PathLike[str],
545 *,
546 chunk_size: int | None = None,
547 ) -> None:
548 """Streams the output to the given file.
549
550 Accepts a filename or any path-like object, e.g. pathlib.Path
551 """
552 path = anyio.Path(file)
553 async with await path.open(mode="wb") as f:
554 async for data in self.iter_bytes(chunk_size):
555 await f.write(data)
556
557
558class MissingStreamClassError(TypeError):
559 def __init__(self) -> None:
560 super().__init__(
561 "The `stream` argument was set to `True` but the `stream_cls` argument was not given. See `openai._streaming` for reference",
562 )
563
564
565class StreamAlreadyConsumed(OpenAIError):
566 """
567 Attempted to read or stream content, but the content has already
568 been streamed.
569
570 This can happen if you use a method like `.iter_lines()` and then attempt
571 to read th entire response body afterwards, e.g.
572
573 ```py
574 response = await client.post(...)
575 async for line in response.iter_lines():
576 ... # do something with `line`
577
578 content = await response.read()
579 # ^ error
580 ```
581
582 If you want this behaviour you'll need to either manually accumulate the response
583 content or call `await response.read()` before iterating over the stream.
584 """
585
586 def __init__(self) -> None:
587 message = (
588 "Attempted to read or stream some content, but the content has "
589 "already been streamed. "
590 "This could be due to attempting to stream the response "
591 "content more than once."
592 "\n\n"
593 "You can fix this by manually accumulating the response content while streaming "
594 "or by calling `.read()` before starting to stream."
595 )
596 super().__init__(message)
597
598
599class ResponseContextManager(Generic[_APIResponseT]):
600 """Context manager for ensuring that a request is not made
601 until it is entered and that the response will always be closed
602 when the context manager exits
603 """
604
605 def __init__(self, request_func: Callable[[], _APIResponseT]) -> None:
606 self._request_func = request_func
607 self.__response: _APIResponseT | None = None
608
609 def __enter__(self) -> _APIResponseT:
610 self.__response = self._request_func()
611 return self.__response
612
613 def __exit__(
614 self,
615 exc_type: type[BaseException] | None,
616 exc: BaseException | None,
617 exc_tb: TracebackType | None,
618 ) -> None:
619 if self.__response is not None:
620 self.__response.close()
621
622
623class AsyncResponseContextManager(Generic[_AsyncAPIResponseT]):
624 """Context manager for ensuring that a request is not made
625 until it is entered and that the response will always be closed
626 when the context manager exits
627 """
628
629 def __init__(self, api_request: Awaitable[_AsyncAPIResponseT]) -> None:
630 self._api_request = api_request
631 self.__response: _AsyncAPIResponseT | None = None
632
633 async def __aenter__(self) -> _AsyncAPIResponseT:
634 self.__response = await self._api_request
635 return self.__response
636
637 async def __aexit__(
638 self,
639 exc_type: type[BaseException] | None,
640 exc: BaseException | None,
641 exc_tb: TracebackType | None,
642 ) -> None:
643 if self.__response is not None:
644 await self.__response.close()
645
646
647def to_streamed_response_wrapper(func: Callable[P, R]) -> Callable[P, ResponseContextManager[APIResponse[R]]]:
648 """Higher order function that takes one of our bound API methods and wraps it
649 to support streaming and returning the raw `APIResponse` object directly.
650 """
651
652 @functools.wraps(func)
653 def wrapped(*args: P.args, **kwargs: P.kwargs) -> ResponseContextManager[APIResponse[R]]:
654 extra_headers: dict[str, str] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
655 extra_headers[RAW_RESPONSE_HEADER] = "stream"
656
657 kwargs["extra_headers"] = extra_headers
658
659 make_request = functools.partial(func, *args, **kwargs)
660
661 return ResponseContextManager(cast(Callable[[], APIResponse[R]], make_request))
662
663 return wrapped
664
665
666def async_to_streamed_response_wrapper(
667 func: Callable[P, Awaitable[R]],
668) -> Callable[P, AsyncResponseContextManager[AsyncAPIResponse[R]]]:
669 """Higher order function that takes one of our bound API methods and wraps it
670 to support streaming and returning the raw `APIResponse` object directly.
671 """
672
673 @functools.wraps(func)
674 def wrapped(*args: P.args, **kwargs: P.kwargs) -> AsyncResponseContextManager[AsyncAPIResponse[R]]:
675 extra_headers: dict[str, str] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
676 extra_headers[RAW_RESPONSE_HEADER] = "stream"
677
678 kwargs["extra_headers"] = extra_headers
679
680 make_request = func(*args, **kwargs)
681
682 return AsyncResponseContextManager(cast(Awaitable[AsyncAPIResponse[R]], make_request))
683
684 return wrapped
685
686
687def to_custom_streamed_response_wrapper(
688 func: Callable[P, object],
689 response_cls: type[_APIResponseT],
690) -> Callable[P, ResponseContextManager[_APIResponseT]]:
691 """Higher order function that takes one of our bound API methods and an `APIResponse` class
692 and wraps the method to support streaming and returning the given response class directly.
693
694 Note: the given `response_cls` *must* be concrete, e.g. `class BinaryAPIResponse(APIResponse[bytes])`
695 """
696
697 @functools.wraps(func)
698 def wrapped(*args: P.args, **kwargs: P.kwargs) -> ResponseContextManager[_APIResponseT]:
699 extra_headers: dict[str, Any] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
700 extra_headers[RAW_RESPONSE_HEADER] = "stream"
701 extra_headers[OVERRIDE_CAST_TO_HEADER] = response_cls
702
703 kwargs["extra_headers"] = extra_headers
704
705 make_request = functools.partial(func, *args, **kwargs)
706
707 return ResponseContextManager(cast(Callable[[], _APIResponseT], make_request))
708
709 return wrapped
710
711
712def async_to_custom_streamed_response_wrapper(
713 func: Callable[P, Awaitable[object]],
714 response_cls: type[_AsyncAPIResponseT],
715) -> Callable[P, AsyncResponseContextManager[_AsyncAPIResponseT]]:
716 """Higher order function that takes one of our bound API methods and an `APIResponse` class
717 and wraps the method to support streaming and returning the given response class directly.
718
719 Note: the given `response_cls` *must* be concrete, e.g. `class BinaryAPIResponse(APIResponse[bytes])`
720 """
721
722 @functools.wraps(func)
723 def wrapped(*args: P.args, **kwargs: P.kwargs) -> AsyncResponseContextManager[_AsyncAPIResponseT]:
724 extra_headers: dict[str, Any] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
725 extra_headers[RAW_RESPONSE_HEADER] = "stream"
726 extra_headers[OVERRIDE_CAST_TO_HEADER] = response_cls
727
728 kwargs["extra_headers"] = extra_headers
729
730 make_request = func(*args, **kwargs)
731
732 return AsyncResponseContextManager(cast(Awaitable[_AsyncAPIResponseT], make_request))
733
734 return wrapped
735
736
737def to_raw_response_wrapper(func: Callable[P, R]) -> Callable[P, APIResponse[R]]:
738 """Higher order function that takes one of our bound API methods and wraps it
739 to support returning the raw `APIResponse` object directly.
740 """
741
742 @functools.wraps(func)
743 def wrapped(*args: P.args, **kwargs: P.kwargs) -> APIResponse[R]:
744 extra_headers: dict[str, str] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
745 extra_headers[RAW_RESPONSE_HEADER] = "raw"
746
747 kwargs["extra_headers"] = extra_headers
748
749 return cast(APIResponse[R], func(*args, **kwargs))
750
751 return wrapped
752
753
754def async_to_raw_response_wrapper(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[AsyncAPIResponse[R]]]:
755 """Higher order function that takes one of our bound API methods and wraps it
756 to support returning the raw `APIResponse` object directly.
757 """
758
759 @functools.wraps(func)
760 async def wrapped(*args: P.args, **kwargs: P.kwargs) -> AsyncAPIResponse[R]:
761 extra_headers: dict[str, str] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
762 extra_headers[RAW_RESPONSE_HEADER] = "raw"
763
764 kwargs["extra_headers"] = extra_headers
765
766 return cast(AsyncAPIResponse[R], await func(*args, **kwargs))
767
768 return wrapped
769
770
771def to_custom_raw_response_wrapper(
772 func: Callable[P, object],
773 response_cls: type[_APIResponseT],
774) -> Callable[P, _APIResponseT]:
775 """Higher order function that takes one of our bound API methods and an `APIResponse` class
776 and wraps the method to support returning the given response class directly.
777
778 Note: the given `response_cls` *must* be concrete, e.g. `class BinaryAPIResponse(APIResponse[bytes])`
779 """
780
781 @functools.wraps(func)
782 def wrapped(*args: P.args, **kwargs: P.kwargs) -> _APIResponseT:
783 extra_headers: dict[str, Any] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
784 extra_headers[RAW_RESPONSE_HEADER] = "raw"
785 extra_headers[OVERRIDE_CAST_TO_HEADER] = response_cls
786
787 kwargs["extra_headers"] = extra_headers
788
789 return cast(_APIResponseT, func(*args, **kwargs))
790
791 return wrapped
792
793
794def async_to_custom_raw_response_wrapper(
795 func: Callable[P, Awaitable[object]],
796 response_cls: type[_AsyncAPIResponseT],
797) -> Callable[P, Awaitable[_AsyncAPIResponseT]]:
798 """Higher order function that takes one of our bound API methods and an `APIResponse` class
799 and wraps the method to support returning the given response class directly.
800
801 Note: the given `response_cls` *must* be concrete, e.g. `class BinaryAPIResponse(APIResponse[bytes])`
802 """
803
804 @functools.wraps(func)
805 def wrapped(*args: P.args, **kwargs: P.kwargs) -> Awaitable[_AsyncAPIResponseT]:
806 extra_headers: dict[str, Any] = {**(cast(Any, kwargs.get("extra_headers")) or {})}
807 extra_headers[RAW_RESPONSE_HEADER] = "raw"
808 extra_headers[OVERRIDE_CAST_TO_HEADER] = response_cls
809
810 kwargs["extra_headers"] = extra_headers
811
812 return cast(Awaitable[_AsyncAPIResponseT], func(*args, **kwargs))
813
814 return wrapped
815
816
817def extract_response_type(typ: type[BaseAPIResponse[Any]]) -> type:
818 """Given a type like `APIResponse[T]`, returns the generic type variable `T`.
819
820 This also handles the case where a concrete subclass is given, e.g.
821 ```py
822 class MyResponse(APIResponse[bytes]):
823 ...
824
825 extract_response_type(MyResponse) -> bytes
826 ```
827 """
828 return extract_type_var_from_base(
829 typ,
830 generic_bases=cast("tuple[type, ...]", (BaseAPIResponse, APIResponse, AsyncAPIResponse)),
831 index=0,
832 )
833