Separate BaseProvider for faster import

[krassowski]

References

• Closes Performance: move BaseProvider to a separate python module #1337

User-facing changes

IPython users who define a custom model by inheriting from BaseProvider will be able to start IPython about 33% faster; this is significant both in relative and absolute terms (while for me it is only 0.4 faster, I run the benchmarks on a very fast CPU which is not necessarily representative of the general user base).

Import	Before	After
from jupyter_ai_magics import BaseProvider	1.342 s ± 0.021 s	1.046 s ± 0.014 s
from jupyter_ai_magics.providers import BaseProvider	1.344 s ± 0.016 s	1.049 s ± 0.010 s
from jupyter_ai_magics.base_provider import BaseProvider	(not possible)	902.5 ms ± 10.0 ms

Before

$ hyperfine "python -c 'from jupyter_ai_magics import BaseProvider'"
Benchmark 1: python -c 'from jupyter_ai_magics import BaseProvider'
  Time (mean ± σ):      1.342 s ±  0.021 s    [User: 1.188 s, System: 0.157 s]
  Range (min … max):    1.315 s …  1.388 s    10 runs
$ hyperfine "python -c 'from jupyter_ai_magics.providers import BaseProvider'"
Benchmark 1: python -c 'from jupyter_ai_magics.providers import BaseProvider'
  Time (mean ± σ):      1.344 s ±  0.016 s    [User: 1.194 s, System: 0.153 s]
  Range (min … max):    1.318 s …  1.368 s    10 runs

After

$ hyperfine "python -c 'from jupyter_ai_magics import BaseProvider'"
Benchmark 1: python -c 'from jupyter_ai_magics import BaseProvider'
  Time (mean ± σ):      1.046 s ±  0.014 s    [User: 0.914 s, System: 0.136 s]
  Range (min … max):    1.022 s …  1.072 s    10 runs
$ hyperfine "python -c 'from jupyter_ai_magics.providers import BaseProvider'"
Benchmark 1: python -c 'from jupyter_ai_magics.providers import BaseProvider'
  Time (mean ± σ):      1.049 s ±  0.010 s    [User: 0.924 s, System: 0.129 s]
  Range (min … max):    1.028 s …  1.058 s    10 runs
$ hyperfine "python -c 'from jupyter_ai_magics.base_provider import BaseProvider'"
Benchmark 1: python -c 'from jupyter_ai_magics.base_provider import BaseProvider'
  Time (mean ± σ):     902.5 ms ±  10.0 ms    [User: 788.1 ms, System: 118.4 ms]
  Range (min … max):   887.7 ms … 914.9 ms    10 runs

jupyter-ai-magics users who use %load_ext jupyter_ai_magics at startup will not see a speedup yet because self.providers = get_lm_providers() in AiMagics init causes all providers to be loaded. They would only see a speedup if they disabled some of the providers (which previously would not have been the case as all providers would be imported anyways).

To be done

• vendor import_attr instead of importing it from langchain_core._import_utils to make sure it does not break in a point release of langchain (preserving the attribution to langchain team)
• fix typing actually, mypy is failing on other PRs too..

[dlqqq]

@krassowski Thank you for opening this PR! Really appreciate you taking the lead on debugging the module import latency & for providing objective data. I remember you had mentioned in another thread that LangChain recently improved their module imports to be faster; do you know if that is as simple as upgrading our version floor on langchain? Would be curious to see if that makes an impact too.

I've been tasked with a 2-week coding sprint to work on another project right now. Sorry if I'm slow to respond for the next week. I'll be back on Monday to give this a full review.

[dlqqq]

This static assertion is great, but it won't run until the module is imported. Can we add a unit test for this to guarantee that CI will block release if this assertion excepts?

def test_top_level_module():
  from jupyter_ai_magics import * as exports
  assert hasattr(exports, 'BaseProvider')
  ... # etc

I know there are other tests that import from jupyter_ai_magics, but having an explicit one reduces the chance that this causes a runtime error for users even if we disable/skip other tests in the future.

[dlqqq]

This looks like a private module, which means it may break in patch releases. Can we re-implement this in jupyter_ai_magics to prevent that in advance?

[krassowski]

Yes, that's the plan! I took a shortcut for PoC.

[dlqqq]

When consumers import from jupyter_ai_magics, do they still get type inference for mypy & IDE intellisense? I would test this myself but I'm short on time this week.

[krassowski]

It depends on the editor. The explicit __all__ helps, but yes not all editros will autocomplete. I think adding typing stubs could help but I am not sure if it is worth the effort.

I have not checked mypy yet.

We could recommend importing from specific modules instead of the top level which will in 100% preserve completions and mypy checks (and in case of base provider be faster too).

[krassowski]

I remember you had mentioned in another thread that LangChain recently improved their module imports to be faster; do you know if that is as simple as upgrading our version floor on langchain?

Yup, this was langchain-core, and the timings in this PR are based on the latest version including the fix. In fact this PR in the current state directly imports the lazy import machinery from langchain_core (which we should not do as it is private, we should vendor a copy instead), because this approach is required to ensure the imports are not triggered in __init__.py.

Separately of this PR if you want to ensure all users get the faster imports then yes adding floor level pin should be enough as it was released in a minor version.

Have a good coding spring! I saw your status which is why I did not ping - I really appreciate a quick look from you!

[dlqqq]

BTW, do you have a good reference on how Python module imports work? I'd definitely like to review those resources before I review your PR fully. I'm surprised that you & LangChain maintainers had to add manual code changes to get dynamic fast imports.

[krassowski]

The relevant bits are:

• https://peps.python.org/pep-0562/ looking at it we may want to implement __dir__ too
• https://docs.python.org/3/reference/import.html#regular-packages