From 3d024de4edbc39d0cb6128e770d1112fe26fa38f Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 30 Aug 2023 21:06:09 -0700 Subject: [PATCH 1/2] Upgrades LangChain to 0.0.277 --- .../jupyter_ai_magics/embedding_providers.py | 2 +- .../jupyter_ai_magics/parsers.py | 2 +- .../jupyter_ai_magics/providers.py | 2 +- packages/jupyter-ai-magics/pyproject.toml | 2 +- .../jupyter_ai/chat_handlers/learn.py | 19 +++++++++++++++---- packages/jupyter-ai/jupyter_ai/extension.py | 6 +++--- packages/jupyter-ai/jupyter_ai/handlers.py | 2 +- packages/jupyter-ai/jupyter_ai/models.py | 2 +- packages/jupyter-ai/pyproject.toml | 2 +- 9 files changed, 25 insertions(+), 14 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 index bdfd7012c..1d8f8f8ab 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py @@ -7,7 +7,7 @@ OpenAIEmbeddings, ) from langchain.embeddings.base import Embeddings -from pydantic import BaseModel, Extra +from langchain.pydantic_v1 import BaseModel, Extra class BaseEmbeddingsProvider(BaseModel): diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py index a6acf3525..5c2d431b5 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -1,7 +1,7 @@ from typing import Literal, Optional, get_args import click -from pydantic import BaseModel +from langchain.pydantic_v1 import BaseModel FORMAT_CHOICES_TYPE = Literal[ "code", "html", "image", "json", "markdown", "math", "md", "text" diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index 93616a655..9d76e2860 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -23,8 +23,8 @@ ) from langchain.llms.sagemaker_endpoint import LLMContentHandler from langchain.llms.utils import enforce_stop_tokens +from langchain.pydantic_v1 import BaseModel, Extra, root_validator from langchain.utils import get_from_dict_or_env -from pydantic import BaseModel, Extra, root_validator class EnvAuthStrategy(BaseModel): diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index 1cd7a513f..eee529873 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ "ipython", "pydantic", "importlib_metadata>=5.2.0", - "langchain==0.0.223", + "langchain==0.0.277", "typing_extensions>=4.5.0", "click~=8.0", "jsonpath-ng>=1.5.3,<2", diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py index 963eae6ce..7d88beac2 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py @@ -4,6 +4,7 @@ from typing import Any, Awaitable, 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 ( @@ -29,7 +30,7 @@ METADATA_SAVE_PATH = os.path.join(INDEX_SAVE_DIR, "metadata.json") -class LearnChatHandler(BaseChatHandler, BaseRetriever): +class LearnChatHandler(BaseChatHandler): def __init__( self, root_dir: str, dask_client_future: Awaitable[DaskClient], *args, **kwargs ): @@ -266,9 +267,6 @@ def load_metadata(self): j = json.loads(f.read()) self.metadata = IndexMetadata(**j) - def get_relevant_documents(self, query: str) -> List[Document]: - raise NotImplementedError() - async def aget_relevant_documents( self, query: str ) -> Coroutine[Any, Any, List[Document]]: @@ -291,3 +289,16 @@ def get_embedding_model(self): return None return em_provider_cls(**em_provider_args) + + +class Retriever(BaseRetriever): + learn_chat_handler: LearnChatHandler = None + + def _get_relevant_documents(self, query: str) -> List[Document]: + raise NotImplementedError() + + async def _aget_relevant_documents( + self, query: str + ) -> Coroutine[Any, Any, List[Document]]: + docs = await self.learn_chat_handler.aget_relevant_documents(query) + return docs diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index f9d0fee56..a6564679b 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -1,6 +1,7 @@ import time from dask.distributed import Client as DaskClient +from jupyter_ai.chat_handlers.learn import Retriever from jupyter_ai_magics.utils import get_em_providers, get_lm_providers from jupyter_server.extension.application import ExtensionApp @@ -93,9 +94,8 @@ def initialize_settings(self): dask_client_future=dask_client_future, ) help_chat_handler = HelpChatHandler(**chat_handler_kwargs) - ask_chat_handler = AskChatHandler( - **chat_handler_kwargs, retriever=learn_chat_handler - ) + retriever = Retriever(learn_chat_handler=learn_chat_handler) + ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever) self.settings["jai_chat_handlers"] = { "default": default_chat_handler, "/ask": ask_chat_handler, diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py index 95db5198a..c8f3b2803 100644 --- a/packages/jupyter-ai/jupyter_ai/handlers.py +++ b/packages/jupyter-ai/jupyter_ai/handlers.py @@ -10,7 +10,7 @@ from jupyter_ai.chat_handlers import BaseChatHandler from jupyter_server.base.handlers import APIHandler as BaseAPIHandler from jupyter_server.base.handlers import JupyterHandler -from pydantic import ValidationError +from langchain.pydantic_v1 import ValidationError from tornado import web, websocket from tornado.web import HTTPError diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py index ca2030881..0512ed1e0 100644 --- a/packages/jupyter-ai/jupyter_ai/models.py +++ b/packages/jupyter-ai/jupyter_ai/models.py @@ -1,7 +1,7 @@ from typing import Any, Dict, List, Literal, Optional, Union from jupyter_ai_magics.providers import AuthStrategy, Field -from pydantic import BaseModel +from langchain.pydantic_v1 import BaseModel DEFAULT_CHUNK_SIZE = 2000 DEFAULT_CHUNK_OVERLAP = 100 diff --git a/packages/jupyter-ai/pyproject.toml b/packages/jupyter-ai/pyproject.toml index 7c7792dd3..ccad30f8c 100644 --- a/packages/jupyter-ai/pyproject.toml +++ b/packages/jupyter-ai/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "openai~=0.26", "aiosqlite>=0.18", "importlib_metadata>=5.2.0", - "langchain==0.0.223", + "langchain==0.0.277", "tiktoken", # required for OpenAIEmbeddings "jupyter_ai_magics", "dask[distributed]", From 9bc14386d5b05373d979922da39641d194ac688c Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Fri, 1 Sep 2023 12:31:13 -0700 Subject: [PATCH 2/2] Pinned Pydantic version, updated pydantic references --- .../jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py | 2 +- packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py | 2 +- packages/jupyter-ai-magics/jupyter_ai_magics/providers.py | 2 +- packages/jupyter-ai-magics/pyproject.toml | 2 +- packages/jupyter-ai/jupyter_ai/handlers.py | 2 +- packages/jupyter-ai/jupyter_ai/models.py | 2 +- packages/jupyter-ai/pyproject.toml | 2 +- 7 files changed, 7 insertions(+), 7 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 index 1d8f8f8ab..bdfd7012c 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py @@ -7,7 +7,7 @@ OpenAIEmbeddings, ) from langchain.embeddings.base import Embeddings -from langchain.pydantic_v1 import BaseModel, Extra +from pydantic import BaseModel, Extra class BaseEmbeddingsProvider(BaseModel): diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py index 5c2d431b5..a6acf3525 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -1,7 +1,7 @@ from typing import Literal, Optional, get_args import click -from langchain.pydantic_v1 import BaseModel +from pydantic import BaseModel FORMAT_CHOICES_TYPE = Literal[ "code", "html", "image", "json", "markdown", "math", "md", "text" diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index 9d76e2860..93616a655 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -23,8 +23,8 @@ ) from langchain.llms.sagemaker_endpoint import LLMContentHandler from langchain.llms.utils import enforce_stop_tokens -from langchain.pydantic_v1 import BaseModel, Extra, root_validator from langchain.utils import get_from_dict_or_env +from pydantic import BaseModel, Extra, root_validator class EnvAuthStrategy(BaseModel): diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index eee529873..38d885e7a 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -22,7 +22,7 @@ dynamic = ["version", "description", "authors", "urls", "keywords"] dependencies = [ "ipython", - "pydantic", + "pydantic~=1.0", "importlib_metadata>=5.2.0", "langchain==0.0.277", "typing_extensions>=4.5.0", diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py index c8f3b2803..95db5198a 100644 --- a/packages/jupyter-ai/jupyter_ai/handlers.py +++ b/packages/jupyter-ai/jupyter_ai/handlers.py @@ -10,7 +10,7 @@ from jupyter_ai.chat_handlers import BaseChatHandler from jupyter_server.base.handlers import APIHandler as BaseAPIHandler from jupyter_server.base.handlers import JupyterHandler -from langchain.pydantic_v1 import ValidationError +from pydantic import ValidationError from tornado import web, websocket from tornado.web import HTTPError diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py index 0512ed1e0..ca2030881 100644 --- a/packages/jupyter-ai/jupyter_ai/models.py +++ b/packages/jupyter-ai/jupyter_ai/models.py @@ -1,7 +1,7 @@ from typing import Any, Dict, List, Literal, Optional, Union from jupyter_ai_magics.providers import AuthStrategy, Field -from langchain.pydantic_v1 import BaseModel +from pydantic import BaseModel DEFAULT_CHUNK_SIZE = 2000 DEFAULT_CHUNK_OVERLAP = 100 diff --git a/packages/jupyter-ai/pyproject.toml b/packages/jupyter-ai/pyproject.toml index ccad30f8c..71ae54ff4 100644 --- a/packages/jupyter-ai/pyproject.toml +++ b/packages/jupyter-ai/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ dependencies = [ "jupyter_server>=1.6,<3", "jupyterlab~=4.0", - "pydantic", + "pydantic~=1.0", "openai~=0.26", "aiosqlite>=0.18", "importlib_metadata>=5.2.0",