use a persistent runloop in library client

This commit is contained in:
Ashwin Bharambe 2025-07-18 11:45:36 -07:00
parent a2f460871b
commit e90ba95ae0

View file

@ -128,6 +128,8 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
self.skip_logger_removal = skip_logger_removal self.skip_logger_removal = skip_logger_removal
self.provider_data = provider_data self.provider_data = provider_data
self.loop = asyncio.new_event_loop()
def initialize(self): def initialize(self):
if in_notebook(): if in_notebook():
import nest_asyncio import nest_asyncio
@ -136,7 +138,7 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
if not self.skip_logger_removal: if not self.skip_logger_removal:
self._remove_root_logger_handlers() self._remove_root_logger_handlers()
return asyncio.run(self.async_client.initialize()) return self.loop.run_until_complete(self.async_client.initialize())
def _remove_root_logger_handlers(self): def _remove_root_logger_handlers(self):
""" """
@ -149,10 +151,7 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
logger.info(f"Removed handler {handler.__class__.__name__} from root logger") logger.info(f"Removed handler {handler.__class__.__name__} from root logger")
def request(self, *args, **kwargs): def request(self, *args, **kwargs):
# NOTE: We are using AsyncLlamaStackClient under the hood loop = self.loop
# A new event loop is needed to convert the AsyncStream
# from async client into SyncStream return type for streaming
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop) asyncio.set_event_loop(loop)
if kwargs.get("stream"): if kwargs.get("stream"):
@ -169,7 +168,6 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
pending = asyncio.all_tasks(loop) pending = asyncio.all_tasks(loop)
if pending: if pending:
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
loop.close()
return sync_generator() return sync_generator()
else: else:
@ -179,7 +177,6 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
pending = asyncio.all_tasks(loop) pending = asyncio.all_tasks(loop)
if pending: if pending:
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
loop.close()
return result return result