From ebbb57beae795cd5937262bb13fb1115b46e92a3 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Tue, 14 Jan 2025 10:49:51 -0800 Subject: [PATCH] more correct event loop handling for stream --- llama_stack/distribution/library_client.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/llama_stack/distribution/library_client.py b/llama_stack/distribution/library_client.py index dacbc18b8..0c124e64b 100644 --- a/llama_stack/distribution/library_client.py +++ b/llama_stack/distribution/library_client.py @@ -144,8 +144,6 @@ class LlamaStackAsLibraryClient(LlamaStackClient): print(f"Removed handler {handler.__class__.__name__} from root logger") def request(self, *args, **kwargs): - async_response = asyncio.run(self.async_client.request(*args, **kwargs)) - if kwargs.get("stream"): # NOTE: We are using AsyncLlamaStackClient under the hood # A new event loop is needed to convert the AsyncStream @@ -155,8 +153,11 @@ class LlamaStackAsLibraryClient(LlamaStackClient): def sync_generator(): try: + async_stream = loop.run_until_complete( + self.async_client.request(*args, **kwargs) + ) while True: - chunk = loop.run_until_complete(async_response.__anext__()) + chunk = loop.run_until_complete(async_stream.__anext__()) yield chunk except StopAsyncIteration: pass @@ -165,7 +166,7 @@ class LlamaStackAsLibraryClient(LlamaStackClient): return sync_generator() else: - return async_response + return asyncio.run(self.async_client.request(*args, **kwargs)) class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):