Class StreamingResponseWithStatusCode

gpustack/api/responses.py:12–109

Variation of StreamingResponse that can dynamically decide the HTTP status code, based on the returns from the content iterator (parameter 'content'). Expects the content to yield tuples of (content: str, status_code: int), instead of just content as it was in the original StreamingResponse

class StreamingResponseWithStatusCode(StreamingResponse):
    '''
    Variation of StreamingResponse that can dynamically decide the HTTP status code, based on the returns from the content iterator (parameter 'content').
    Expects the content to yield tuples of (content: str, status_code: int), instead of just content as it was in the original StreamingResponse.
    The parameter status_code in the constructor is ignored, but kept for compatibility with StreamingResponse.
    '''

    async def stream_response(self, send: Send) -> None:
        try:
            first_chunk_content, headers, self.status_code = (
                await self.body_iterator.__anext__()
            )

            if not isinstance(first_chunk_content, bytes):
                first_chunk_content = first_chunk_content.encode(self.charset)

            asgi_headers: List[Tuple[bytes, bytes]] = [
                (key.encode("latin-1"), value.encode("latin-1"))
                for key, value in headers.items()
            ]
            await send(
                {
                    "type": "http.response.start",
                    "status": self.status_code,
                    "headers": asgi_headers,
                }
            )
            await send(
                {
                    "type": "http.response.body",
                    "body": first_chunk_content,
                    "more_body": True,
                }
            )

            async for chunk_content, _, _ in self.body_iterator:
                if not isinstance(chunk_content, bytes):
                    chunk_content = chunk_content.encode(self.charset)
                await send(
                    {
                        "type": "http.response.body",
                        "body": chunk_content,
                        "more_body": True,
                    }
                )

            await send({"type": "http.response.body", "body": b"", "more_body": False})

        except StopAsyncIteration:
            self.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
            await send(
                {
                    "type": "http.response.start",
                    "status": self.status_code,
                    "headers": self.raw_headers,
                }
            )