Add pydantic v2 support (#181)

This PR adds pydantic v2 support to LangServe. With pydantic v2 all endpoints will work, but openapi docs will not be generated for these endpoints. (At least at the moment.) --------- Co-authored-by: jakerachleff <[email protected]>
langchain-ai · Nov 8, 2023 · 0dd55c8 · 0dd55c8
1 parent 2f8d046
commit 0dd55c8
Show file tree

Hide file tree

Showing 7 changed files with 272 additions and 160 deletions.
diff --git a/.github/workflows/langserve_ci.yml b/.github/workflows/langserve_ci.yml
@@ -39,6 +39,12 @@ jobs:
       working-directory: .
     secrets: inherit
 
+  pydantic-compatibility:
+    uses:
+      ./.github/workflows/_pydantic_compatibility.yml
+    with:
+      working-directory: .
+    secrets: inherit
   test:
     runs-on: ubuntu-latest
     defaults:
@@ -51,7 +57,7 @@ jobs:
           - "3.9"
           - "3.10"
           - "3.11"
-    name: Python ${{ matrix.python-version }} extended tests
+    name: Python ${{ matrix.python-version }} tests
     steps:
       - uses: actions/checkout@v3
 

diff --git a/langserve/pydantic.py b/langserve/pydantic.py
@@ -0,0 +1,10 @@
+import pydantic
+
+
+def _get_pydantic_version() -> int:
+    """Get the pydantic major version."""
+    return int(pydantic.__version__.split(".")[0])
+
+
+# Code is written to support both version 1 and 2
+PYDANTIC_MAJOR_VERSION = _get_pydantic_version()
diff --git a/langserve/schema.py b/langserve/schema.py
@@ -2,13 +2,16 @@
 from typing import Dict, List, Optional, Union
 from uuid import UUID
 
-try:
-    from pydantic.v1 import BaseModel
-except ImportError:
-    from pydantic import BaseModel
+from langserve.pydantic import PYDANTIC_MAJOR_VERSION
 
+if PYDANTIC_MAJOR_VERSION == 2:
+    from pydantic.v1 import BaseModel as BaseModelV1
+else:
+    from pydantic import BaseModel as BaseModelV1
+from pydantic import BaseModel
 
-class CustomUserType(BaseModel):
+
+class CustomUserType(BaseModelV1):
     """Inherit from this class to create a custom user type.
 
     Use a custom user type if you want the data to de-serialize
@@ -27,7 +30,7 @@ class CustomUserType(BaseModel):
     """
 
 
-class SharedResponseMetadata(BaseModel):
+class SharedResponseMetadata(BaseModelV1):
     """
     Any response metadata should inherit from this class. Response metadata
     represents non-output data that may be useful to some clients, but