perf(startup): defer provider SDK imports to first use (#372)

Author: mbektas

CHANGELOG.md

@@ -10,6 +10,10 @@ For each release we list user-facing changes grouped as **Added**, **Changed**,
 
 ## [Unreleased]
 
+### Changed
+
+- **Provider SDKs load on first use instead of at module import** (#370). `import notebook_intelligence` no longer imports `litellm`, `openai`, `ollama`, or the `anthropic` SDK; `litellm`, `openai`, and `anthropic` load the first time their provider is actually used (for Claude mode that includes the client construction and model refresh NBI runs at startup), while `ollama` still loads during extension startup when the provider enumerates local models. This roughly halves the server-extension import time (a cost the Jupyter server pays on every start), with the biggest effect on Windows machines where antivirus scanning amplifies the many-small-file SDK imports (#368). When NBI does load litellm, it now defaults `LITELLM_LOCAL_MODEL_COST_MAP=true` so litellm reads its bundled model-cost map rather than fetching it over HTTP at import; set the env var to `false` to restore the fetch.
+
 ## [5.1.0] - UNRELEASED
 
 5.1.0 builds on 5.0.x with a focus on Claude-mode agent visibility. Tool calls the agent runs now render as persistent status cards with inline diffs and collapsible grouping, the generating indicator can cycle custom verbs, and cancelling a turn tears down the whole process tree the agent spawned instead of leaking it. It also adds two opt-in security guardrails (an MCP stdio-command allowlist and a default-token-password check on shared filesystems) and an always-visible mode for chat feedback. No traitlet, env-var, REST route, or on-disk-format renames or removals; every new admin surface is opt-in and listed below.

docs/admin-guide.md

@@ -122,6 +122,7 @@ The full surface, in one table.
 | `NBI_GHE_SUBDOMAIN`                              | str  | `""`                        | env                                | GitHub Enterprise subdomain for GitHub Copilot users on a GHE tenant. Empty selects github.com.                                                                                                                                                                                                                                                                                                                                                                          |
 | `NBI_GITHUB_ENTERPRISE_HOSTS`                    | csv  | `""`                        | env                                | Comma-separated hostnames the plugin marketplace detector treats as GitHub. Cookie-domain shape: bare token (`github.acme.com`) matches exactly; leading-dot token (`.acme.com`) matches any subdomain of `acme.com`. Independent of `NBI_GHE_SUBDOMAIN`, which only configures the Copilot OAuth tenant. Required so `allow_github_plugin_import = False` actually gates GHE marketplace adds and so the `GITHUB_TOKEN` / `gh auth token` chain injects on GHE sources. |
 | `NBI_LOG_LEVEL`                                  | str  | `INFO`                      | env                                | Python logging level for the `notebook_intelligence` logger.                                                                                                                                                                                                                                                                                                                                                                                                             |
+| `LITELLM_LOCAL_MODEL_COST_MAP`                   | bool | `true` (NBI default)        | env                                | litellm setting that NBI defaults to `true` when it loads litellm, so litellm reads the model-cost map bundled with the installed package instead of fetching it over HTTP (litellm's own default), which stalls on proxied networks. Set to `false` before starting JupyterLab to restore litellm's remote fetch.                                                                                                                                                       |
 | `NBI_DISABLE_OUTPUT_SCRUB`                       | bool | unset                       | env                                | When set (`1` / `true` / `yes` / `on`), disables the shell-tool output scrubber so raw stdout/stderr (including any env-var values that leak) is sent through to chat. Default off; the scrubber redacts values for sensitive-named env vars (`TOKEN`, `SECRET`, `API_KEY`, ...) plus tokens with well-known credential prefixes (`ghp_`, `sk-ant-`, `AKIA`, ...). Opt out only when debugging credential helpers where the redaction interferes.                        |
 | `GITHUB_TOKEN`, `GH_TOKEN`                       | str  | unset                       | env                                | Used (in that order) by user-initiated skill imports and GitHub-sourced plugin marketplace adds for GitHub auth. Falls back to `gh` CLI auth.                                                                                                                                                                                                                                                                                                                            |
 | `NBI_*_POLICY`                                   | str  | `user-choice`               | env                                | Lock individual Settings panel toggles. See [README → Admin policies](../README.md#admin-policies) for the full list of `*_POLICY` env vars and matching traitlets, including `NBI_SKILLS_MANAGEMENT_POLICY`, `NBI_CLAUDE_MCP_MANAGEMENT_POLICY`, `NBI_CLAUDE_PLUGINS_MANAGEMENT_POLICY`, `NBI_TERMINAL_DRAG_DROP_POLICY`, and `NBI_REFRESH_OPEN_FILES_ON_DISK_CHANGE_POLICY`.                                                                                           |

notebook_intelligence/claude.py

@@ -10,20 +10,21 @@
 from queue import Queue
 import threading
 import time
-from typing import Any, Optional
+from typing import Any, Optional, TYPE_CHECKING
 import uuid
 import re
 from anyio.abc import Process
-from anthropic import Anthropic
 from notebook_intelligence.api import AskUserQuestionData, BackendMessageType, CancelToken, ChatCommand, ChatModel, ChatRequest, ChatResponse, ClaudeToolType, CompletionContext, ConfirmationData, Host, InlineCompletionModel, MarkdownData, ProgressData, SignalImpl, ToolCallData
 from notebook_intelligence.base_chat_participant import BaseChatParticipant
 from notebook_intelligence._version import __version__ as NBI_VERSION
 import base64
 import logging
 from claude_agent_sdk import AssistantMessage, PermissionResultAllow, PermissionResultDeny, TextBlock, ToolResultBlock, ToolUseBlock, UserMessage, create_sdk_mcp_server, ClaudeAgentOptions, ClaudeSDKClient, tool
-from anthropic.types.text_block import TextBlock as AnthropicTextBlock
 
-from notebook_intelligence.util import ThreadSafeWebSocketConnector, _emit, get_jupyter_root_dir, resolve_claude_cli_path, safe_jupyter_path, terminate_process_tree
+from notebook_intelligence.util import ThreadSafeWebSocketConnector, _emit, get_jupyter_root_dir, import_litellm, resolve_claude_cli_path, safe_jupyter_path, terminate_process_tree
+
+if TYPE_CHECKING:
+    from anthropic import Anthropic
 
 log = logging.getLogger(__name__)
 
@@ -345,7 +346,7 @@ def get_claude_models() -> list[dict]:
 def _get_context_window(model_id: str) -> int:
     """Get context window size for a model using litellm's model database."""
     try:
-        import litellm
+        litellm = import_litellm()
         info = litellm.get_model_info(model_id)
         return info.get("max_input_tokens", 200000)
     except Exception:
@@ -364,8 +365,9 @@ def _normalize_anthropic_credential(value: Any) -> str | None:
     return value.strip() or None
 
 
-def _create_anthropic_client(api_key: str = None, base_url: str = None) -> Anthropic:
+def _create_anthropic_client(api_key: str = None, base_url: str = None) -> "Anthropic":
     """Create an Anthropic client with normalized credentials and default headers."""
+    from anthropic import Anthropic
     api_key = _normalize_anthropic_credential(api_key)
     base_url = _normalize_anthropic_credential(base_url)
     return Anthropic(
@@ -544,6 +546,8 @@ def inline_completions(self, prefix, suffix, language, filename, context: Comple
         if cancel_token.is_cancel_requested:
             return ''
 
+        from anthropic.types.text_block import TextBlock as AnthropicTextBlock
+
         message = self._client.messages.create(
             model=self._model_id,
             max_tokens=10000,

notebook_intelligence/llm_providers/litellm_compatible_llm_provider.py

@@ -3,7 +3,7 @@
 import json
 from typing import Any
 from notebook_intelligence.api import ChatModel, EmbeddingModel, InlineCompletionModel, LLMProvider, CancelToken, ChatResponse, CompletionContext, LLMProviderProperty
-import litellm
+from notebook_intelligence.util import import_litellm
 
 DEFAULT_CONTEXT_WINDOW = 4096
 
@@ -37,6 +37,7 @@ def context_window(self) -> int:
             return DEFAULT_CONTEXT_WINDOW
 
     def completions(self, messages: list[dict], tools: list[dict] = None, response: ChatResponse = None, cancel_token: CancelToken = None, options: dict = {}) -> Any:
+        litellm = import_litellm()
         stream = response is not None
         model_id = self.get_property("model_id").value
         base_url = self.get_property("base_url").value
@@ -111,6 +112,7 @@ def context_window(self) -> int:
             return DEFAULT_CONTEXT_WINDOW
 
     def inline_completions(self, prefix, suffix, language, filename, context: CompletionContext, cancel_token: CancelToken) -> str:
+        litellm = import_litellm()
         model_id = self.get_property("model_id").value
         base_url = self.get_property("base_url").value
         api_key_prop = self.get_property("api_key")

notebook_intelligence/llm_providers/ollama_llm_provider.py

@@ -3,7 +3,6 @@
 import json
 from typing import Any
 from notebook_intelligence.api import ChatModel, EmbeddingModel, InlineCompletionModel, LLMProvider, CancelToken, ChatResponse, CompletionContext
-import ollama
 import logging
 
 from notebook_intelligence.util import extract_llm_generated_code
@@ -37,6 +36,7 @@ def context_window(self) -> int:
         return self._context_window
 
     def completions(self, messages: list[dict], tools: list[dict] = None, response: ChatResponse = None, cancel_token: CancelToken = None, options: dict = {}) -> Any:
+        import ollama
         stream = response is not None
         completion_args = {
             "model": self._model_id, 
@@ -102,6 +102,7 @@ def context_window(self) -> int:
         return self._context_window
 
     def inline_completions(self, prefix, suffix, language, filename, context: CompletionContext, cancel_token: CancelToken) -> str:
+        import ollama
         has_suffix = suffix.strip() != ""
         if has_suffix:
             prompt = self._prompt_template.format(prefix=prefix, suffix=suffix.strip())
@@ -171,6 +172,7 @@ def embedding_models(self) -> list[EmbeddingModel]:
 
     def update_chat_model_list(self):
         try:
+            import ollama
             response = ollama.list()
             models = response.models
             self._chat_models = []

notebook_intelligence/llm_providers/openai_compatible_llm_provider.py

@@ -5,7 +5,6 @@
 import re
 from typing import Any
 from notebook_intelligence.api import ChatModel, EmbeddingModel, InlineCompletionModel, LLMProvider, CancelToken, ChatResponse, CompletionContext, LLMProviderProperty
-from openai import OpenAI, omit
 
 INLINE_COMPLETION_SYSTEM_PROMPT = """You are a code completion assistant. Your task is to generate intelligent autocomplete suggestions for the code at the cursor position for given language and active file type. This is not an interactive session, don't ask for clarifying questions, always generate a suggestion. Don't include any explanations for your response, just generate the code. Don't return any thinking or reasoning, just generate the code. You are given a code snippet with a prefix and a suffix. You need to generate a suggestion for the code that fits best in place of <CURSOR/>. You should return only the code that fits best in place of <CURSOR/>. You should provide multiline code if needed. Enclose the code in triple backticks, just return the code in language. You should not return any other text, just the code. DO NOT INCLUDE THE PREFIX OR SUFFIX IN THE RESPONSE. .ipynb files are Jupyter notebook files and for notebook files, you generate suggestions for a cell within the notebook. A cell can be a code cell with code or a markdown cell with markdown text. If the language is markdown, only return markdown text. If you need to install a Python package within a notebook cell code (for .ipynb files), use %pip install <package_name> instead of !pip install <package_name>. Follow the tags very carefully for proper spacing and indentations."""
 
@@ -55,6 +54,7 @@ def context_window(self) -> int:
             return DEFAULT_CONTEXT_WINDOW
 
     def completions(self, messages: list[dict], tools: list[dict] = None, response: ChatResponse = None, cancel_token: CancelToken = None, options: dict = {}) -> Any:
+        from openai import OpenAI, omit
         stream = response is not None
         model_id = self.get_property("model_id").value
         base_url_prop = self.get_property("base_url")
@@ -152,6 +152,8 @@ def inline_completions(self, prefix, suffix, language, filename, context: Comple
         if cancel_token.is_cancel_requested:
             return ''
 
+        from openai import OpenAI
+
         model_id = self.get_property("model_id").value
         base_url_prop = self.get_property("base_url")
         base_url = base_url_prop.value if base_url_prop is not None else None

notebook_intelligence/util.py

@@ -26,6 +26,20 @@
 _jupyter_root_dir: str = None
 _enabled_tools: Set[str] = None
 
+def import_litellm():
+    """Import litellm on first use instead of at server-extension import.
+
+    litellm is over a second of import time (the largest single chunk of
+    NBI's startup) and, by default, fetches its model-cost map over HTTP at
+    import, which can stall on proxied networks. Defaulting
+    LITELLM_LOCAL_MODEL_COST_MAP to true makes it read the map bundled with
+    the installed litellm instead; set the env var to false to restore the
+    remote fetch.
+    """
+    os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "true")
+    import litellm
+    return litellm
+
 def set_jupyter_root_dir(root_dir: str):
     global _jupyter_root_dir
     _jupyter_root_dir = root_dir

pyproject.toml

@@ -92,6 +92,14 @@ test = ["pytest>=7.0,<9", "pytest-timeout>=2.3.0", "psutil", "pyyaml"]
 # job. Other tests in the suite run in milliseconds, so 30s is
 # generous headroom rather than a target.
 timeout = 30
+# Make `from tests.conftest import ...` resolvable on its own. Without
+# this, those imports only worked because the eager `import litellm`
+# pulled in litellm/proxy/proxy_cli.py, which calls
+# sys.path.append(os.getcwd()) at import time; deferring the litellm
+# import (#370) removed that accidental path entry and broke collection
+# under a non-editable install (editable installs already put the repo
+# root on sys.path, which is why local runs never noticed).
+pythonpath = ["."]
 
 [tool.hatch.version]
 source = "nodejs"

tests/test_claude_models.py

@@ -78,7 +78,7 @@ def _make_mock_model(self, model_id, display_name):
         return m
 
     @patch("notebook_intelligence.claude._get_context_window", return_value=200000)
-    @patch("notebook_intelligence.claude.Anthropic")
+    @patch("anthropic.Anthropic")
     def test_fetches_and_caches_models(self, mock_anthropic_cls, mock_ctx_window):
         from notebook_intelligence.claude import fetch_claude_models, get_claude_models
 
@@ -98,7 +98,7 @@ def test_fetches_and_caches_models(self, mock_anthropic_cls, mock_ctx_window):
         assert len(get_claude_models()) == 2
 
     @patch("notebook_intelligence.claude._get_context_window", return_value=150000)
-    @patch("notebook_intelligence.claude.Anthropic")
+    @patch("anthropic.Anthropic")
     def test_uses_litellm_context_window(self, mock_anthropic_cls, mock_ctx_window):
         from notebook_intelligence.claude import fetch_claude_models
 
@@ -112,7 +112,7 @@ def test_uses_litellm_context_window(self, mock_anthropic_cls, mock_ctx_window):
         mock_ctx_window.assert_called_once_with("claude-sonnet-4-6")
 
     @patch("notebook_intelligence.claude._get_context_window", return_value=200000)
-    @patch("notebook_intelligence.claude.Anthropic")
+    @patch("anthropic.Anthropic")
     def test_cache_is_mutated_in_place(self, mock_anthropic_cls, mock_ctx_window):
         """Verify the cache list is mutated, not replaced, so importers keep a valid reference."""
         from notebook_intelligence.claude import fetch_claude_models, _claude_models_cache
@@ -129,7 +129,7 @@ def test_cache_is_mutated_in_place(self, mock_anthropic_cls, mock_ctx_window):
         assert len(original_list) == 1
 
     @patch("notebook_intelligence.claude._get_context_window", return_value=200000)
-    @patch("notebook_intelligence.claude.Anthropic")
+    @patch("anthropic.Anthropic")
     def test_clears_old_cache_on_refresh(self, mock_anthropic_cls, mock_ctx_window):
         from notebook_intelligence.claude import fetch_claude_models, get_claude_models, _claude_models_cache
 
@@ -147,7 +147,7 @@ def test_clears_old_cache_on_refresh(self, mock_anthropic_cls, mock_ctx_window):
         assert len(models) == 1
         assert models[0]["id"] == "new-model"
 
-    @patch("notebook_intelligence.claude.Anthropic")
+    @patch("anthropic.Anthropic")
     def test_returns_existing_cache_on_api_failure(self, mock_anthropic_cls):
         from notebook_intelligence.claude import fetch_claude_models, get_claude_models, _claude_models_cache
 
@@ -162,7 +162,7 @@ def test_returns_existing_cache_on_api_failure(self, mock_anthropic_cls):
         assert result[0]["id"] == "cached-model"
 
     @patch("notebook_intelligence.claude._get_context_window", return_value=200000)
-    @patch("notebook_intelligence.claude.Anthropic")
+    @patch("anthropic.Anthropic")
     def test_empty_api_key_passed_as_none(self, mock_anthropic_cls, mock_ctx_window):
         from notebook_intelligence.claude import fetch_claude_models