fix: made several slow imports and top-level calls lazy (#97)

Makefile

@@ -104,3 +104,6 @@ cloc-server:
 
 cloc-tests:
 	cloc tests/*.py --by-file
+
+bench-importtime:
+	time poetry run python -X importtime -m gptme --model openrouter --non-interactive 2>&1 | grep "import time" | cut -d'|' -f 2- | sort -n

gptme/llm_anthropic.py

@@ -1,25 +1,28 @@
 from collections.abc import Generator
-from typing import Literal, TypedDict
-from typing_extensions import Required
+from typing import TYPE_CHECKING, Literal, TypedDict
 
-from anthropic import Anthropic
+from typing_extensions import Required
 
 from .constants import TEMPERATURE, TOP_P
 from .message import Message, len_tokens, msgs2dicts
 
-anthropic: Anthropic | None = None
+if TYPE_CHECKING:
+    from anthropic import Anthropic
+
+anthropic: "Anthropic | None" = None
 
 
 def init(config):
     global anthropic
     api_key = config.get_env_required("ANTHROPIC_API_KEY")
+    from anthropic import Anthropic  # fmt: skip
     anthropic = Anthropic(
         api_key=api_key,
         max_retries=5,
     )
 
 
-def get_client() -> Anthropic | None:
+def get_client() -> "Anthropic | None":
     return anthropic
 
 

gptme/llm_openai.py

@@ -1,12 +1,14 @@
 import logging
 from collections.abc import Generator
-
-from openai import AzureOpenAI, OpenAI
+from typing import TYPE_CHECKING
 
 from .constants import TEMPERATURE, TOP_P
 from .message import Message, msgs2dicts
 
-openai: OpenAI | None = None
+if TYPE_CHECKING:
+    from openai import OpenAI
+
+openai: "OpenAI | None" = None
 logger = logging.getLogger(__name__)
 
 
@@ -19,6 +21,7 @@
 
 def init(llm: str, config):
     global openai
+    from openai import AzureOpenAI, OpenAI  # fmt: skip
 
     if llm == "openai":
         api_key = config.get_env_required("OPENAI_API_KEY")
@@ -44,7 +47,7 @@ def init(llm: str, config):
     assert openai, "LLM not initialized"
 
 
-def get_client() -> OpenAI | None:
+def get_client() -> "OpenAI | None":
     return openai
 
 

gptme/tools/__init__.py

@@ -1,6 +1,7 @@
 import logging
 from collections.abc import Generator
 from dataclasses import dataclass
+from typing import Callable
 from xml.etree import ElementTree
 
 from ..message import Message
@@ -9,8 +10,9 @@
 from .browser import tool as browser_tool
 from .gh import tool as gh_tool
 from .patch import tool as patch_tool
-from .python import execute_python, register_function
-from .python import tool as python_tool
+from .python import execute_python
+from .python import get_tool as get_python_tool
+from .python import register_function
 from .read import tool as tool_read
 from .save import execute_save, tool_append, tool_save
 from .shell import execute_shell
@@ -33,22 +35,17 @@
     "all_tools",
 ]
 
-
-all_tools: list[ToolSpec] = [
-    tool
-    for tool in [
-        tool_read,
-        tool_save,
-        tool_append,
-        patch_tool,
-        python_tool,
-        shell_tool,
-        subagent_tool,
-        tmux_tool,
-        browser_tool,
-        gh_tool,
-    ]
-    if tool.available
+all_tools: list[ToolSpec | Callable[[], ToolSpec]] = [
+    tool_read,
+    tool_save,
+    tool_append,
+    patch_tool,
+    get_python_tool,
+    shell_tool,
+    subagent_tool,
+    tmux_tool,
+    browser_tool,
+    gh_tool,
 ]
 loaded_tools: list[ToolSpec] = []
 
@@ -69,6 +66,10 @@ def execute(self, ask: bool) -> Generator[Message, None, None]:
 def init_tools() -> None:
     """Runs initialization logic for tools."""
     for tool in all_tools:
+        if not isinstance(tool, ToolSpec):
+            tool = tool()
+        if not tool.available:
+            continue
         if tool in loaded_tools:
             continue
         load_tool(tool)
@@ -204,7 +205,7 @@ def get_tooluse_xml(content: str) -> Generator[ToolUse, None, None]:
 
 def get_tool(tool_name: str) -> ToolSpec:
     """Returns a tool by name."""
-    for tool in all_tools:
+    for tool in loaded_tools:
         if tool.name == tool_name:
             return tool
     raise ValueError(f"Tool '{tool_name}' not found")

gptme/tools/python.py

@@ -10,78 +10,13 @@
 from logging import getLogger
 from typing import Literal, TypeVar, get_origin
 
-from IPython.terminal.embed import InteractiveShellEmbed
-from IPython.utils.capture import capture_output
-
 from ..message import Message
 from ..util import ask_execute, print_preview, transform_examples_to_chat_directives
 from .base import ToolSpec
 
 logger = getLogger(__name__)
 
 
-@functools.lru_cache
-def get_installed_python_libraries() -> set[str]:
-    """Check if a select list of Python libraries are installed."""
-    candidates = [
-        "numpy",
-        "pandas",
-        "matplotlib",
-        "seaborn",
-        "scipy",
-        "scikit-learn",
-        "statsmodels",
-        "pillow",
-    ]
-    installed = set()
-    for candidate in candidates:
-        try:
-            __import__(candidate)
-            installed.add(candidate)
-        except ImportError:
-            pass
-    return installed
-
-
-python_libraries = get_installed_python_libraries()
-python_libraries_str = "\n".join(f"- {lib}" for lib in python_libraries)
-
-
-instructions = f"""
-When you send a message containing Python code (and is not a file block), it will be executed in a stateful environment.
-Python will respond with the output of the execution.
-
-The following libraries are available:
-{python_libraries_str}
-""".strip()
-
-# TODO: get this working again (needs to run get_functions_prompt() after all functions are registered)
-_unused = """
-The following functions are available in the REPL:
-{get_functions_prompt()}
-"""
-
-examples = """
-#### Results of the last expression will be displayed, IPython-style:
-User: What is 2 + 2?
-Assistant:
-```ipython
-2 + 2
-```
-System: Executed code block.
-```stdout
-4
-```
-
-#### The user can also run Python code with the /python command:
-
-User: /python 2 + 2
-System: Executed code block.
-```stdout
-4
-```
-""".strip()
-
 # IPython instance
 _ipython = None
 
@@ -131,6 +66,7 @@ def get_functions_prompt() -> str:
 
 def _get_ipython():
     global _ipython
+    from IPython.terminal.embed import InteractiveShellEmbed  # fmt: skip
     if _ipython is None:
         _ipython = InteractiveShellEmbed()
         _ipython.push(registered_functions)
@@ -156,6 +92,7 @@ def execute_python(code: str, ask: bool, args=None) -> Generator[Message, None,
     _ipython = _get_ipython()
 
     # Capture the standard output and error streams
+    from IPython.utils.capture import capture_output  # fmt: skip
     with capture_output() as captured:
         # Execute the code
         result = _ipython.run_cell(code, silent=False, store_history=False)
@@ -186,6 +123,29 @@ def execute_python(code: str, ask: bool, args=None) -> Generator[Message, None,
     yield Message("system", "Executed code block.\n\n" + output)
 
 
+@functools.lru_cache
+def get_installed_python_libraries() -> set[str]:
+    """Check if a select list of Python libraries are installed."""
+    candidates = [
+        "numpy",
+        "pandas",
+        "matplotlib",
+        "seaborn",
+        "scipy",
+        "scikit-learn",
+        "statsmodels",
+        "pillow",
+    ]
+    installed = set()
+    for candidate in candidates:
+        try:
+            __import__(candidate)
+            installed.add(candidate)
+        except ImportError:
+            pass
+    return installed
+
+
 def check_available_packages():
     """Checks that essentials like numpy, pandas, matplotlib are available."""
     expected = ["numpy", "pandas", "matplotlib"]
@@ -199,17 +159,57 @@ def check_available_packages():
         )
 
 
+examples = """
+#### Results of the last expression will be displayed, IPython-style:
+User: What is 2 + 2?
+Assistant:
+```ipython
+2 + 2
+```
+System: Executed code block.
+```stdout
+4
+```
+
+#### The user can also run Python code with the /python command:
+
+User: /python 2 + 2
+System: Executed code block.
+```stdout
+4
+```
+""".strip()
+
 __doc__ += transform_examples_to_chat_directives(examples)
 
-tool = ToolSpec(
-    name="python",
-    desc="Execute Python code",
-    instructions=instructions,
-    examples=examples,
-    init=init_python,
-    execute=execute_python,
-    block_types=[
-        "python",
-        "ipython",
-    ],  # ideally, models should use `ipython` and not `python`, but they don't
-)
+
+def get_tool() -> ToolSpec:
+    python_libraries = get_installed_python_libraries()
+    python_libraries_str = "\n".join(f"- {lib}" for lib in python_libraries)
+
+    instructions = f"""
+    When you send a message containing Python code (and is not a file block), it will be executed in a stateful environment.
+    Python will respond with the output of the execution.
+
+    The following libraries are available:
+    {python_libraries_str}
+    """.strip()
+
+    # TODO: get this working again (needs to run get_functions_prompt() after all functions are registered)
+    _unused = """
+    The following functions are available in the REPL:
+    {get_functions_prompt()}
+    """
+
+    return ToolSpec(
+        name="python",
+        desc="Execute Python code",
+        instructions=instructions,
+        examples=examples,
+        init=init_python,
+        execute=execute_python,
+        block_types=[
+            "python",
+            "ipython",
+        ],  # ideally, models should use `ipython` and not `python`, but they don't
+    )