diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py index 75c8fa0a3..a539b90b9 100644 --- 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 index 31870fcb5..951bea385 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -12,20 +12,19 @@ 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, @@ -33,8 +32,6 @@ 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 index 3f67c07d2..e1a22c9cc 100644 --- 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),