From a2170353af47015dbe2f057b147a20fd0ce81681 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Mon, 9 Dec 2024 09:37:52 -0800
Subject: [PATCH] better detection for jupyter

---
 llama_stack/distribution/library_client.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/llama_stack/distribution/library_client.py b/llama_stack/distribution/library_client.py
index 64cd343d4..693e2f56c 100644
--- a/llama_stack/distribution/library_client.py
+++ b/llama_stack/distribution/library_client.py
@@ -33,16 +33,17 @@ from llama_stack.distribution.stack import (
 T = TypeVar("T")
 
 
-def is_jupyter():
-    """Check if we're running in a Jupyter notebook"""
+def in_notebook():
     try:
-        shell = get_ipython().__class__.__name__  # type: ignore
-        if shell == "ZMQInteractiveShell":  # Jupyter notebook or qtconsole
-            return True
-        else:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
             return False
-    except NameError:  # Probably standard Python interpreter
+    except ImportError:
         return False
+    except AttributeError:
+        return False
+    return True
 
 
 def stream_across_asyncio_run_boundary(
@@ -115,7 +116,7 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
         self.pool_executor = ThreadPoolExecutor(max_workers=4)
 
     def initialize(self):
-        if is_jupyter():
+        if in_notebook():
             import nest_asyncio
 
             nest_asyncio.apply()