From 7b7552029f560be955104165ef531f2ab19733c2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Thu, 10 Oct 2024 13:15:02 +0200
Subject: [PATCH] fix: use importlib.util.find_spec instead of attempting
 costly imports at startup, reducing startup time by ~1s

---
 gptme/tools/python.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/gptme/tools/python.py b/gptme/tools/python.py
index 61f8f689..80a1e863 100644
--- a/gptme/tools/python.py
+++ b/gptme/tools/python.py
@@ -6,6 +6,7 @@
 
 import dataclasses
 import functools
+import importlib.util
 import re
 import types
 from collections.abc import Callable, Generator
@@ -43,6 +44,9 @@
 def register_function(func: T) -> T:
     """Decorator to register a function to be available in the IPython instance."""
     registered_functions[func.__name__] = func
+    # if ipython is already initialized, push the function to it to make it available
+    if _ipython is not None:
+        _ipython.push({func.__name__: func})
     return func
 
 
@@ -150,17 +154,15 @@ def get_installed_python_libraries() -> set[str]:
         "matplotlib",
         "seaborn",
         "scipy",
-        "scikit-learn",
+        "sklearn",
         "statsmodels",
-        "pillow",
+        "PIL",
     ]
     installed = set()
     for candidate in candidates:
-        try:
-            __import__(candidate)
+        if importlib.util.find_spec(candidate):
             installed.add(candidate)
-        except ImportError:
-            pass
+
     return installed