migrate import statements deprecated by langchain v0.1

dlqqq · Jan 10, 2024 · b931852 · b931852
1 parent 8f8d59e
commit b931852
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 11 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
@@ -1,4 +1,3 @@
-import os
 from typing import ClassVar, List
 
 from jupyter_ai_magics.providers import (
@@ -8,7 +7,7 @@
     Field,
     MultiEnvAuthStrategy,
 )
-from langchain.embeddings import (
+from langchain_community.embeddings import (
     BedrockEmbeddings,
     CohereEmbeddings,
     GPT4AllEmbeddings,

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -12,29 +12,26 @@
     Dict,
     List,
     Literal,
-    Mapping,
     Optional,
     Union,
 )
 
 from jsonpath_ng import parse
-from langchain.chat_models import (
+from langchain_community.chat_models import (
     AzureChatOpenAI,
     BedrockChat,
     ChatAnthropic,
     QianfanChatEndpoint,
 )
 from langchain.chat_models.base import BaseChatModel
-from langchain.llms import (
+from langchain_community.llms import (
     AI21,
     Anthropic,
     Bedrock,
     Cohere,
     GPT4All,
     HuggingFaceHub,
     OpenAI,
-    OpenAIChat,
-    QianfanLLMEndpoint,
     SagemakerEndpoint,
 )
 from langchain.llms.sagemaker_endpoint import LLMContentHandler

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
@@ -1,10 +1,9 @@
 import argparse
 import json
 import os
-from typing import Any, Awaitable, Coroutine, List, Optional, Tuple
+from typing import Any, Coroutine, List, Optional, Tuple
 
 from dask.distributed import Client as DaskClient
-from jupyter_ai.config_manager import ConfigManager
 from jupyter_ai.document_loaders.directory import get_embeddings, split
 from jupyter_ai.document_loaders.splitter import ExtensionSplitter, NotebookSplitter
 from jupyter_ai.models import (
@@ -22,7 +21,7 @@
     PythonCodeTextSplitter,
     RecursiveCharacterTextSplitter,
 )
-from langchain.vectorstores import FAISS
+from langchain_community.vectorstores import FAISS
 
 from .base import BaseChatHandler, SlashCommandRoutingType
 
@@ -141,7 +140,7 @@ def _build_list_response(self):
     async def learn_dir(
         self, path: str, chunk_size: int, chunk_overlap: int, all_files: bool
     ):
-        dask_client = await self.dask_client_future
+        dask_client: DaskClient = await self.dask_client_future
         splitter_kwargs = {"chunk_size": chunk_size, "chunk_overlap": chunk_overlap}
         splitters = {
             ".py": PythonCodeTextSplitter(**splitter_kwargs),