From d289afdbde9a282244c3e6962b09f6cd7d8b9434 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Wed, 6 Nov 2024 11:00:34 -0800 Subject: [PATCH] Fix exception in server when client SSE connection closes --- llama_stack/distribution/server/server.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py index 2560f4070..16c0fd0e0 100644 --- a/llama_stack/distribution/server/server.py +++ b/llama_stack/distribution/server/server.py @@ -209,7 +209,8 @@ async def maybe_await(value): async def sse_generator(event_gen): try: - async for item in await event_gen: + event_gen = await event_gen + async for item in event_gen: yield create_sse_event(item) await asyncio.sleep(0.01) except asyncio.CancelledError: @@ -229,7 +230,6 @@ async def sse_generator(event_gen): def create_dynamic_typed_route(func: Any, method: str): - async def endpoint(request: Request, **kwargs): await start_trace(func.__name__)