Unable to use 1.x in SageMaker Studio

[dlqqq]

Description

In SageMaker Studio, running Python 3.9, jupyter-ai~=1.0 fails to initialize with the message:

[W 2023-11-20 23:05:51.475 ServerApp] jupyter_ai | error adding extension (enabled: True): issubclass() arg 1 must be a class
    Traceback (most recent call last):
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_server/extension/manager.py", line 320, in add_extension
        extpkg = ExtensionPackage(name=extension_name, enabled=enabled)
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_server/extension/manager.py", line 166, in __init__
        super().__init__(*args, **kwargs)
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/traitlets/traitlets.py", line 1238, in __init__
        super_kwargs[key] = value
      File "/opt/conda/envs/studio/lib/python3.9/contextlib.py", line 126, in __exit__
        next(self.gen)
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/traitlets/traitlets.py", line 1348, in hold_trait_notifications
        value = trait._cross_validate(self, getattr(self, name))
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/traitlets/traitlets.py", line 729, in _cross_validate
        value = obj._trait_validators[self.name](obj, proposal)
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/traitlets/traitlets.py", line 1132, in __call__
        return self.func(*args, **kwargs)
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_server/extension/manager.py", line 175, in _validate_name
        self._module, self._metadata = get_metadata(name)
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_server/extension/utils.py", line 55, in get_metadata
        module = importlib.import_module(package_name)
      File "/opt/conda/envs/studio/lib/python3.9/importlib/__init__.py", line 127, in import_module
        return _bootstrap._gcd_import(name[level:], package, level)
      File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
      File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
      File "<frozen importlib._bootstrap>", line 986, in _find_and_load_unlocked
      File "<frozen importlib._bootstrap>", line 680, in _load_unlocked
      File "<frozen importlib._bootstrap_external>", line 850, in exec_module
      File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_ai/__init__.py", line 2, in <module>
        from jupyter_ai_magics import load_ipython_extension, unload_ipython_extension
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_ai_magics/__init__.py", line 4, in <module>
        from .embedding_providers import (
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_ai_magics/embedding_providers.py", line 4, in <module>
        from jupyter_ai_magics.providers import (
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/jupyter_ai_magics/providers.py", line 21, in <module>
        from langchain.chat_models import (
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/langchain/chat_models/__init__.py", line 20, in <module>
        from langchain.chat_models.anthropic import ChatAnthropic
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/langchain/chat_models/anthropic.py", line 3, in <module>
        from langchain.callbacks.manager import (
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/langchain/callbacks/__init__.py", line 10, in <module>
        from langchain.callbacks.aim_callback import AimCallbackHandler
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/langchain/callbacks/aim_callback.py", line 5, in <module>
        from langchain.schema import AgentAction, AgentFinish, LLMResult
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/langchain/schema/__init__.py", line 2, in <module>
        from langchain.schema.agent import AgentAction, AgentFinish
      File "/opt/conda/envs/studio/lib/python3.9/site-packages/langchain/schema/agent.py", line 9, in <module>
        class AgentAction(Serializable):
      File "pydantic/main.py", line 198, in pydantic.main.ModelMetaclass.__new__
      File "pydantic/fields.py", line 506, in pydantic.fields.ModelField.infer
      File "pydantic/fields.py", line 436, in pydantic.fields.ModelField.__init__
      File "pydantic/fields.py", line 552, in pydantic.fields.ModelField.prepare
      File "pydantic/fields.py", line 668, in pydantic.fields.ModelField._type_analysis
      File "/opt/conda/envs/studio/lib/python3.9/typing.py", line 852, in __subclasscheck__
        return issubclass(cls, self.__origin__)
    TypeError: issubclass() arg 1 must be a class

Related issues

From langchain:

• TypeError: issubclass() arg 1 must be a class when using langchain in azure langchain-ai/langchain#7548
• import langchain with python<=3.9 fails langchain-ai/langchain#5113

[dlqqq]

This issue occurs in a SageMaker Studio environment running pydantic==1.10.2

Notably, LangChain users claim that the issue is fixed in pydantic>=1.10.8. I will test upgrading this package and see if that allows jupyter-ai to start.

[dlqqq]

I ran the following commands in an interactive shell through the JupyterLab terminal:

% conda activate studio
% pip install pydantic~=1.10.8
% restart-jupyter-server

This installs pydantic==1.10.13, and allows Jupyter AI to work in SageMaker Studio.

Interestingly, this version is fully compatible with the studio environment. SageMaker Studio should bump the version of pydantic used in the studio environment to either the latest v1 or v2.