[WIP] Migrate from LangChain to LiteLLM (major upgrade)

.github/workflows/lint.yml

@@ -3,8 +3,8 @@ on:
   - pull_request
 
 jobs:
-  lint_ts:
-    name: Lint TypeScript source
+  lint_frontend:
+    name: Stylelint, Prettier, ESLint
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -17,3 +17,7 @@ jobs:
         run: jlpm
       - name: Lint TypeScript source
         run: jlpm lerna run lint:check
+
+  # TODO
+  # precommit:
+  #   name: pre-commit check

.github/workflows/unit-tests.yml

@@ -19,7 +19,7 @@ jobs:
       matrix:
         include:
           - dependency-type: minimum
-            python-version: "3.9"
+            python-version: "3.10"
           - dependency-type: standard
             python-version: "3.12"
     steps:

.gitignore

@@ -142,3 +142,6 @@ packages/**/_version.py
 # Ignore local chat files & local .jupyter/ dir
 *.chat
 .jupyter/
+
+# Ignore secrets in '.env'
+.env

conftest.py

@@ -1,13 +1,20 @@
+from __future__ import annotations
+import asyncio
 from pathlib import Path
-
 import pytest
+from traitlets.config import Config, LoggingConfigurable
+import logging
+from jupyter_server.services.contents.filemanager import AsyncFileContentsManager
+from typing import TYPE_CHECKING
 
-pytest_plugins = ("jupyter_server.pytest_plugin",)
+if TYPE_CHECKING:
+    from jupyter_server.serverapp import ServerApp
 
+pytest_plugins = ("jupyter_server.pytest_plugin",)
 
 @pytest.fixture
-def jp_server_config(jp_server_config):
-    return {"ServerApp": {"jpserver_extensions": {"jupyter_ai": True}}}
+def jp_server_config(jp_server_config, tmp_path):
+    return Config({"ServerApp": {"jpserver_extensions": {"jupyter_ai": True}}, "ContentsManager": {"root_dir": str(tmp_path)}})
 
 
 @pytest.fixture(scope="session")
@@ -21,9 +28,45 @@ def static_test_files_dir() -> Path:
         / "static"
     )
 
