Skip to content

Commit

Permalink
fix: pre-commit changes
Browse files Browse the repository at this point in the history
  • Loading branch information
NarekA committed Oct 27, 2023
1 parent c4afd99 commit c654c14
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 6 deletions.
6 changes: 5 additions & 1 deletion jina/serve/runtimes/gateway/http_fastapi_app_docarrayv2.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,11 @@ def add_streaming_routes(
)
async def streaming_get(request: Request, body: input_doc_model = None):
body = body or dict(request.query_params)
body = input_doc_model.parse_obj(body) if docarray_v2 else Document.from_dict(body)
body = (
input_doc_model.parse_obj(body)
if docarray_v2
else Document.from_dict(body)
)

async def event_generator():
async for doc, error in streamer.stream_doc(
Expand Down
8 changes: 4 additions & 4 deletions jina/serve/runtimes/gateway/request_handling.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,14 +175,13 @@ def _http_fastapi_default_app(
)
)


async def _load_balance(self, request: 'aiohttp.web_request.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:

Expand All @@ -192,8 +191,9 @@ async def _load_balance(self, request: 'aiohttp.web_request.Request'):
if payload:
request_kwargs['json'] = payload


async with session.get(url=target_url, **request_kwargs) as response:
async with session.get(
url=target_url, **request_kwargs
) as response:
# Create a StreamResponse with the same headers and status as the target response
stream_response = web.StreamResponse(
status=response.status,
Expand Down
5 changes: 4 additions & 1 deletion tests/integration/streaming/test_streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,10 @@ async def test_streaming_deployment(protocol, include_gateway):
client = Client(port=port, protocol=protocol, asyncio=True)
i = 0
async for doc in client.stream_doc(
on='/hello', inputs=Document(text='hello world'), return_type=Document, input_type=Document
on='/hello',
inputs=Document(text='hello world'),
return_type=Document,
input_type=Document,
):
docs.append(doc.text)
i += 1
Expand Down

0 comments on commit c654c14

Please sign in to comment.