feat: adapt pydantic v2

Author: JoanFM

jina/_docarray.py

@@ -8,3 +8,9 @@
     from docarray import Document, DocumentArray
 
     docarray_v2 = False
+
+
+import pydantic
+
+is_pydantic_v2 = pydantic.__version__.startswith('2.')
+

jina/serve/runtimes/gateway/graph/topology_graph.py

@@ -7,7 +7,7 @@
 
 import grpc.aio
 
-from jina._docarray import DocumentArray, docarray_v2
+from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2
 from jina.constants import __default_endpoint__
 from jina.excepts import InternalNetworkError
 from jina.logging.logger import JinaLogger
@@ -20,7 +20,10 @@
     from docarray import DocList
     from docarray.documents.legacy import LegacyDocument
 
-    from jina.serve.runtimes.helper import _create_pydantic_model_from_schema
+    if not is_pydantic_v2:
+        from jina.serve.runtimes.helper import _create_pydantic_model_from_schema
+    else:
+
 
     legacy_doc_schema = LegacyDocument.schema()
 

jina/serve/runtimes/helper.py

@@ -1,7 +1,7 @@
 import copy
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-from jina._docarray import docarray_v2
+from jina._docarray import docarray_v2, is_pydantic_v2
 
 _SPECIFIC_EXECUTOR_SEPARATOR = '__'
 
@@ -79,7 +79,7 @@ def _parse_specific_params(parameters: Dict, executor_name: str):
     return parsed_params
 
 
-if docarray_v2:
+if docarray_v2 and not is_pydantic_v2:
     from docarray import BaseDoc, DocList
     from docarray.typing import AnyTensor
     from pydantic import create_model

jina/serve/runtimes/worker/request_handling.py

@@ -20,7 +20,7 @@
 
 from google.protobuf.struct_pb2 import Struct
 
-from jina._docarray import DocumentArray, docarray_v2
+from jina._docarray import DocumentArray, docarray_v2, is_pydantic_v2
 from jina.constants import __default_endpoint__
 from jina.excepts import BadConfigSource, RuntimeTerminated
 from jina.helper import get_full_version
@@ -1013,21 +1013,24 @@ async def endpoint_discovery(self, empty, context) -> jina_pb2.EndpointsProto:
         if docarray_v2:
             from docarray.documents.legacy import LegacyDocument
 
-            from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list
+            if not is_pydantic_v2:
+                from jina.serve.runtimes.helper import _create_aux_model_doc_list_to_list as create_pure_python_type_model
+            else:
+                from docarray.utils.create_dynamic_doc_class import create_pure_python_type_model
 
             legacy_doc_schema = LegacyDocument.schema()
             for endpoint_name, inner_dict in schemas.items():
                 if inner_dict['input']['model'].schema() == legacy_doc_schema:
                     inner_dict['input']['model'] = legacy_doc_schema
                 else:
-                    inner_dict['input']['model'] = _create_aux_model_doc_list_to_list(
+                    inner_dict['input']['model'] = create_pure_python_type_model(
                         inner_dict['input']['model']
                     ).schema()
 
                 if inner_dict['output']['model'].schema() == legacy_doc_schema:
                     inner_dict['output']['model'] = legacy_doc_schema
                 else:
-                    inner_dict['output']['model'] = _create_aux_model_doc_list_to_list(
+                    inner_dict['output']['model'] = create_pure_python_type_model(
                         inner_dict['output']['model']
                     ).schema()