+class MockAiExtension(LoggingConfigurable):
+    """Mock AiExtension class for testing purposes."""
+
+    serverapp: ServerApp
+
+    def __init__(self, *args, serverapp: ServerApp, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.serverapp = serverapp
+        self._log = None
+
+    @property
+    def log(self) -> logging.Logger:
+        return self.serverapp.log
+
+    @property
+    def event_loop(self) -> asyncio.AbstractEventLoop:
+        return self.serverapp.io_loop.asyncio_loop
+
+    @property
+    def contents_manager(self) -> AsyncFileContentsManager:
+        return self.serverapp.contents_manager
+
 
 @pytest.fixture
-def jp_ai_staging_dir(jp_data_dir: Path) -> Path:
-    staging_area = jp_data_dir / "scheduler_staging_area"
-    staging_area.mkdir()
-    return staging_area
+def mock_ai_extension(jp_server_config, jp_configurable_serverapp) -> MockAiExtension:
+    """
+    Returns a mocked `AiExtension` object that can be passed as the `parent`
+    argument to objects normally initialized by `AiExtension`. This should be
+    passed to most of the "manager singletons" like `ConfigManager`,
+    `PersonaManager`, and `EnvSecretsManager`.
+
+    See `MockAiExtension` in `conftest.py` for a complete description of the
+    attributes, properties, and methods available. If something is missing,
+    please feel free to add to it in your PR.
+
+    Returns:
+        A `MockAiExtension` instance that can be passed as the `parent` argument
+        to objects normally initialized by `AiExtension`.
+    """
+    serverapp = jp_configurable_serverapp()
+    return MockAiExtension(config=jp_server_config, serverapp=serverapp)
+

docs/source/users/index.md

@@ -762,7 +762,7 @@ We currently support the following language model providers:
 To configure a default model you can use the IPython `%config` magic:
 
 ```python
-%config AiMagics.default_language_model = "anthropic:claude-v1.2"
+%config AiMagics.initial_language_model = "anthropic:claude-v1.2"
 ```
 
 Then subsequent magics can be invoked without typing in the model:
@@ -772,10 +772,10 @@ Then subsequent magics can be invoked without typing in the model:
 Write a poem about C++.
 ```
 
-You can configure the default model for all notebooks by specifying `c.AiMagics.default_language_model` tratilet in `ipython_config.py`, for example:
+You can configure the default model for all notebooks by specifying `c.AiMagics.initial_language_model` tratilet in `ipython_config.py`, for example:
 
 ```python
-c.AiMagics.default_language_model = "anthropic:claude-v1.2"
+c.AiMagics.initial_language_model = "anthropic:claude-v1.2"
 ```
 
 The location of `ipython_config.py` file is documented in [IPython configuration reference](https://ipython.readthedocs.io/en/stable/config/intro.html).
@@ -965,18 +965,18 @@ produced the following Python error:
 Write a new version of this code that does not produce that error.
 ```
 
-As a shortcut for explaining errors, you can use the `%ai error` command, which will explain the most recent error using the model of your choice.
+As a shortcut for explaining and fixing errors, you can use the `%ai fix` command, which will explain the most recent error using the model of your choice.
 
 ```
-%ai error anthropic:claude-v1.2
+%ai fix anthropic:claude-v1.2
 ```
 
 ### Creating and managing aliases
 
-You can create an alias for a model using the `%ai register` command. For example, the command:
+You can create an alias for a model using the `%ai alias` command. For example, the command:
 
 ```
-%ai register claude anthropic:claude-v1.2
+%ai alias claude anthropic:claude-v1.2
 ```
 
 will register the alias `claude` as pointing to the `anthropic` provider's `claude-v1.2` model. You can then use this alias as you would use any other model name:
@@ -1001,10 +1001,10 @@ prompt = PromptTemplate(
 chain = LLMChain(llm=llm, prompt=prompt)
 ```
 
-… and then use `%ai register` to give it a name:
+… and then use `%ai alias` to give it a name:
 
 ```
-%ai register companyname chain
+%ai alias companyname chain
 ```
 
 You can change an alias's target using the `%ai update` command:
@@ -1013,10 +1013,10 @@ You can change an alias's target using the `%ai update` command:
 %ai update claude anthropic:claude-instant-v1.0
 ```
 
-You can delete an alias using the `%ai delete` command:
+You can delete an alias using the `%ai dealias` command:
 
 ```
-%ai delete claude
+%ai dealias claude
 ```
 
 You can see a list of all aliases by running the `%ai list` command.
@@ -1103,7 +1103,7 @@ the selections they make in the settings panel will take precedence over these v
 
 Specify default language model
 ```bash
-jupyter lab --AiExtension.default_language_model=bedrock-chat:anthropic.claude-v2
+jupyter lab --AiExtension.initial_language_model=bedrock-chat:anthropic.claude-v2
 ```
 
 Specify default embedding model

packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py

@@ -1,89 +1,20 @@
+from __future__ import annotations
+
 from typing import TYPE_CHECKING
 
-from ._import_utils import import_attr as _import_attr
 from ._version import __version__
 
 if TYPE_CHECKING:
-    # same as dynamic imports but understood by mypy
-    from .embedding_providers import (
-        BaseEmbeddingsProvider,
-        GPT4AllEmbeddingsProvider,
-        HfHubEmbeddingsProvider,
-        QianfanEmbeddingsEndpointProvider,
-    )
-    from .exception import store_exception
-    from .magics import AiMagics
-    from .providers import (
-        AI21Provider,
-        BaseProvider,
-        GPT4AllProvider,
-        HfHubProvider,
-        QianfanProvider,
-        TogetherAIProvider,
-    )
-else:
-    _exports_by_module = {
-        "embedding_providers": [
-            "BaseEmbeddingsProvider",
-            "GPT4AllEmbeddingsProvider",
-            "HfHubEmbeddingsProvider",
-            "QianfanEmbeddingsEndpointProvider",
-        ],
-        "exception": ["store_exception"],
-        "magics": ["AiMagics"],
-        # expose model providers on the package root
-        "providers": [
-            "AI21Provider",
-            "BaseProvider",
-            "GPT4AllProvider",
-            "HfHubProvider",
-            "QianfanProvider",
-            "TogetherAIProvider",
-        ],
-    }
+    from IPython.core.interactiveshell import InteractiveShell
 
-    _modules_by_export = {
-        import_name: module
-        for module, imports in _exports_by_module.items()
-        for import_name in imports
-    }
-
-    def __getattr__(export_name: str) -> object:
-        module_name = _modules_by_export.get(export_name)
-        result = _import_attr(export_name, module_name, __spec__.parent)
-        globals()[export_name] = result
-        return result
 
+def load_ipython_extension(ipython: InteractiveShell):
+    from .exception import store_exception
+    from .magics import AiMagics
 
-def load_ipython_extension(ipython):
-    ipython.register_magics(__getattr__("AiMagics"))
-    ipython.set_custom_exc((BaseException,), __getattr__("store_exception"))
+    ipython.register_magics(AiMagics)
+    ipython.set_custom_exc((BaseException,), store_exception)
 
 
-def unload_ipython_extension(ipython):
+def unload_ipython_extension(ipython: InteractiveShell):
     ipython.set_custom_exc((BaseException,), ipython.CustomTB)
-
-
-# required to preserve backward compatibility with `from jupyter_ai_magics import *`
-__all__ = [
-    "__version__",
-    "load_ipython_extension",
-    "unload_ipython_extension",
-    "BaseEmbeddingsProvider",
-    "GPT4AllEmbeddingsProvider",
-    "HfHubEmbeddingsProvider",
-    "QianfanEmbeddingsEndpointProvider",
-    "store_exception",
-    "AiMagics",
-    "AI21Provider",
-    "BaseProvider",
-    "GPT4AllProvider",
-    "HfHubProvider",
-    "QianfanProvider",
-    "TogetherAIProvider",
-]
-
-
-def __dir__():
-    # Allows more editors (e.g. IPython) to complete on `jupyter_ai_magics.<tab>`
-    return list(__all__)