From 9bc14386d5b05373d979922da39641d194ac688c Mon Sep 17 00:00:00 2001
From: Piyush Jain <piyushjain@duck.com>
Date: Fri, 1 Sep 2023 12:31:13 -0700
Subject: [PATCH] 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",