From bfd2106fc3f55b763f61e12db98101ae127847e0 Mon Sep 17 00:00:00 2001
From: Dinesh Yeduguru <yvdinesh@gmail.com>
Date: Mon, 6 Jan 2025 13:01:55 -0800
Subject: [PATCH] add flag to skip remove of logger handlers

---
 llama_stack/distribution/library_client.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/llama_stack/distribution/library_client.py b/llama_stack/distribution/library_client.py
index 0292717d6..fdb2be5cd 100644
--- a/llama_stack/distribution/library_client.py
+++ b/llama_stack/distribution/library_client.py
@@ -171,6 +171,7 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
     def __init__(
         self,
         config_path_or_template_name: str,
+        skip_logger_removal: bool = False,
         custom_provider_registry: Optional[ProviderRegistry] = None,
     ):
         super().__init__()
@@ -178,12 +179,14 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
             config_path_or_template_name, custom_provider_registry
         )
         self.pool_executor = ThreadPoolExecutor(max_workers=4)
+        self.skip_logger_removal = skip_logger_removal
 
     def initialize(self):
         if in_notebook():
             import nest_asyncio
 
             nest_asyncio.apply()
+        if not self.skip_logger_removal:
             self._remove_root_logger_handlers()
 
         return asyncio.run(self.async_client.initialize())