Repository: jina-ai/jina
File: jina/serve/runtimes/gateway/request_handling.py (Python)

Code before:
                tracer_provider=tracer_provider,
            )
        )

    async def _load_balance(self, request):
        import aiohttp
        from aiohttp import web

        target_server = next(self.load_balancer_servers)
        target_url = f'{target_server}{request.path_qs}'

        try:
            async with aiohttp.ClientSession() as session:

                request_kwargs = {
                    'headers': request.headers,
                    'params': request.query,
                }
                try:
                    payload = await request.content.read()
                    if payload:
                        request_kwargs['data'] = payload
                except Exception:
                    self.logger.debug('No JSON payload found in request')

                async with _RequestContextManager(
                    session._request(request.method, target_url, **request_kwargs)
                ) as response:
                    # Looking for application/octet-stream, text/event-stream, text/stream
                    if response.content_type.endswith('stream'):

                        # Create a StreamResponse with the same headers and status as the target response
                        stream_response = web.StreamResponse(
                            status=response.status,
                            headers=response.headers,
                        )

                        # Prepare the response to send headers
                        await stream_response.prepare(request)

                        # Stream the response from the target server to the client
                        async for chunk in response.content.iter_any():
                            await stream_response.write(chunk)

                        # Close the stream response once all chunks are sent
                        await stream_response.write_eof()
                        return stream_response
                    content = await response.read()
                    return web.Response(
                        body=content,
                        status=response.status,

Review comment (suggestion):
@NarekA why not directly use the `request` method from the session?, it's already wrapping the context manager for you.

```
    def request(
        self, method: str, url: StrOrURL, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP request."""
        return _RequestContextManager(self._request(method, url, **kwargs))
        ```

Code after:
        )

    async def _load_balance(self, request):
        import aiohttp
        from aiohttp import web

        target_server = next(self.load_balancer_servers)
        target_url = f'{target_server}{request.path_qs}'

        try:
            async with aiohttp.ClientSession() as session:

                request_kwargs = {}
                try:
                    payload = await request.json()
                    if payload:
                        request_kwargs['json'] = payload
                except Exception:
                    self.logger.debug('No JSON payload found in request')

                async with session.request(
                    request.method,
                    url=target_url,
                    auto_decompress=False,
                    **request_kwargs,
                ) as response:
                    # Create a StreamResponse with the same headers and status as the target response
                    stream_response = web.StreamResponse(
                        status=response.status,
                        headers=response.headers,
                    )

                    # Prepare the response to send headers
                    await stream_response.prepare(request)

                    # Stream the response from the target server to the client
                    async for chunk in response.content.iter_any():
                        await stream_response.write(chunk)

                    # Close the stream response once all chunks are sent
                    await stream_response.write_eof()
                    return stream_response

        except aiohttp.ClientError as e:
            return web.Response(text=f'Error: {str(e)}', status=500)

    def _websocket_fastapi_default_app(self, tracing, tracer_provider):
        from jina.helper import extend_rest_interface
        from jina.serve.runtimes.gateway.websocket_fastapi_app import get_fastapi_app

        return extend_rest_interface(