Proposal: Replacing LangChain with LiteLLM

[dlqqq]

Problem

There are several issues with using langchain as the model library of choice in JAI:

• Hard requirement on LangChain makes it difficult for contributors to extend model support

• Models are only listed if their provider dependencies are installed (Chat panel fails without selected provider dependencies installed #680, Inconsistent experience with respect to non-installed third-party packages #840)

• Magics & chat authentication are different ([v3-future] Enhance and generalize secret/key management? #1237)

• LangChain modules are slow to import & load from entry points (jupyter_ai import time is too slow #1115)

This issue proposes a solution to #1312.

Proposed Solution

We propose that Jupyter AI migrates to use litellm instead of langchain as a model library to provide universal LLM support.

litellm offers a much more simple dependency tree and is noticeably faster to import. It also provides access to a vast array of models out-of-the-box, without the need for partner packages. This entirely solves 3 major issues in JAI v2:

1. Dealing with provider dependencies: litellm provides everything out-of-the-box.
2. Forcing users to wait for us to update the list of models regularly for each provider: litellm maintains its own lists.
3. Forcing developers to create a package to add a model via entry points API: litellm supports custom models

As Jupyter AI v3 will only provide Jupyternaut, the requirement on langchain can be dropped if this proposal is accepted. For models that are not available through our litellm implementation, developers can write custom personas. These are easy to write and can be defined in any .jupyter/ directory.

Additional context

Also would fix #1308.

[dlqqq]

cc @krassowski

Wanted to loop you into our latest initiative. I know you were frustrated with Jupyter AI's import times, and this change should be a substantial improvement once implemented. Feedback welcomed! 👋

[ellisonbg]

I think this is a great direction, especially because we are getting the whole LangChain dependency right now just for the model abstraction. That is a lot of code to pull in just for model access. Obviously, we have to figure out all the details, but overall +1 to this

[dlqqq]

Since this involves a lot of changes across the frontend & backend, I will create a new litellm branch for everybody to contribute to. I will start that branch by an initial PR, and we can iterate on the litellm branch until it is ready to be merged into main. 👍

[dlqqq]

Closing this issue as no objections were raised against this proposal. Development will be tracked in the litellm feature branch, described in #1426.